import itertools
import numpy as np
import matplotlib.pyplot as plt
import matplotlib.animation as ani
from matplotlib.lines import Line2D

#色の準備
col_bg  = "powderblue"
col_r   = "orangered"
col_k   = "0.2"
col_w   = "wheat"
col_ng  = "crimson"
col_ok  = "dodgerblue"
col_txt = "0.4"

#紙の準備
fig = plt.figure(facecolor=col_bg)
fig.suptitle("2023東大数学　理系第２問　文系第３問",
             color="0.6",x=0.97,y=0.05,ha="right",va="center")
fig.subplots_adjust(top=0.9,bottom=0.2)
ax1 = fig.add_subplot(211)
ax1.set_xlim(-1,12)
ax1.set_ylim(-1,1)
ax1.set_aspect("equal")
ax1.axis("off")
ax2 = fig.add_subplot(212)
ax2.set_xlim(-8,6)
ax2.set_ylim(-2,2)
ax2.set_aspect("equal")
ax2.axis("off")

#玉を準備
tama = ax1.scatter([4,5,6,4,5,6,7,4,5,6,7,8],
                   [-1]*3+[-3]*4+[-5]*5,
                   facecolors=[col_k]*3+[col_r]*4+[col_w]*5,
                   s=400,clip_on=False)

#数え上げ領域の準備
objs = []
batsu1,= ax1.plot([0.5,2.5],[-2.5,-3.5],color=col_ng,linewidth=3,solid_capstyle="round",clip_on=False)
batsu2,= ax1.plot([0.5,2.5],[-3.5,-2.5],color=col_ng,linewidth=3,solid_capstyle="round",clip_on=False)
objs = objs + [batsu1] + [batsu2]
objs +=  ax2.plot([-7,-1],[0,0],color=col_txt)
objs +=  ax2.plot([ 1, 3],[0,0],color=col_txt)
objs += [ax2.text(0,0,"＝",color=col_txt,size=20,ha="center",va="center")]
objs += [ax2.text(-6.0,-.80,"赤",color="w",size=16,ha="center",va="center",fontfamily="HGSoeiKakupoptai",
                  bbox=dict(fc=col_r,boxstyle="circle",linewidth=0,pad=0.4))]
objs += [ax2.text(-6.0,0.75,"黒",color="w",size=16,ha="center",va="center",fontfamily="HGSoeiKakupoptai",
                  bbox=dict(fc=col_k,boxstyle="circle",linewidth=0,pad=0.4))]
objs += [ax2.text(-3.3,-.85,"が隣り合わない",color=col_txt,size=14,ha="center",va="center",fontfamily="HGSoeiKakupoptai")]
objs += [ax2.text(-3.3,0.70,"が隣り合わない",color=col_txt,size=14,ha="center",va="center",fontfamily="HGSoeiKakupoptai")]
nump  =  ax2.text( 1.1,-.85,"for number",color=col_txt,size=20,ha="left",va="center",fontfamily="fantasy")
numq  =  ax2.text( 1.1,0.60,"for number",color=col_txt,size=20,ha="left",va="center",fontfamily="fantasy")
mkup  =  ax2.text( 3.3,0.65,"$\\uparrow$",color=col_ok,size=22,ha="center",va="center")
txok  =  ax2.text(-1,0,"( ｀д´)b ｵｯｹｰ!",color=col_ok,size=24,ha="center",va="center")
objs += [nump,numq,mkup,txok]
for obj in objs:
    obj.set_visible(False)

#減らずグチたち
commentdic={
"ha" : "center",
"va" : "center",
"size" : 20,
"color" : "white",
"zorder" : 10,
"fontfamily": "Meiryo",
"linespacing" : 2,
}
commentbgdic={
"facecolor" : "black",
"pad" : 300,
}
msg_opening1 = "絶対に\n負けられない 戦いが\n\nここにはある"
msg_opening2 = "今は昔\n黒玉３個、赤玉４個、白玉５個\nがありました"
msg_opening3 = "おばあさんは\n赤玉が隣り合わないよう\n一列に並べました"
msg_opening4 = "このとき\n黒玉が隣り合っていたら\nダメです"
msg_opening5 = "赤玉も黒玉も\n隣り合っていなかったら\nOKです！"
msg_opening6 = "今から\n無の時間が\nやってきます"
msg_opening7 = "行くぞ！"
msg_ending1  = "（答）\n$\\dfrac{\\:103\\:}{\\:168\\:}$"
msg_ending2  = "場合の数は\n退屈との戦いです"
msg_ending3  = "-----  END  -----"
comment = fig.text(0.5,0.5,"",fontdict=commentdic,bbox=commentbgdic)

#訂正線
strike_through = Line2D([0.29,0.29],[0.56,0.56],
                        color="w",linewidth=2.5,zorder=11)
fig.add_artist(strike_through)
txt_opening1 = fig.text(0.43,0.47,"数え切る",
                        zorder=11,alpha=0,fontdict=commentdic)

#エンディング用
endings = []
endings +=  ax2.plot([4.8,6.4],[0,0],color=col_txt,clip_on=False)
endings += [ax2.text(4,0,"＝",color=col_txt,size=20,ha="center",va="center",fontfamily="Meiryo")]
endings += [ax2.text(5.6,-.85,str(168),color=col_txt,size=20,ha="center",va="center",fontfamily="fantasy")]
endings += [ax2.text(5.6,0.60,str(103),color=col_txt,size=20,ha="center",va="center",fontfamily="fantasy")]
for obj in endings:
    obj.set_alpha(0)

#まずはパイソン君に数えさせます
l1 = list(range(12))
cases = []
count = 0
for r in itertools.combinations(l1,4):
    l2 = [i for i in l1 if i not in r]
    for k in itertools.combinations(l2,3):
        if 1 not in np.diff(r):
            if 1 not in np.diff(k):
                count += 1
                cases.append([r,k,count])
            else:
                cases.append([r,k,count])

#お気に入りのイージング関数
def easing(x):
    return -(np.cos(np.pi*x)-1)/2

#無の時間を作り出す関数
def show_case(t):
    
    i = min(t,len(cases)-1)
    r = cases[i][0]
    k = cases[i][1]
    
    colors = []
    for j in range(12):
        if j in r:
            color = col_r
        elif j in k:
            color = col_k
        else:
            color = col_w
        colors.append(color)
    tama.set_facecolors(colors)
    
    q = cases[i][2]
    numq.set_text("{:^5}".format(q))
    nump.set_text("{:^5}".format(i+1))
    
    if 1 in np.diff(k):
        batsux = k[np.argmin(np.diff(k)!=1)]
        batsu1.set_data([batsux-0.5,batsux+1.5],[0.5,-0.5])
        batsu2.set_data([batsux-0.5,batsux+1.5],[-0.5,0.5])
        for obj in [batsu1,batsu2]:
            obj.set_visible(True)
        mkup.set_visible(False)
    else:
        for obj in [batsu1,batsu2]:
            obj.set_visible(False)
        if t < len(cases):
            mkup.set_visible(True)
        else:
            mkup.set_visible(False)

#オープニング
def opening(t):
    if t==0:
        comment.set_text(msg_opening1)
    elif t>=20 and t<=40:
        k = easing((t-20)/20)
        strike_through.set_xdata([0.29,0.29+0.27*k])
    elif t>=60 and t<=80:
        k = easing((t-60)/20)
        txt_opening1.set_alpha(k)
    elif t==120:
        comment.set_text(msg_opening2)
        for obj in [strike_through,txt_opening1]:
            obj.set_visible(False)
    elif t==160:
        comment.set_visible(False)
    elif t==200:
        comment.set_text(msg_opening3)
        comment.set_visible(True)
        tama.set_offsets(list(zip(l1,[-3]*12)))
        tama.set_facecolors([col_r,col_k,col_k,col_r,col_w,col_w,col_r,col_k,col_w,col_r,col_w,col_w])
    elif t==240:
        comment.set_visible(False)
    elif t==280:
        comment.set_text(msg_opening4)
        comment.set_visible(True)
        for obj in [batsu1,batsu2]:
            obj.set_visible(True)
    elif t==320:
        comment.set_visible(False)
    elif t==360:
        comment.set_text(msg_opening5)
        comment.set_visible(True)
        for obj in [batsu1,batsu2]:
            obj.set_visible(False)
        tama.set_facecolors([col_r,col_k,col_w,col_r,col_k,col_w,col_r,col_k,col_w,col_r,col_w,col_w])
        txok.set_visible(True)
    elif t==400:
        comment.set_visible(False)
    elif t==440:
        comment.set_text(msg_opening6)
        comment.set_visible(True)
        tama.set_offsets(list(zip(l1,[0]*12)))
        for obj in objs:
            obj.set_visible(True)
        txok.set_visible(False)
    elif t==480:
        comment.set_text(msg_opening7)
    elif t==499:
        show_case(0)
        comment.set_visible(False)

#エンディング
def ending(t):
    if t<40:
        alpha = easing(t/39)
        for obj in endings:
            obj.set_alpha(alpha)
    elif t<80:
        pass
    elif t==80:
        comment.set_text(msg_ending1)
        comment.set_visible(True)
    elif t==120:
        comment.set_text(msg_ending2)
    elif t==160:
        comment.set_text(msg_ending3)

#上演
def update(i):
    if i<500:
        opening(i)
    elif i<7600:
        k = i-500
        show_case(k)
    else:
        k = i-7600
        ending(k)

mov = ani.FuncAnimation(fig,update,7800,interval=100)
plt.show()