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

#紙の準備
fig = plt.figure()
fig.canvas.draw()
fig.suptitle("2021京大数学　理系第３問",
             color="0.5",ha="right",x=0.95,y=0.95,zorder=-1)
width  = 8
height = 4
ax = fig.add_axes((0.35,0,0.6,1))
ax.set_xlim(0,width)
ax.set_ylim(0,height)
ax.set_aspect("equal")
ax.axis("off")

#フォント辞書たち
linedic={
"usetex" : True,
"clip_on" : True,
"animated" : True,
"fontsize" : 18,
"verticalalignment" : "baseline",
}
txtdic={
"ha" : "center",
"va" : "center",
"size" : 20,
"color" : "white",
"animated" : True,
"fontfamily": "Meiryo",
}
texdic={
"ha" : "center",
"va" : "center",
"size" : 22,
"color" : "white",
"usetex" : True,
"animated" : True,
}
txtbgdic={
"facecolor" : "black",
"pad" : 300,
}

#台本(計算以外)
text1=[
"無限級数",
"$\\displaystyle \\sum_{n=0}^\\infty \\, \\left(\\frac{\\;1\\;}{\\;2\\;}\\right)^{\\!\\!n} \\cos \\frac{n\\pi}{6}$",
"を求めよ",
]
text2=[
"オイラーの公式",
"$\\displaystyle e^{i\\theta} = \\cos\\theta + i\\sin\\theta$",
"を使います",
]
a0a=('fig.text(0.70,0.15,'
     '"大学で習うものなので、",'
     'color="w",alpha={})')
a0b=('fig.text(0.70,0.10,'
     '"ホントは反則なんだけど...",'
     'color="w",alpha={})')
text3=[
"すると",
"$\\displaystyle \\cos\\theta = \\frac{\\;e^{i\\theta}+e^{-i\\theta}\\;}{\\;2\\;}$",
"となるので...".rjust(10),
]
text4=[
"さっきの $f(\\theta)$ に",
"$\\displaystyle \\theta = \\frac{\\;\\pi\\;}{\\;6\\;}$",
"を代入すると...".rjust(10),
]
text5=[
"答えは",
"$\\displaystyle \\frac{\\;14+3\\sqrt{3}\\;}{\\;13\\;}$",
"です！".rjust(5),
]
text6=[
"ちなみに",
"$\\displaystyle \\theta = 0$",
"を代入すると...".rjust(10),
]
text7=[
"あと",
"$\\displaystyle \\theta = \\pi$",
"を代入すると...".rjust(10),
]

#台本(計算部分)
line1=[
r"= \sum_{n=0}^\infty \, \left(\frac{\;1\;}{\;2\;}\right)^{\!\!n} \left(\frac{\;e^{in\theta}+e^{-in\theta}\;}{\;2\;}\right)",
r"= \frac{\;1\;}{\;2\;}\sum_{n=0}^\infty \left(\frac{\;e^{i\theta}\;}{\;2\;}\right)^{\!\!n} + \frac{\;1\;}{\;2\;}\sum_{n=0}^\infty \left(\frac{\;e^{-i\theta}\;}{\;2\;}\right)^{\!\!n}",
r"= \frac{\;1\;}{\;2\;}\cdot\frac{\;1\;}{\;1-\,^{e^{i\theta}}\!/_2\;} + \frac{\;1\;}{\;2\;}\cdot\frac{\;1\;}{\;1-\,^{e^{-i\theta}}\!/_2\;}",
r"= \frac{\;1\;}{\;2-e^{i\theta}\;} + \frac{\;1\;}{\;2-e^{-i\theta}\;}",
r"= \frac{\;(2-e^{i\theta})+(2-e^{-i\theta})\;}{\;(2-e^{i\theta})\,(2-e^{-i\theta})\;}",
r"= \frac{\;4-e^{i\theta}-e^{-i\theta}\;}{\;4-2\,e^{i\theta}-2\,e^{-i\theta}+e^{i\theta}\cdot e^{-i\theta}\;}",
r"= \frac{\;4-(e^{i\theta}+e^{-i\theta})\;}{\;5-2\,(e^{i\theta}+e^{-i\theta})\;}",
r"= \frac{\;4-2\cos\theta\;}{\;5-4\cos\theta\;}",
]
line2=[
r"= f\left(\frac{\;\pi\;}{\;6\;}\right) = \frac{\;4-2\cos\frac{\;\pi\;}{\;6\;}\;}{\;5-4\cos\frac{\;\pi\;}{\;6\;}\;}",
r"= \frac{\;4-2\cdot\frac{\;\sqrt{3}\;}{\;2\;}\;}{\;5-4\cdot\frac{\;\sqrt{3}\;}{\;2\;}\;}",
r"= \frac{\;4-\sqrt{3}\;}{\;5-2\sqrt{3}\;}",
r"= \frac{\;(4-\sqrt{3}\,)(5+2\sqrt{3}\,)\;}{\;(5-2\sqrt{3}\,)(5+2\sqrt{3}\,)\;}",
r"= \frac{\;14+3\sqrt{3}\;}{\;13\;}",
]
line3=[
r"= f(0) = \frac{\;4-2\cos 0\;}{\;5-4\cos 0\;}",
r"= \frac{\;4-2\;}{\;5-4\;} = 2",
]
line4=[
r"= f(\pi) = \frac{\;4-2\cos\pi\;}{\;5-4\cos\pi\;}",
r"= \frac{\;4+2\;}{\;5+4\;} = \frac{\;2\;}{\;3\;}",
]

#定数の準備
base = height/5
delta = 0.1
colors = [(1,1,1,alpha) for alpha in np.arange(0,11)/10]
cmap = ListedColormap(colors)

#与式とか
f1  = ax.text(-4,base+height/2,
              "$\\displaystyle \\sum_{n=0}^\\infty \\, \\left(\\frac{\\;1\\;}{\\;2\\;}\\right)^{\\!\\!n} \\cos n\\theta$",
              fontdict=linedic,clip_on=False)
a1a =('ax.text(0,base-height/2,'
      '"$\\displaystyle =: f(\\\\theta)$",'
      'fontdict=linedic,clip_on=False,alpha={})')
a1b =('ax.text(2,base-height/2,'
      '"とおきます",'
      'fontsize=14,alpha={})')
f2  = ax.text(-4,base+height/2,
              "$\\displaystyle \\sum_{n=0}^\\infty \\, \\left(\\frac{\\;1\\;}{\\;2\\;}\\right)^{\\!\\!n} \\cos \\frac{n\\pi}{6}$",
              fontdict=linedic,clip_on=False)
f3a = ax.text(-4,base+height,
              "$\\displaystyle \\sum_{n=0}^\\infty \\, \\left(\\frac{\\;1\\;}{\\;2\\;}\\right)^{\\!\\!n} \\cos 0$",
              fontdict=linedic,clip_on=False)
f3b = ax.text( 0,base+height,
              "$= 1+\\frac{\\;1\\;}{\\;2\\;}+\\frac{\\;1\\;}{\\;4\\;}+\\frac{\\;1\\;}{\\;8\\;}+\\frac{\\;1\\;}{\\;16\\;}+\\frac{\\;1\\;}{\\;32\\;}+\\cdots$",
              fontdict=linedic,clip_on=False)
a3a =('ax.text(0,base-height/2,'
      '"$\\displaystyle =\\\\frac{{\\;1\\;}}{{\\;1-\\\\frac{{1}}{{2}}\\;}}$",'
      'fontdict=linedic,clip_on=False,alpha={})')
a3b =('ax.text(3.5,base/2-height/2,'
      '"無限等比級数\\nの公式と一致！",'
      'fontsize=14,alpha={})')
f4a = ax.text(-4,base+height,
              "$\\displaystyle \\sum_{n=0}^\\infty \\, \\left(\\frac{\\;1\\;}{\\;2\\;}\\right)^{\\!\\!n} \\cos\\pi$",
              fontdict=linedic,clip_on=False)
f4b = ax.text( 0,base+height,
              "$= 1-\\frac{\\;1\\;}{\\;2\\;}+\\frac{\\;1\\;}{\\;4\\;}-\\frac{\\;1\\;}{\\;8\\;}+\\frac{\\;1\\;}{\\;16\\;}-\\frac{\\;1\\;}{\\;32\\;}+\\cdots$",
              fontdict=linedic,clip_on=False)
a4a =('ax.text(0,base-height/2,'
      '"$\\displaystyle =\\\\frac{{\\;1\\;}}{{\\;1-(-\\\\frac{{1}}{{2}})\\;}}$",'
      'fontdict=linedic,clip_on=False,alpha={})')

#数式をLATEX変換する関数
def show_line(y,line,alpha=1.0):
    img = []
    line = "$\\displaystyle " + line + "$" 
    text = ax.text(0,y,line,alpha=alpha,fontdict=linedic)
    img += [text]
    bbox = text.get_window_extent().transformed(ax.transData.inverted())
    right = bbox.x1
    return img,right
        
#数式が徐々に現れる関数
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 = ax.pcolormesh(X,Y,C,cmap=cmap,zorder=10)
    return img

#枠外の与式などを付け加える関数
def add_objs_to_imgs(imgs,img,objs):
    for obj in objs:
        img += [obj]
    imgs.append(img)

#最初の１行目を表示する関数
def show_first_line(imgs,line,objs):
    
    i = 0
    while True:
        img,right = show_line(base+height/2,line)
        X,Y = np.mgrid[-2:right+4:delta,height/2:3*height/2:height/2]
        img += [appearing(i,X,Y)]
        add_objs_to_imgs(imgs,img,objs)
        i += 1
        if 2*i+20 > len(X)-1:
            break

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

    for i in range(21):
        alpha = max((10-i)/10,0)
        img1,right1 = show_line(base+height/2+height*i/40,line1,alpha)
        img2,right2 = show_line(base+height*i/40,line2)
        img = img1 + img2
        add_objs_to_imgs(imgs,img,objs)

#数式を次々表示させる関数
def load_script(imgs,lines,objs):
    
    for i in range(20):
        imgs.append(objs)
    show_first_line(imgs,lines[0],objs)
    
    N = len(lines)
    for k in range(N-2):
        show_second_line(imgs,lines[k],lines[k+1],objs)
        go_next_line(imgs,lines[k],lines[k+1],objs)
        
    show_second_line(imgs,lines[N-2],lines[N-1],objs)
    img = imgs[-1]
    for i in range(20):
        imgs.append(img)

#お気に入りのイージング関数
def easing(x):
    if x<0.5:
        return 2*x**2
    else:
        return 1-(-2*x+2)**2/2
    
#３連続コメント
def triple_text(imgs,text):
    
    for k in range(90):
        if   k<21:
            alpha1 = easing(k/20)
            alpha2 = 0
            alpha3 = 0
        elif k<42:
            alpha1 = 1
            alpha2 = easing((k-21)/20)
            alpha3 = 0
        elif k<63:
            alpha1 = 1
            alpha2 = 1
            alpha3 = easing((k-42)/20)
        else:
            alpha1 = 1
            alpha2 = 1
            alpha3 = 1
        img =  [fig.text(0.5,0.5," ",bbox=txtbgdic)]
        img += [fig.text(0.5,0.66,text[0],alpha=alpha1,fontdict=txtdic)]
        img += [fig.text(0.5,0.50,text[1],alpha=alpha2,fontdict=texdic)]
        img += [fig.text(0.5,0.33,text[2],alpha=alpha3,fontdict=txtdic)]
        imgs.append(img)

#補足説明
def appendix(imgs,objs):

    for k in range(70):
        if   k<21:
            alpha1 = easing(k/20)
            alpha2 = 0
        elif k<42:
            alpha1 = 1
            alpha2 = easing((k-21)/20)
        else:
            alpha1 = 1
            alpha2 = 1
        obj1 = eval(objs[0].format(alpha1))
        obj2 = eval(objs[1].format(alpha2))
        img = imgs[-(1+k)] + [obj1,obj2]
        imgs.append(img)

#上演
imgs=[]
triple_text(imgs,text1)
triple_text(imgs,text2)
appendix   (imgs,[a0a,a0b])
triple_text(imgs,text3)
load_script(imgs,line1,[f1])
appendix   (imgs,[a1a,a1b])
triple_text(imgs,text4)
load_script(imgs,line2,[f2])
triple_text(imgs,text5)
triple_text(imgs,text6)
load_script(imgs,line3,[f3a,f3b])
appendix   (imgs,[a3a,a3b])
triple_text(imgs,text7)
load_script(imgs,line4,[f4a,f4b])
appendix   (imgs,[a4a,a3b])

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