import numpy as np
import matplotlib.pyplot as plt
import matplotlib.animation as ani
from matplotlib.lines import Line2D
from matplotlib.gridspec import GridSpec
from matplotlib.colors import ListedColormap
from scipy.spatial.transform import Rotation
from mpl_toolkits.mplot3d.art3d import Poly3DCollection

#紙の準備
fig = plt.figure()
fig.suptitle("2020灘中算数　１日目１０番",
             color="0.5",ha="right",x=0.96,y=0.96)
gs = GridSpec(2,2,height_ratios=[5,2],
              bottom=0.2,hspace=0,left=0.1,right=0.9)
ax2 = fig.add_subplot(gs[1,0])
ax3 = fig.add_subplot(gs[1,1])
for ax in [ax2,ax3]:
    ax.set_xlim(0,8)
    ax.set_ylim(-3,0)
    ax.set_aspect("equal")
    ax.axis("off")
ax1 = fig.add_subplot(gs[0,:],projection='3d',computed_zorder=False)
ax1.view_init(10,-90)
ax1.set_xlim(-2.4,2.6)
ax1.set_ylim(-2.5,2.5)
ax1.set_zlim(-4.0,1.0)
ax1.set_box_aspect((1,1,1))
ax1.axis("off")

#色の準備
a1 = 0.2
a2 = 0.8
fc1 = (0.85,0.65,0.12)
fc2 = (0.55,0.27,0.07)

#回転後の座標を取得する関数
n = 61
rotvec = 2*np.pi*np.array([-3/5,0,4/5])
def rot(p0):
    for i in range(n):
        rvi = rotvec*i/(n-1)
        R = Rotation.from_rotvec(rvi).as_matrix()
        rotxyz = R @ np.array(p0)
        if i==0:
            rotxyzs = rotxyz
        else:
            rotxyzs = np.vstack([rotxyzs,rotxyz])
    return rotxyzs

#回転後の座標
rd = rot([0,0,0])
ra = rot([-3,0,0])
rb = rot([-3,0,-4])
rm = rot([0,0,-4])
rc = rot([3,0,-4])

#回転中の側面を取得する関数
def make_verts(r1,r2):
    verts = []
    for i in range(1,n):
        vert = np.array([r1[i-1],r2[i-1],r2[i],r1[i]])
        verts.append(vert)
    return verts

#回転中の側面
vDA = make_verts(rd,ra)
vAB = make_verts(ra,rb)
vBM = make_verts(rb,rm)
vMC = make_verts(rm,rc)
V = [vDA,vAB,vBM,vMC]

#回転体クラス
class Kaitenmen(Poly3DCollection):
    
    def __init__(self,vs0):
        super().__init__([],fc=fc1+(a1,),clip_on=False)
        self.vs0 = vs0
        ax1.add_collection(self)
        
    def gururi(self,i):
        self.set_verts(self.vs0[:i])
    
    def slider(self,a,slivec):
        slivec = np.array(slivec)
        new_vs = []
        for old_v in self.vs0:
            slider = np.tile(a*slivec,(len(old_v),1))
            new_v = old_v + slider
            new_vs.append(new_v)
        self.set_verts(new_vs)
        
        if a==1:
            self.vs0 = new_vs
    
    def mawaru(self,a):
        A = np.array([[np.cos(a*np.pi)-1,0,2*np.sin(a*np.pi)],
                      [0,-1,0],
                      [-2*np.sin(a*np.pi),0,np.cos(a*np.pi)-1]])
        slivec = A @ np.array([3,0,-4])
        new_vs = []
        for old_v in self.vs0:
            slider = np.tile(slivec,(len(old_v),1))
            new_v = old_v + slider
            new_vs.append(new_v)
        self.set_verts(new_vs)
        
        if a==1:
            self.vs0 = new_vs
    
    def rotate(self,a,rotvec,shaft):
        ar = a*np.array(rotvec)
        R = Rotation.from_rotvec(ar).as_matrix()
        new_vs = []
        for old_v in self.vs0:
            b = np.tile(shaft,(len(old_v),1))
            new_v = (old_v-b) @ R + b
            new_vs.append(new_v)
        self.set_verts(new_vs)
        
    def slirot(self,a,slivec,rotvec,shaft):
        slivec = np.array(slivec)
        ar = a*np.array(rotvec)
        R = Rotation.from_rotvec(ar).as_matrix()
        new_vs = []
        for old_v in self.vs0:
            b = np.tile(shaft,(len(old_v),1))
            slider = np.tile(a*slivec,(len(old_v),1))
            new_v = (old_v-b) @ R + b + slider
            new_vs.append(new_v)
        self.set_verts(new_vs)

#回転体インスタンス
sDA = Kaitenmen(vDA)
sAB = Kaitenmen(vAB)
sBM = Kaitenmen(vBM)
sMC = Kaitenmen(vMC)
S = [sDA,sAB,sBM,sMC]

#板たち
ABCD = Poly3DCollection([[ra[0],rb[0],rc[0],rd[0]]],
                        fc=fc2+(a2,),clip_on=False)
R345 = Poly3DCollection([[[-10,0,-4],[-5.2,0,-4],[-5.2,0,2.4]],
                         [[0.4,0,-4],[ 5.2,0,-4],[ 5.2,0,-.4]]],
                        fc="k",alpha=0,clip_on=False)
for obj in [ABCD,R345]:
    ax1.add_collection(obj)

#数式原稿
lines1=[
r"=\dfrac{24}{5}\times\dfrac{24}{5}\times\:\,\pi\:\,\times\dfrac{32}{5}\times\dfrac{1}{3}\times\dfrac{7}{8}",
r"\,\fallingdotseq\,\dfrac{24}{5}\times\dfrac{24}{5}\times\dfrac{22}{7}\times\dfrac{32}{5}\times\dfrac{1}{3}\times\dfrac{7}{8}",
r"=\dfrac{\,24\times22\times32\,}{5\times5\times5}",
]
lines2=[
r"=\dfrac{24}{5}\times\dfrac{24}{5}\times\:\,\pi\:\,\times\dfrac{18}{5}\times\dfrac{1}{3}\times\dfrac{7}{8}",
r"\,\fallingdotseq\,\dfrac{24}{5}\times\dfrac{24}{5}\times\dfrac{22}{7}\times\dfrac{18}{5}\times\dfrac{1}{3}\times\dfrac{7}{8}",
r"=\dfrac{\,24\times22\times18\,}{5\times5\times5}",
]

#イージング関数
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

#数式クラス
colors = [(1,1,1,alpha) for alpha in np.arange(0,11)/10]
cmap = ListedColormap(colors)
linedic = dict(va="baseline",clip_on=True,fontsize=11)
shasendic = dict(color="k",lw=1)
class Lines:
    
    def __init__(self,ax,lines):
        self.ax = ax
        self.TXT = []
        self.PCM = []
        for i,line in enumerate(lines):
            txt = ax.text(0,-(i+1)*1.5+0.6,"$"+line+"$",**linedic)
            X,Y = np.mgrid[-2:10:0.1,-(i+1)*1.5:-i*1.5+0.1:1.5]
            C = np.ones((len(X)-1,1))
            C[0] = 0
            pcm = ax.pcolormesh(X,Y,C,cmap=cmap,zorder=10)
            self.TXT.append(txt)
            self.PCM.append(pcm)
        self.SHASEN = ax.plot(3.40,-2.35,7.05,-1.80,
                              0.65,-1.80,6.00,-2.35,7.05,-2.35,
                              **shasendic)
    
    def appearing(self,i,n):
        C = np.ones(len(self.PCM[n-1].get_array()))
        C[:2*i+1] = 0
        C[2*i+1:2*i+20] = np.arange(0.05,1.00,0.05)
        self.PCM[n-1].set_array(C)
    
    def go_next_line(self,i):
        a = easing(i/19)
        self.TXT[0].set_alpha(1-a)
        self.ax.set_ylim(-1.5*a-3,-1.5*a)
    
    def draw_shasen(self,i,sl):
        a = easing(i/19)
        for shasen in self.SHASEN[sl]:
            x0 = shasen.get_xdata()[0]
            y0 = shasen.get_ydata()[0]
            shasen.set_data([[x0,x0+0.6*a],[y0,y0-0.3*a]])
    
    def show_lines(self,i):
        if i<50:
            k = i
            self.appearing(k,1)
        elif i<100:
            k = i-50
            self.appearing(k,2)
        elif i<120:
            k = i-100
            self.go_next_line(k)
        elif i<140:
            k = i-120
            self.draw_shasen(k,slice(0,2))
        elif i<160:
            k = i-140
            self.draw_shasen(k,slice(2,5))
        elif i<200:
            k = i-160
            self.appearing(k,3)
        
        if i==199:
            for obj in self.TXT+self.SHASEN:
                obj.set_color("0.7")
        
#数式インスタンス     
lL = Lines(ax2,lines1)
lR = Lines(ax3,lines2)

#結論パート
txt3  = fig.text(0.50,0.15,"プリン２個で",ha="center",**linedic)
txt4a = fig.text(0.36,0.08,r"$\dfrac{\,24\times22\times50\,}{5\times5\times5}$",**linedic)
txt4b = fig.text(0.52,0.08,r"$=\mathbf{\dfrac{\,1056\,}{5}}$",**linedic)
txt4s = fig.text(0.515,0.11,r"$2$",**linedic)
for obj in [txt3,txt4a,txt4b,txt4s]:
    obj.set_alpha(0)
shasen1 = Line2D([0.425,0.425],[0.085,0.085],**shasendic)
shasen2 = Line2D([0.470,0.470],[0.085,0.085],**shasendic)
shasen3 = Line2D([0.485,0.485],[0.120,0.120],**shasendic)
for obj in [shasen1,shasen2,shasen3]:
    fig.add_artist(obj)

#長さ表記クラス
pdic = dict(color="0.6",linewidth=1,zorder=-10,clip_on=False)
tdic = dict(color="0.3",size=11,ha="center",va="center")
class Length:
    
    def __init__(self):
        self.objs = []
    
    def add(self,edge0,edge1,bump,latex):
        e0 = np.array([edge0]).T
        e1 = np.array([edge1]).T
        N = int(np.linalg.norm(e1-e0)*20)+1
        t = np.array([np.linspace(0,1,N)])
        p = e0 @ (1-t) + e1 @ t
        b = np.array([bump]).T @ (2*np.sqrt(t*(1-t)))
        p = p + b
        self.objs += ax1.plot(*p[:,:N//2-10],**pdic)
        self.objs += ax1.plot(*p[:,N//2+11:],**pdic)
        self.objs += [ax1.text(*p[:,N//2],latex,**tdic)]
    
    def set_alphas(self,a):
        for obj in self.objs:
            obj.set_alpha(a)

#長さ表記インスタンス
l1 = Length()
l1.add([0,0,0],[-3,0,0],[0,0,3/4],"$3$")
l1.add([-3,0,0],[-3,0,-4],[-3/4,0,0],"$4$")
l1.add([-3,0,-4],[3,0,-4],[0,0,-3/4],"$6$")
l1.add([3,0,-4],[0,0,0],[12/20,0,9/20],"$5$")
l2 = Length()
l2.add([-10,0,-4],[-5.2,0,-4],[0,0,-3/4],"$\\dfrac{24}{5}$")
l2.add([-5.2,0,-4],[-5.2,0,2.4],[3/4,0,0],"$\\dfrac{32}{5}$")
l2.add([-5.2,0,2.4],[-10,0,-4],[-12/20,0,9/20],"$8$")
l2.set_alphas(0)
l3 = Length()
l3.add([0.4,0,-4],[5.2,0,-4],[0,0,-3/4],"$\\dfrac{24}{5}$")
l3.add([5.2,0,-4],[5.2,0,-.4],[3/4,0,0],"$\\dfrac{18}{5}$")
l3.add([5.2,0,-.4],[0.4,0,-4],[-9/20,0,12/20],"$6$")
l3.set_alphas(0)

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

#台本
periods = {"P":10,"A0":30,"G1":n-1,
           "A1":30,"S1":30,"S2":60,"S3":30,
           "A2":30,"S4":30,"R1":30,
           "A3":30,"L1":200,"L2":200,
           "T1":30,"T2":30,"T3":20,"T4":30,
           "E":30}
endings = np.cumsum(list(periods.values()))
endings = np.append(0,endings)

#上演
def update(i):
    
    k,a,key = section(i)

    if "A0"==key:
        l1.set_alphas(1-a)
    elif "G1"==key:
        ABCD.set_verts([[ra[k+1],rb[k+1],rc[k+1],rd[k+1]]])
        for s in S:
            s.gururi(k+1)
    elif "A1"==key:
        ABCD.set_alpha(a2*(1-a))
    elif "S1"==key:
        sMC.slider(a,[1.5,0,-2])
    elif "S2"==key:
        sMC.mawaru(a)
    elif "S3"==key:
        sMC.slider(a,[1.5,0,-2])
    elif "A2"==key:
        if k==1:
            sAB.vs0 += [ra]
            sBM.vs0 += [rm]
            for s in [sAB,sBM]:
                s.set_verts(s.vs0)
        for s in [sDA,sMC]:
            s.set_facecolor(fc1+((1-a)*a1,))
        for s in [sAB,sBM]:
            s.set_facecolor([fc1+(a1,)]*(n-1)+[(0,0,0,a*a1)])
    elif "S4"==key:
        sAB.slider(a,[-7,0,0])
        sBM.slider(a,[3,0,0])
    elif "R1"==key:
        sAB.rotate(a,[0,-np.arcsin(3/5),0],[-10,0,-4])
        sBM.slirot(a,[10,0,0],[0,np.pi-np.arcsin(3/5),0],[0,0,-4])
    elif "A3"==key:
        R345.set_alpha(a*a1)
        l2.set_alphas(a)
        l3.set_alphas(a)
    elif "L1"==key:
        lL.show_lines(k)
    elif "L2"==key:
        lR.show_lines(k)
    elif "T1"==key:
        txt3.set_alpha(a)
    elif "T2"==key:
        txt4a.set_alpha(a)
    elif "T3"==key:
        txt4s.set_alpha(a)
        for shasen in [shasen1,shasen2,shasen3]:
            x0 = shasen.get_xdata()[0]
            y0 = shasen.get_ydata()[0]
            shasen.set_data([[x0,x0+0.03*a],[y0,y0-0.025*a]])
    elif "T4"==key:
        txt4b.set_alpha(a)

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