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

#初期設定
plt.rcParams["mathtext.fontset"]="stix"
linedic={
"fontsize" : 18,
"verticalalignment" : "baseline",
}
phrasedic={
"ha" : "center",
"va" : "center",
"size" : 20,
"color" : "white",
"zorder" : 100,
"visible" : False,
"linespacing": 2.4,
"fontfamily": "Meiryo",
}
base=0.6
colors = [(1,1,1,alpha) for alpha in np.arange(0,11)/10]
cmap = ListedColormap(colors)
seq_pcm = []
seq_prs = []

#紙の準備
fig = plt.figure()
fig.suptitle("2018京大数学　理系第２問　文系第３問",
             color="0.5",ha="right",x=0.96,y=0.96)

#１枚目を書きます
fig.subplots_adjust(left=0.04,right=0.96)
ax1 = fig.add_subplot(131)
ax2 = fig.add_subplot(132)
ax3 = fig.add_subplot(133)
for ax,mod in zip([ax1,ax2,ax3],[0,1,2]):
    ax.set_xlim(0,8)
    ax.set_ylim(0,8)
    ax.set_aspect("equal")
    ax.axis("off")
    ax.text(1.1,6+base,"$n\\equiv"+str(mod)+"$",fontdict=linedic)
    ax.text(4.0,6+base,"のとき",fontsize=14,fontfamily="Meiryo")
    ax.text(1.2,4+base,"$n^3-7n+9$",fontdict=linedic)
    ax.text(0.0,2+base,"$\\equiv"+str(mod)+"^3-7\\times"+str(mod)+"+9$",fontdict=linedic)
    ax.text(0.0,0+base,"$\\equiv"+str(mod**3-7*mod+9)+"\\equiv 0$",fontdict=linedic)
    for i in range(4):
        X,Y = np.mgrid[-2:10:0.1,8-(i+1)*2:8-i*2+0.1:2]
        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)
ax3.text(7.7,-5.5,"※ mod3 で考えています",
         color="white",ha="right",fontsize=12,fontfamily="Meiryo",
         bbox=dict(boxstyle="round",color="0.5",pad=0.5))

#２枚目を書きます
def set_2nd_stage():
    fig.subplots_adjust(left=0.3,right=0.7)
    ax4 = fig.add_subplot(111)
    ax4.set_xlim(0,12)
    ax4.set_ylim(0,8)
    ax4.set_aspect("equal")
    ax4.axis("off")
    ax4.text(2.5,6+base,"$n^3-7n+9=3$",fontdict=linedic)
    ax4.text(2.5,4+base,"$n^3-7n+6=0$",fontdict=linedic)
    ax4.text(0.5,2+base,"$(n-1)(n-2)(n+3)=0$",fontdict=linedic)
    ax4.text(3.3,0+base,"$n=1,2,-3$",fontdict=linedic)
    ax4.set_title("最後に三次方程式を解きます",
                  pad=20,fontsize=12,fontfamily="Meiryo")
    for i in range(4):
        X,Y = np.mgrid[-2:14:0.1,8-(i+1)*2:8-i*2+0.1:2]
        C = np.ones(len(X)-1)
        C[0] = 0
        C = np.reshape(C,(len(X)-1,1))
        pcm = ax4.pcolormesh(X,Y,C,cmap=cmap,zorder=10)
        seq_pcm.append(pcm)

#減らずグチたち
phrases = []
phrases.append("""
$n^3-7n+9$ を素数にしなさい
（ ただし、$n$ は整数 ）
""")
phrases.append("""
困ったなあ...
""")
phrases.append("""
mod3 参上！
""")
phrases.append("""
というわけで、結局
$n^3-7n+9$ は３の倍数です
""")
phrases.append("""
３の倍数である素数は
３しかありません
""")
phrases.append("""
（答）
$n=1,\\,2,\\,-3$
""")
phrases.append("""
おしまい
(´・ω・`)
""")
for phrase in phrases:
    prs = fig.text(0.5,0.5,phrase,fontdict=phrasedic,
                   bbox=dict(color="black",pad=300))
    seq_prs.append(prs)

#数式が徐々に現れる関数
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)

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

#台本の構成
periods = {"P1":60,"P2":30,"P3":30,
           "S1":10,"A01":50,"A02":50,"A03":50,"A04":40,
           "S2":10,"A05":50,"A06":50,"A07":50,"A08":40,
           "S3":10,"A09":50,"A10":50,"A11":50,"A12":50,
           "P4":50,"P5":50,
           "S4":10,"A13":70,"A14":70,"A15":70,"A16":70,
           "P6":60,"P7":30}
endings = np.cumsum(list(periods.values()))

#上演
def update(i):
    
    key,k = section(i)
    if "A" in key:
        appearing(k,seq_pcm[int(key[1:])-1])
            
    if k == 0:
        if "P" in key:
            seq_prs[int(key[1])-2].set_visible(False)
            seq_prs[int(key[1])-1].set_visible(True)
        elif key=="S1":
            seq_prs[2].set_visible(False)
        elif key=="S2":
            for txt in ax1.texts:
                txt.set_color("0.5")
            ax1.plot([8.8,8.8],[0,8],color="0.5",clip_on=False,lw=1)
        elif key=="S3":
            for txt in ax2.texts:
                txt.set_color("0.5")
            ax2.plot([8.8,8.8],[0,8],color="0.5",clip_on=False,lw=1)
        elif key=="S4":
            for ax in fig.axes:
                ax.remove()
            seq_prs[4].set_visible(False)
            set_2nd_stage()
            
mov = ani.FuncAnimation(fig,update,endings[-1],interval=100)
plt.show()