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.suptitle("2014阪大数学　理系第３問",
             color="0.5",ha="right",x=0.96,y=0.96)
fig.subplots_adjust(left=0.0,right=1.0)
gs = gridspec.GridSpec(2,31)
ax1u = fig.add_subplot(gs[0, 1:10],visible=False)
ax1d = fig.add_subplot(gs[1, 1:11])
ax2u = fig.add_subplot(gs[0,11:20])
ax2d = fig.add_subplot(gs[1,11:21])
ax3u = fig.add_subplot(gs[0,21:30],visible=False)
ax3d = fig.add_subplot(gs[1,21:31])

#グラフの準備
for ax in [ax1u,ax2u,ax3u]:
    ax.set_xlim(-1,11)
    ax.set_ylim(-0.1,1.1)
    ax.axes.xaxis.set_visible(False)
    ax.axes.yaxis.set_visible(False)
    ax.spines["top"].set_visible(False)
    ax.spines["right"].set_visible(False)
    ax.spines["left"].set_position(("data",0))
    ax.spines["bottom"].set_position(("data",0))
    ax.spines["left"].set_color("0.5")
    ax.spines["bottom"].set_color("0.5")
    ax.text(6,0.15,"$\\cdots\\cdots$",ha="center",color="0.3")
    for n in [1,2,3,4,9,10]:
        ax.fill([n-1,n,n,n-1],[0,0,1/np.sqrt(n),1/np.sqrt(n)],
                fill=False,ec="0.3",hatch="////",zorder=10)
    for n in [1,2,3,4]:
        ax.text(n,-0.1,"$"+str(n)+"$",ha="center",color="0.3",fontsize=8)
    ax.annotate("$39999$",xy=(9,0),ha="center",color="0.3",fontsize=8,xytext=(6.5,-0.15),
                arrowprops=dict(arrowstyle = "->",color="0.5",connectionstyle="angle,angleA=0,angleB=90"))
    ax.annotate("$40000$",xy=(10,0),ha="center",color="0.3",fontsize=8,xytext=(7.0,-0.25),
                arrowprops=dict(arrowstyle = "->",color="0.5",connectionstyle="angle,angleA=0,angleB=90"))
for n in [1,2,3,4]:
    ax2u.plot([n+0.3,n**1.1+1.5],[1/np.sqrt(n),1/np.sqrt(n)],color="0.5",linestyle="dotted",linewidth=1)
    if n == 1:
        ax2u.text(2.5,1,"$1$",va="center",color="0.3",fontsize=8)
    else:
        ax2u.text(n**1.1+1.8,1/np.sqrt(n),"$\\frac{1}{\\sqrt{"+str(n)+"}}$",va="baseline",color="0.3",fontsize=8)

#数式のスペースの準備
for ax in [ax1d,ax2d,ax3d]:
    ax.set_xlim(0,8)
    ax.set_ylim(0,8)
    ax.axis("off")

#曲線を描きます
x1 = np.linspace(-0.9,11,120)
x3 = np.linspace( 0.1,11,110)
y1 = 1/np.sqrt(x1+1)
y3 = 1/np.sqrt(x3)
ax1u.plot(x1,y1,color="0.5",linewidth=0.5)
ax3u.plot(x3,y3,color="0.5",linewidth=0.5)
x1 = np.append(np.linspace(0,9,91),10)
x3 = np.append(0,np.linspace(1,10,91))
y1 = np.append(1/np.sqrt(x1[:-1]+1),1/np.sqrt(10))
y3 = np.append(1,1/np.sqrt(x3[1:]))
ax1u.fill_between(x1,y1,ec="b",fc=(0,0,1,0.2),zorder=20)
ax3u.fill_between(x3,y3,ec="r",fc=(1,0,0,0.2),zorder=20)
ax1u.text(5,0.5,"$y=\\frac{1}{\\sqrt{x+1}}$",color="b",fontsize=14,usetex=True)
ax3u.text(6,0.5,"$y=\\frac{1}{\\sqrt{x}}$",color="r",fontsize=14,usetex=True)

#フォント辞書たち
linedic={
"usetex" : True,
"clip_on" : True,
"fontsize" : 12,
"verticalalignment" : "baseline",
}
phrasedic={
"ha" : "center",
"va" : "center",
"size" : 20,
"color" : "white",
"zorder" : 100,
"visible" : False,
"linespacing": 2.4,
"fontfamily": "Meiryo",
}
futogodic={
"ha" : "center",
"size" : 32,
"color" : "0.3",
"alpha" : 0,
}

#数式たち
lines1=[
r"= \int_{0}^{39999} \!\!\!\!\!\! \frac{1}{\sqrt{x+1}} \, dx + \frac{1}{\sqrt{40000}}",
r"= \left[\, 2\sqrt{x+1} \,\,\right]_{0}^{39999} \!\! + \frac{1}{200}",
r"= 2\,(\sqrt{40000}-\sqrt{1}\,) + \frac{1}{200}",
r"= 398.005",
]
lines2=[
r"= \frac{1}{\sqrt{1}} + \frac{1}{\sqrt{2}} + \frac{1}{\sqrt{3}} + \frac{1}{\sqrt{4}} +",
r"\cdots\cdots + \frac{1}{\sqrt{39999}} + \frac{1}{\sqrt{40000}}",
r"= \sum_{n=1}^{40000} \frac{1}{\sqrt{n}}",
]
lines3=[
r"= \frac{1}{\sqrt{1}} + \int_{1}^{40000} \!\! \frac{1}{\sqrt{x}} \, dx",
r"= 1 + \left[\, 2\sqrt{x} \,\,\right]_{1}^{40000}",
r"= 1 + 2\,(\sqrt{40000}-\sqrt{1}\,)",
r"= 399",
]

#減らずグチたち
phrase1="""
$\\sum_{n=1}^{40000}\\dfrac{1}{\\sqrt{n}}$
の整数部分は？
"""
phrase2="""
えっ？
"""
phrase3="""
つべこべ言わず
積分
しましょう
"""
phrase4="""
$398.005\\,<\\,\\sum_{n=1}^{40000}\\dfrac{1}{\\sqrt{n}}\\,<\\,399$
"""
phrase5="""
（答）
$398$
"""
phrase6="""
見てくれて
サン・キュ・パッ
(*´▽`人)
"""

#カラーマップの準備
colors = [(1,1,1,alpha) for alpha in np.arange(0,11)/10]
cmap = ListedColormap(colors)

#数式を書き込んでいきます
seq_txt = []
seq_pcm = []
for ax,lines,left,lh\
in zip([ax2d,ax1d,ax3d],
       [lines2,lines1,lines3],
       [0.3,0.0,0.6],
       [2.7,2.2,2.2]):
    for i in range(len(lines)):
        latex = "$\\displaystyle " + lines[i] + "$"
        txt = ax.text(left,8-(i+1)*lh+lh*0.4,latex,fontdict=linedic)
        seq_txt.append(txt)
        X,Y = np.mgrid[-2:12:0.1,8-(i+1)*lh:8-i*lh+0.1:lh]
        C = np.ones(len(X)-1)
        C[0] = 0
        C = np.reshape(C,(len(X)-1,1))
        pcm = ax.pcolormesh(X,Y,C,cmap=cmap,zorder=10)
        seq_pcm.append(pcm)

#減らずグチを埋め込んでいきます        
seq_prs = []
for phrase in [phrase1,phrase2,phrase3,
               phrase4,phrase5,phrase6]:
    prs = fig.text(0.5,0.5,phrase,fontdict=phrasedic,
                   bbox=dict(color="black",pad=300))
    seq_prs.append(prs)

#不等号を埋め込んでいきます
futogo1 = fig.text(0.33,0.66,"$<$",fontdict=futogodic)
futogo2 = fig.text(0.65,0.66,"$<$",fontdict=futogodic)

#数式が徐々に現れる関数
def appearing(i,pcm):
    C = np.ones(len(pcm.get_array()))
    C[:2*i+1] = 0
    C[2*i+1:2*i+20] = np.arange(0.05,1.00,0.05)
    pcm.set_array(C)

#台本の構成
periods = [50,30,30,
           20,60,60,60,
           20,60,60,60,60,
           20,60,60,60,60,80,
           50,50,30]
endings = np.cumsum(periods)
n_appearing = [4,5,6,8,9,10,11,13,14,15,16]

#セクションを判定する関数
def section(i): 
    n = np.argmin(~(i<endings))
    if n == 0:
        k = i
    else:
        k = i-endings[n-1]
    return n,k

#上演
def update(i):
    
    n,k = section(i)
    if n in n_appearing:
        appearing(k,seq_pcm[n_appearing.index(n)])
    elif n == 17:
        alpha = max(min((k-20)/20,1),0)
        futogo1.set_alpha(alpha)
        futogo2.set_alpha(alpha)
    
    if i == 0:
        seq_prs[0].set_visible(True)
    elif i == endings[0]:
        seq_prs[0].set_visible(False)
        seq_prs[1].set_visible(True)
    elif i == endings[1]:
        seq_prs[1].set_visible(False)
        seq_prs[2].set_visible(True)
    elif i == endings[2]:
        seq_prs[2].set_visible(False)
    elif i == endings[6]:
        ax1u.set_visible(True)
        for txt in seq_txt[0:2]:
            txt.set_color("0.5")
    elif i == endings[11]:
        ax3u.set_visible(True)
        for txt in seq_txt[3:6]:
            txt.set_color("0.5")
    elif i == endings[16]:
        for txt in seq_txt[7:10]:
            txt.set_color("0.5")
    elif i == endings[-4]:
        seq_prs[3].set_visible(True)
    elif i == endings[-3]:
        seq_prs[3].set_visible(False)
        seq_prs[4].set_visible(True)
    elif i == endings[-2]:
        seq_prs[4].set_visible(False)
        seq_prs[5].set_visible(True)

mov = ani.FuncAnimation(fig,update,endings[-1],interval=100)
plt.show()