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

#紙の準備
def ax_settings(ax,width,height):
    ax.set_xlim(0,width)
    ax.set_ylim(0,height)
    ax.set_aspect("equal")
    ax.axis("off")

fig = plt.figure()
fig.canvas.draw()
fig.suptitle("2021京大数学　文系第５問",
             color="0.5",ha="right",x=0.95,y=0.95,zorder=-1)
width=4
height=8
ax0 = fig.add_subplot(131)
ax_settings(ax0,width,height)
ax1 = fig.add_subplot(132)
ax_settings(ax1,width,height)
ax2 = fig.add_subplot(133)
ax_settings(ax2,width,height)

#フォントの準備
fdic={
"animated" : True,
"fontsize" : 14,
"fontfamily": "Meiryo",
"verticalalignment" : "baseline",
}
fdic_wc={
"ha" : "center",
"va" : "center",
"color" : "white",
"animated" : True,
"fontsize" : 20,
"fontfamily": "Meiryo",
}
boxdic={
"facecolor" : "black",
"pad" : 300,
}

#台本
comment0="""\
$p$ が素数$\\;\\Rightarrow\\; p^4+14$ は素数でない

を示せ
"""
comment1="""\
(´･ω･`;)　えっ
"""
comment2="""\
落ち着くんだ
"""
comment3="""\
あるゆる整数は

3 で割ったときの余りに応じて

余りが 0 か 1 か 2 のときに

に分けられる
"""
comment4="""\
つまり
"""
comment5="""\
あるゆる整数は

3 を法として

$p\\equiv0\\;,\\;p\\equiv1\\;,\\;p\\equiv2$

のどれかになるってわけさ
"""
lines0=[
r"$p\equiv0$ のとき",
r"$\quad\; p=3$ のみ",
r"( ∵ $p$ は素数 )",
r"$\quad\; p^4+14$",
r"$=3^4+14$",
r"$=95$",
r"$=19\times5$",
" 5 で割り切れる",
"→ 素数じゃない",
]
lines1=[
r"$p\equiv1$ のとき",
"",
"",
r"$\quad\; p^4+14$",
r"$\equiv 1^4+14$",
r"$\equiv 15$",
r"$\equiv 0$    (mod 3)",
" 3 で割り切れる",
"→ 素数じゃない",
]
lines2=[
r"$p\equiv2$ のとき",
"",
"",
r"$\quad\; p^4+14$",
r"$\equiv 2^4+14$",
r"$\equiv 30$",
r"$\equiv 0$    (mod 3)",
" 3 で割り切れる",
"→ 素数じゃない",
]
comment10="""\
Q.E.D.

(｡･ω･｡) /
"""

#定数のセッティング
base=0.3
delta=0.1
colors = [(1,1,1,alpha) for alpha in np.arange(0,11)/10]
cmap = ListedColormap(colors)

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

#過去の場合分け
def oldax(ax,lines):
    oldimg = [ax.text(width/2,height+base,lines[0],
                      fontdict=fdic,color="0.5",ha="center")]
    for i in range(len(lines)-1):
        n = i+1
        oldimg += [ax.text(0,height-n+base,lines[n],
                           fontdict=fdic,color="0.5")]
    oldimg += ax.plot([1.1*width,1.1*width],[0,1.1*height],
                      color="0.5",linestyle="dashed",clip_on=False)
    return oldimg

#今の場合分け
def newax(imgs,oldimg,ax,lines):
    
    for i in range(0,len(lines)):
        
        k = 0
        while True:
            newimg = oldimg.copy()
            newimg += [ax.text(width/2,height+base,lines[0],
                           fontdict=fdic,ha="center")]
            if i == 0:
                imgs.append(newimg)
                lingering(imgs,20)
                break
            for j in range(1,i):
                newimg += [ax.text(0,height-j+base,lines[j],fontdict=fdic)]
            if lines[i]=="":
                break
            else:
                last_line = ax.text(0,height-i+base,lines[i],fontdict=fdic)
                bbox = last_line.get_window_extent().transformed(ax.transData.inverted())
                newimg += [last_line]
                X,Y = np.mgrid[-2:bbox.x1+2:delta,height-i:height-i+2:1]
                newimg += [appearing(ax,k,X,Y)]
                imgs.append(newimg)
                k += 1
                if 2*k+20 > len(X)-1:
                    break
                
    lingering(imgs,20)

#コメントの部分の関数
def interlude(imgs,str_comment,frames):
    for i in range(frames):
        interlude=ax1.text(width/2,height/2,str_comment,
                          fontdict=fdic_wc,bbox=boxdic)
        imgs.append([interlude])

#上演
imgs = []
interlude(imgs,comment0,60)
interlude(imgs,comment1,30)
interlude(imgs,comment2,20)
interlude(imgs,comment3,80)
interlude(imgs,comment4,20)
interlude(imgs,comment5,80)
oldimg = []
newax(imgs,oldimg,ax0,lines0)
oldimg = oldax(ax0,lines0)
newax(imgs,oldimg,ax1,lines1)
oldimg = oldax(ax0,lines0) + oldax(ax1,lines1)
newax(imgs,oldimg,ax2,lines2)  
interlude(imgs,comment10,50) 

mov=ani.ArtistAnimation(fig, imgs, 100)
plt.show()