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

#全事象の準備
masus = [(x,y) for x in range(1,4) for y in range(1,4)]
cases = [v for v in itertools.combinations(masus,3)]

#各場合のポイント集計
points,aligns = [],[]
for case in cases:
    point = 0
    align = []
    for i in range(2):
        if case[0][i] == case[1][i] == case[2][i]:
            point += 3
            align += [(i,case[0][i])]
        elif case[0][i] == case[1][i]:
            point += 1
            align += [(i,case[0][i])]
        elif case[1][i] == case[2][i]:
            point += 1
            align += [(i,case[1][i])]
        elif case[2][i] == case[0][i]:
            point += 1
            align += [(i,case[2][i])]
    points.append(point)
    aligns.append(align)

#紙の準備
fig = plt.figure()
fig.suptitle("2019一橋大数学　前期第５問",
             color="0.5",x=0.97,y=0.05,ha="right",va="center")
gs  = fig.add_gridspec(3,14,height_ratios=(5,1,1),wspace=0,
                       left=0.2,right=0.8,bottom=0.15,top=0.95)
line1 = fig.add_artist(lines.Line2D([0.0,1.0],[0.460,0.460],color="0.8",linewidth=0.8))
line2 = fig.add_artist(lines.Line2D([0.1,0.9],[0.275,0.275],color="0.8",linewidth=1.2))

#上部の準備
waku = fig.add_subplot(gs[0,:])
waku.set_xlim(0,4)
waku.set_ylim(0,4)
waku.set_aspect("equal")
waku.axis("off")
wakudic = dict(color="0.5",linewidth=2)
for k in np.linspace(0.5,3.5,4):
    waku.plot([0.5,3.5],[k,k],**wakudic)
    waku.plot([k,k],[0.5,3.5],**wakudic)

#カウンタークラス
class Counter:
    
    mojidic = dict(ha="center",color="0.4",
                   fontsize=10,fontweight=600)
    sujidic = dict(clip_on=True,fontsize=18,color="0.3",
                   ha="center",va="center",fontfamily="Ink Free")
    
    def __init__(self,row,col,moji):
        self.ax1 = fig.add_subplot(gs[row,col])
        self.ax2 = fig.add_subplot(gs[row,col+1])
        for ax in [self.ax1,self.ax2]:
            ax.set_xlim(0,3/4)
            ax.set_ylim(0,1)
            ax.set_aspect("equal")
            ax.tick_params(labelbottom=False,bottom=False,
                           labelright=False,right=False,
                           labelleft=False,left=False,
                           labeltop=False,top=False)
            ax.spines[:].set_color("0.7")
        self.ax1.spines["right"].set_visible(False)
        self.ax2.spines["left"].set_linestyle("--")
        for i in range(1,10):
            self.ax1.text(3/8,i+1/2,str(i),**self.sujidic)
        for i in range(100):
            self.ax2.text(3/8,i+1/2,str(i%10),**self.sujidic)
        if moji[-1] == "P":
            self.title = self.ax2.set_title(moji,loc="left",**self.mojidic)
        else:
            self.title = fig.text(0.5,0.11,moji,**self.mojidic)
            
    def count_up(self,upto,a):
        self.ax2.set_ylim(upto-1+a,upto+a)
        if upto%10==0:
            y = upto//10
            self.ax1.set_ylim(y-1+a,y+a)
            
    def blink(self,t):
        if t%4 < 2:
            self.title.set_alpha(0)
        else:
            self.title.set_alpha(1)
 
#下部の準備
Counters = [Counter(1,4*i,str(i)+" P") for i in range(4)]
Counters+= [Counter(2,6,"全事象 COUNT")]
whiteseal = plt.Rectangle((0.05,0.10),0.9,0.35,
                          fill=True,color='w',zorder=10,
                          transform=fig.transFigure,figure=fig)
fig.patches.append(whiteseal)

#お気に入りのイージング関数
def easing(x):
    if x < 0.5:
        return 2*x**2
    else:
        return 1-(-2*x+2)**2/2

#線を引くための関数
linedic = dict(color="lightsteelblue",linewidth=5,solid_capstyle="round")
Lines = waku.plot([],[],**linedic) + waku.plot([],[],**linedic)
def draw_line(i1i2,a):
    i1,i2 = i1i2
    if a==0:
        Lines[i1].set_alpha(1)

    if i1==0:
        Lines[i1].set_data([0.4,0.4+3.2*a],[4-i2,4-i2])
    else:
        Lines[i1].set_data([i2,i2],[3.6-3.2*a,3.6])

#丸を並べるための関数
marudic = dict(s=600,c='none',marker='o',linewidths=4,edgecolors='slategrey')
Marus = waku.scatter([],[],**marudic)
def arrange(case):
    xys = []
    for i in range(3):
        x = case[i][1]
        y = 4-case[i][0]
        xys.append((x,y))
    return xys

#台本
duration_dict = {0:20,1:30,2:40,3:30}
durations = [100]+[duration_dict[p] for p in points]+[80]
endings = np.cumsum(durations)

#本番
def show_case(k,t):
    remaining = duration_dict[points[k]]-t-1
    
    if t==0:
        xys = arrange(cases[k])
        Marus.set_offsets(xys)
    
    nL = len(aligns[k])
    if t<5:
        a = easing(t/4)
        Marus.set_alpha(a)
    elif t<15 and nL>=1:
        a = easing((t-5)/9)
        draw_line(aligns[k][0],a)
    elif t<25 and nL>=2:
        a = easing((t-15)/9)
        draw_line(aligns[k][1],a)
    elif remaining<5:
        a = remaining/4
        for obj in [Marus]:
            obj.set_alpha(a)
        for align in aligns[k]:
            Lines[align[0]].set_alpha(a)
            
    if t<10:
        a = easing(t/9)
        Counters[-1].count_up(k+1,a)
    elif remaining<15 and remaining>=5:
        p = 14-remaining
        a = easing(p/9)
        upto = points[:k+1].count(points[k])
        Counters[points[k]].count_up(upto,a)
        Counters[points[k]].blink(p)
    elif remaining == 4:
        Counters[points[k]].title.set_alpha(1)

#コメントを準備    
commentdic = dict(ha="center",va="center",
                  fontsize=16,linespacing=2,
                  fontfamily="HGMaruGothicMPRO",zorder=20)
comment = fig.text(0.5,0.275,"３×３のマスに\n◯ を３つ\n描いてきます",**commentdic)

#開幕
def opening(t):
    if t<30:
        pass
    elif t<35:
        a = (t-30)/4
        comment.set_alpha(1-a)
    elif t==35:
        comment.set_text("< ルール >\n・１列に ◯ が３つ並ぶ → ３点加算\n・１列に ◯ が２つ並ぶ → １点加算")
    elif t<40:
        a = (t-35)/4
        comment.set_alpha(a)
    elif t<70:
        pass
    elif t<75:
        a = (t-70)/4
        comment.set_alpha(1-a)
    elif t==75:
        comment.set_text("レッツ・ゴー！")
    elif t<80:
        a = (t-75)/4
        comment.set_alpha(a)
    elif t<90:
        pass
    elif t>=90:
        a = (t-90)/9
        for obj in [comment,whiteseal]:
            obj.set_alpha(1-a)

#閉幕
def ending(t):
    if t==50:
        comment.set_text("END")
        comment.set_position((0.5,0.5))
        comment.set_color("w")
        comment.set_bbox(dict(facecolor="k",pad=300))
        comment.set_alpha(1)
        
#上演
def update(i):
    
    if i < endings[0]:
        opening(i)
        
    elif i < endings[-2]:
        k = np.argmin(~(i<endings[1:]))
        show_case(k,i-endings[k])
        
    else:
        ending(i-endings[-2])

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