import numpy as np
import matplotlib.pyplot as plt
import matplotlib.animation as ani
import matplotlib.gridspec as gridspec
from matplotlib.colors import ListedColormap

#紙の準備
fig = plt.figure()
fig.canvas.draw()
fig.suptitle("2022京大数学　文系第３問",
             color="0.5",ha="right",x=0.95,y=0.95,zorder=-1)
gs = gridspec.GridSpec(5,10)
#放物線とかを描く紙
ax1 = fig.add_subplot(gs[1:,0:4])
ax1.set_xlim(-2,5)
ax1.set_ylim(-2,5)
ax1.axes.xaxis.set_visible(False)
ax1.axes.yaxis.set_visible(False)
ax1.spines["top"].set_visible(False)
ax1.spines["right"].set_visible(False)
ax1.spines["left"].set_position(("data",0))
ax1.spines["bottom"].set_position(("data",0))
ax1.spines["left"].set_color("0.5")
ax1.spines["bottom"].set_color("0.5")
ax1.plot(1,0,">",color="0.6",transform=ax1.get_yaxis_transform(),clip_on=False)
ax1.plot(0,1,"^",color="0.6",transform=ax1.get_xaxis_transform(),clip_on=False)
ax1.set_aspect("equal")
#積分過程を書く紙
ax2 = fig.add_subplot(gs[:,4:])
width = 8
height = 8
ax2.set_xlim(0,width)
ax2.set_ylim(0,height)
ax2.set_aspect("equal")
ax2.axis("off")

#フォント辞書
linedic={
"usetex" : True,
"clip_on" : True,
"animated" : True,
"fontsize" : 16,
"verticalalignment" : "baseline",
}
#積分過程
lines=[
r"= \int_{-1}^\frac{3}{2} \left\{ \frac{1}{4}x^2 \!-\! \left(\!-\frac{1}{2}x-\frac{1}{4}\right) \right\} dx",
r"+ \int_\frac{3}{2}^4 \left\{ \frac{1}{4}x^2 - \left(2x-4\right) \right\} dx",
r"= \frac{1}{4} \int_{-1}^\frac{3}{2} \left( x^2 +2x +1 \right) dx",
r"+ \,\frac{1}{4} \int_\frac{3}{2}^4 \left( x^2 -8x +16 \right) dx",
r"= \frac{1}{4} \int_{-1}^\frac{3}{2} \left( x+1 \right)^2 dx",
r"+ \,\frac{1}{4} \int_\frac{3}{2}^4 \left( x-4 \right)^2 dx",
r"= \frac{1}{4} \left[ \,\frac{1}{3} \left( x+1 \right)^3 \,\right]_{-1}^\frac{3}{2}",
r"+ \,\frac{1}{4} \left[ \,\frac{1}{3} \left( x-4 \right)^3 \,\right]_\frac{3}{2}^4",
r"= \frac{1}{4}\cdot\frac{1}{3}\left( \frac{3}{2}+1 \right)^3 - 0",
r"+ \,0 - \frac{1}{4}\cdot\frac{1}{3} \left( \frac{3}{2}-4 \right)^3",
r"= \frac{1}{4}\cdot\frac{1}{3}\left(\frac{5}{2}\right)^3-\frac{1}{4}\cdot\frac{1}{3}\left(\!-\frac{5}{2}\right)^3",
r"= \frac{125}{96}+\frac{125}{96} = \frac{125}{48}",
]

#定数の準備
base = height/10
delta = 0.1
colors = [(1,1,1,alpha) for alpha in np.arange(0,11)/10]
cmap = ListedColormap(colors)
#放物線と直線
x = np.linspace(-2,5,71)
yc = x**2/4
ax1.plot(x,yc,"0.4",zorder=10)
ax1.text(-2,1.2,"$\\displaystyle y=\\frac{x^2}{4}$",
         color="0.2",ha="center",va="bottom",usetex=True)
y1 = -1/2*x-1/4
ax1.plot(x,y1,"0.4",zorder=10)
ax1.text(7/2,-2,"$\\displaystyle y=-\\frac{x}{2}-\\frac{1}{4}$",
         color="0.2",ha="center",va="top",usetex=True)
y2 = 2*x-4
ax1.plot(x,y2,"0.4",zorder=10)
ax1.text(1.2,-2.4,"$\\displaystyle y=2x-4$",
         color="0.2",ha="right",va="top",usetex=True)
#求める面積
Lspan = (x>-1) & (x<3/2+0.01)
ax1.fill_between(x[Lspan],y1[Lspan],yc[Lspan],zorder=5,
                 fc="1.0",ec="0.6",hatch="/////")
Rspan = (x>3/2-0.01) & (x<4)
ax1.fill_between(x[Rspan],y2[Rspan],yc[Rspan],zorder=5,
                 fc="1.0",ec="0.6",hatch="\\\\\\\\\\")
#共有点たち
ax1.plot(-1,1/4,"o",color="0.4",zorder=10)
ax1.plot( 4,  4,"o",color="0.4",zorder=10)
ax1.plot(3/2,-1,"o",color="0.4",zorder=10)
ax1.text(-1,1/4,"$(-1,\\dfrac{1}{4})$",
         color="0.4",ha="left",va="bottom",backgroundcolor="1.0")
ax1.text( 4,  4,"$(4,4)$  ",
         color="0.4",ha="right",va="center")
ax1.text(3/2,-1,"  $(\\,\\dfrac{3}{2},-\\!1)$",
         color="0.4",ha="left",va="baseline")
#与式
ax2.text(-width/40,base+3*height/4,"$\\displaystyle S$",
         fontdict=linedic,clip_on=False,animated=False,ha="right")

#数式をLATEX変換する関数
def show_line(y,line,alpha=1.0):
    img = []
    line = "$\\displaystyle " + line + "$" 
    text = ax2.text(0,y,line,alpha=alpha,fontdict=linedic)
    img += [text]
    bbox = text.get_window_extent().transformed(ax2.transData.inverted())
    right = bbox.x1
    return img,right

#余韻を残すための関数
def lingering(imgs,frames):
    img = imgs[-1]
    for i in range(frames):
        imgs.append(img)
        
#数式が徐々に現れる関数
def appearing(i,X,Y):
    C = np.zeros(len(X)-1)
    C[2*i+1:2*i+20] = np.arange(0.05,1.00,0.05)
    C[2*i+20:] = 1
    C = np.reshape(C,(len(X)-1,1))
    img = ax2.pcolormesh(X,Y,C,cmap=cmap,zorder=10)
    return img

#最初の1,2行を表示する関数
def show_first_lines(imgs,line1,line2):
    
    i = 0
    while True:
        img1,right1 = show_line(base+3*height/4,line1)
        X,Y = np.mgrid[-2:right1+4:delta,3*height/4:5*height/4:height/4]
        img = img1 + [appearing(i,X,Y)]
        imgs.append(img)
        i += 1
        if 2*i+20 > len(X)-1:
            break
        
    i = 0
    while True:
        img2,right2 = show_line(base+2*height/4,line2)
        X,Y = np.mgrid[-2:right2+4:delta,2*height/4:4*height/4:height/4]
        img = img1 + img2 + [appearing(i,X,Y)]
        imgs.append(img)
        i += 1
        if 2*i+20 > len(X)-1:
            break

#3,4行目を表示する関数
def show_second_lines(imgs,line1,line2,line3,line4):
    
    i = 0
    while True:
        img1,right1 = show_line(base+3*height/4,line1)
        img2,right2 = show_line(base+2*height/4,line2)
        img3,right3 = show_line(base+  height/4,line3)
        img123 = img1 + img2 + img3
        X,Y = np.mgrid[-2:right3+4:delta,height/4:3*height/4:height/4]
        img = img123 + [appearing(i,X,Y)]
        imgs.append(img)
        i += 1
        if 2*i+20 > len(X)-1:
            break
    
    i = 0
    while True:
        img4,right4 = show_line(base,line4)
        img = img123 + img4
        X,Y = np.mgrid[-2:right4+4:delta,0:2*height/4:height/4]
        img = img123 + img4 + [appearing(i,X,Y)]
        imgs.append(img)
        i += 1
        if 2*i+20 > len(X)-1:
            break
        
#次の数式へと進む関数
def go_next_line(imgs,line1,line2,line3,line4):

    for i in range(21):
        alpha = max((10-i)/10,0)
        img1,right1 = show_line(base+3*height/4+height*i/40,line1,alpha)
        img2,right2 = show_line(base+2*height/4+height*i/40,line2,alpha)
        img3,right3 = show_line(base+  height/4+height*i/40,line3)
        img4,right4 = show_line(base           +height*i/40,line4)
        img = img1 + img2 + img3 + img4
        imgs.append(img)

#数式を読み込む関数
def load_script(imgs,lines):
    
    N = len(lines)
    
    show_first_lines(imgs,lines[0],lines[1])
    
    for k in range(0,N-4,2):
        show_second_lines(imgs,lines[k],lines[k+1],lines[k+2],lines[k+3])
        go_next_line(imgs,lines[k],lines[k+1],lines[k+2],lines[k+3])
        
    show_second_lines(imgs,lines[N-4],lines[N-3],lines[N-2],lines[N-1])
    lingering(imgs,100)

#上演
imgs=[]
load_script(imgs,lines)
mov=ani.ArtistAnimation(fig, imgs, 100)
plt.show()