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

#設定項目
linedic={
"va" : "baseline",
"fontfamily" : "Meiryo",
}
tabledic={
"size": 16,
"color" : "w",
"weight" : 600,
"ha" : "center",
"va" : "center",
"fontfamily" : "Meiryo",
}
phrasedic={
"size" : 20,
"color" : "w",
"zorder" : 100,
"ha" : "center",
"va" : "center",
"visible" : False,
"linespacing" : 2.4,
"fontfamily" : "Meiryo",
}
colors = [(1,1,1,alpha) for alpha in np.arange(0,11)/10]
cmap = ListedColormap(colors)
base = 0.25
ew = 2.0

#紙の準備
fig = plt.figure(facecolor="k")
fig.canvas.draw()
fig.suptitle("2016京大数学　理系第２問",
             color="0.5",ha="right",x=0.96,y=0.96)
gs = fig.add_gridspec(1,20)

#タイトルの準備
case_title = fig.text(0.50,0.82,"$\\quad$以下のように場合分けします",
                      size=14,color="w",ha="center",fontdict=linedic)
titles = [
"【A】両方とも３以上の素数の場合",
"【B】片方が２で、もう一方が５以上の素数の場合",
"【C】両方とも２の場合",
"【D】片方が２で、もう一方が３の場合"
]

#白背景
axw = fig.add_subplot(gs[0,2:18],visible=False)
axw.set_xlim(0,8)
axw.set_ylim(0,4)
axw.set_aspect("equal")
axw.axis("off")

#解説パートの４行
line0 = axw.text(0.5,3+base,"$p^q+q^p$",size=18,fontdict=linedic)
line1 = axw.text(2.3,3+base,"",size=18,fontdict=linedic)
line2 = axw.text(2.3,2+base,"",size=18,fontdict=linedic)
line3 = axw.text(0.0,1+base,"",size=13,fontdict=linedic)
line4 = axw.text(4.0,0+base,"",size=15,weight=600,ha="center",fontdict=linedic)
LINES = [line0,line1,line2,line3,line4]

#各行を隠す
pcms = []
for i in range(4):
    X,Y = np.mgrid[-2:10:0.1,3-i:4-i+0.1:1]
    C = np.ones((len(X)-1,1))
    C[0] = 0
    pcm = axw.pcolormesh(X,Y,C,cmap=cmap,zorder=10)
    pcms.append(pcm)
def reset_pcms():
    for i in range(4):
        C = np.ones((len(X)-1,1))
        C[0] = 0
        pcms[i].set_array(C)
        
#解説パートの準備
kaisetsu1 = [
"$=\\,$奇数$\\quad+\\,$奇数    ", #おまけを見せるために、わざと末尾に空白
"$=\\,$奇数$\\,+\\,$奇数$\\,=\\,$偶数",
"よって、$p^q+q^p\\;$は偶数かつ２でもないので",
"素数じゃない！"
]
kaisetsu2 = [
"$\\equiv\\,(-1)\\quad\\,+(\\pm{1})^2$",
"$\\equiv\\,(-1)+1\\,\\equiv\\,0$",
"よって、$p^q+q^p\\;$は３の倍数かつ３でもないので",
"素数じゃない！"
]
kaisetsu3 = [
"$=\\,2^2+2^2\\,=\\,4+4$",
"$=\\,8$",
"",
"素数じゃない！"
]
kaisetsu4 = [
"$=\\,2^3+3^2\\,=\\,8+9$",
"$=\\,17$",
"",
"素数！"
]
KAISETSU = [kaisetsu1,kaisetsu2,kaisetsu3,kaisetsu4]
kaisetsu1hosoku = [axw.text(4.0,4.4,"⇒ $p$ は奇数,",alpha=0,size=14,ha="right",fontdict=linedic)]
kaisetsu1hosoku+= [axw.text(4.0,4.4," $q$ も奇数",alpha=0,size=14,ha="left",fontdict=linedic)]
kaisetsu1komoji = [axw.text(4.0,3.7,"奇数",visible=False,size=10,fontdict=linedic)]
kaisetsu1komoji+= [axw.text(6.2,3.7,"奇数",visible=False,size=10,fontdict=linedic)]
kaisetsu2hosoku = [axw.text(0.0,4.4," mod３",alpha=0,size=14,color="w",fontdict=linedic,
                            bbox=dict(alpha=0,fc="0.5",lw=0,boxstyle="round"))]
kaisetsu2hosoku+= [axw.text(1.6,4.4,"で考えると",alpha=0,size=14,fontdict=linedic)]
kaisetsu2hosoku+= [axw.text(3.8,4.4,"$p=2\\equiv-1$,",alpha=0,size=14,fontdict=linedic)]
kaisetsu2hosoku+= [axw.text(6.4,4.4,"$q\\equiv\\pm1$",alpha=0,size=14,fontdict=linedic)]
kaisetsu2komoji =[axw.text(4.15,3.5,"奇数",visible=False,size=10,fontdict=linedic)]

#各行の長さを計測
def prepare_kaisetsu(kaisetsu):
    xends = []
    for i in range(4):
        LINES[i+1].set_text(kaisetsu[i])
        bbox = LINES[i+1].get_window_extent().transformed(axw.transData.inverted())
        xends.append(int(bbox.x1*5)+10)
        kends = np.cumsum(xends)
    return kends
kends1 = prepare_kaisetsu(kaisetsu1)
kends2 = prepare_kaisetsu(kaisetsu2)
kends3 = prepare_kaisetsu(kaisetsu3)
kends4 = prepare_kaisetsu(kaisetsu4)
KENDS = [kends1,kends2,kends3,kends4]

#黒背景
axk = fig.add_subplot(gs[0,3:17],fc="k")
axk.set_xlim(0,8)
axk.set_ylim(0,6)
axk.set_aspect("equal")
for spine in axk.spines.values():
    spine.set_color("w")
    spine.set_linewidth(ew)
axk.set_xticks([])
axk.set_yticks([])

#表の準備
for i in range(3):
    if i<2:
        xend = 8.0
        yend = 0.0
    else:
        xend = 4.0
        yend = 3.0
    axk.plot([0,xend],[4.5-1.5*i,4.5-1.5*i],color="w",lw=ew)
    axk.plot([2.0+2.0*i,2.0+2.0*i],[yend,6],color="w",lw=ew)
axk.text(1,0.75,"$p\\,\\geqq\\,5$",fontdict=tabledic)
axk.text(1,2.25,"$p=3$",fontdict=tabledic)
axk.text(1,3.75,"$p=2$",fontdict=tabledic)
axk.text(3,5.25,"$q=2$",fontdict=tabledic)
axk.text(5,5.25,"$q=3$",fontdict=tabledic)
axk.text(7,5.25,"$q\\,\\geqq\\,5$",fontdict=tabledic)
cellA = [axk.text(6,1.50,"A",fontdict=tabledic)]
cellB = [axk.text(3,0.75,"B",fontdict=tabledic)]
cellB+= [axk.text(7,3.75,"B",fontdict=tabledic)]
cellC = [axk.text(3,3.75,"C",fontdict=tabledic)]
cellD = [axk.text(3,2.25,"D",fontdict=tabledic)]
cellD+= [axk.text(5,3.75,"D",fontdict=tabledic)]
CELL = [cellA,cellB,cellC,cellD]
hatchA = [patches.Rectangle((4,0.0),4,3.0,fill=False,ec="w",hatch="//")]
hatchB = [patches.Rectangle((2,0.0),2,1.5,fill=False,ec="w",hatch="//")]
hatchB+= [patches.Rectangle((6,3.0),2,1.5,fill=False,ec="w",hatch="//")]
hatchC = [patches.Rectangle((2,3.0),2,1.5,fill=False,ec="w",hatch="//")]
hatchD = [patches.Rectangle((2,1.5),2,1.5,fill=False,ec="w",hatch="//")]
hatchD+= [patches.Rectangle((4,3.0),2,1.5,fill=False,ec="w",hatch="//")]
HATCH = [hatchA,hatchB,hatchC,hatchD]
for hatch in hatchA+hatchB+hatchC+hatchD:
    axk.add_patch(hatch)
    hatch.set_visible(False)

#冒頭と結末
phrase1 = fig.text(0.5,0.5,"$p^q+q^p$ が素数となる\n素数 $p,q$\nを求めたい",fontdict=phrasedic,
                   bbox=dict(color="black",pad=300))
phrase2 = fig.text(0.5,0.5,"以上！",fontdict=phrasedic,
                   bbox=dict(color="black",pad=300))

#表を点滅させる関数
def blink_hatch(k,hatch):
    if k%10 == 0:
        for obj in hatch:
            obj.set_visible(True)
    elif (k+5)%10 == 0:
        for obj in hatch:
            obj.set_visible(False)

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

#解説パートを仕切る関数
def show_kaisetsu(k,kaisetsu,kends):
    n = np.argmin(~(k<kends))
    kends = np.append(0,kends)
    appearing(k-kends[n],pcms[n])
    
#お気に入りのイージング関数
def easing(x):
    if x<0.0:
        return 0
    elif x<0.5:
        return 2*x**2
    elif x<1.0:
        return 1-(-2*x+2)**2/2
    else:
        return 1

#セクションを判定する関数
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,"S1":60,
           "B1":66,"A1":60,"W1":kends1[-1],"N1":50,
           "B2":66,"A2":90,"W2":kends2[-1],"N2":50,
           "B3":66,"W3":kends3[-1],"N3":50,
           "B4":66,"W4":kends4[-1],"N4":50,
           "E1":60,"P2":30}
endings = np.cumsum(list(periods.values()))

#上演
def update(i):
    
    key,k = section(i)
    n = int(key[-1])-1
    
    if "B" in key:
        blink_hatch(k,HATCH[n])
    elif "W" in key:
        show_kaisetsu(k,KAISETSU[n],KENDS[n])
    elif "A1" == key:
        alpha1 = easing(k/29)
        alpha2 = easing((k-30)/29)
        kaisetsu1hosoku[0].set_alpha(alpha1)
        kaisetsu1hosoku[1].set_alpha(alpha2)
    elif "A2" == key:
        alpha1 = easing(k/29)
        alpha2 = easing((k-30)/29)
        alpha3 = easing((k-60)/29)
        kaisetsu2hosoku[0].get_bbox_patch().set_alpha(alpha1)
        kaisetsu2hosoku[0].set_alpha(alpha1)
        kaisetsu2hosoku[1].set_alpha(alpha1)
        kaisetsu2hosoku[2].set_alpha(alpha2)
        kaisetsu2hosoku[3].set_alpha(alpha3)
    
    if k==0:
        if "P1" == key:
            phrase1.set_visible(True)
        elif "S1" == key:
            phrase1.set_visible(False)
        elif "B" in key or "E" in key:
            fig.set_facecolor("k")
            axw.set_visible(False)
            axk.set_visible(True)
            case_title.set_text(titles[n])
            case_title.set_color("w")
            if "B2" in key or "B3" in key or "B4" in key:
                for cell in CELL[n-1]:
                    cell.set_visible(False)
                for hatch in HATCH[n-1]:
                    x = hatch.get_x()
                    y = hatch.get_y()
                    w = hatch.get_width()
                    h = hatch.get_height()
                    axk.plot([x,x+w],[y,y+h],color="w",lw=ew)
                    axk.plot([x,x+w],[y+h,y],color="w",lw=ew)
            elif "E1" == key:
                case_title.set_text("$\\quad$～ まとめ ～")
                for cell in cellD:
                    cell.set_text("OK")
        elif "A" in key or "W" in key:
            fig.set_facecolor("w")
            axk.set_visible(False)
            axw.set_visible(True)
            reset_pcms()
            prepare_kaisetsu(KAISETSU[n])
            case_title.set_color("k")
            if "A1" == key:
                for kmj in kaisetsu1komoji:
                    kmj.set_visible(True)
            elif "A2" == key:
                for hsk in kaisetsu1hosoku:
                    hsk.set_visible(False)
                for kmj in kaisetsu1komoji:
                    kmj.set_visible(False)
                for kmj in kaisetsu2komoji:
                    kmj.set_visible(True)
            elif "W3" == key:
                for hsk in kaisetsu2hosoku:
                    hsk.set_visible(False)
                for kmj in kaisetsu2komoji:
                    kmj.set_visible(False)
        elif "P2" == key:
            phrase2.set_visible(True)

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