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

#紙の準備
fig = plt.figure()
fig.suptitle("2024東大数学　理系第５問",color="0.7",ha="right",x=0.96,y=0.96)
ax1 = fig.add_subplot(111,projection='3d',computed_zorder=False,zorder=10)
ax1.view_init(35,45)
ax1.set_xlim(0-.18,2-.18)
ax1.set_ylim(0,2)
ax1.set_zlim(0,2)
ax1.set_box_aspect((1,1,1))
ax1.axis("off")

#イージング関数は君に決めた
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

#計算式オブジェクト
class Kami:
    
    mathdic = dict(size=13,usetex=True,clip_on=True)
    colors = [(1,1,1,alpha) for alpha in np.arange(0,11)/10]
    cmap = ListedColormap(colors)
    cx = np.arange(-2,10,0.1)
    cn = len(cx)-1
    
    def __init__(self,left,lines):
        self.ax = fig.add_axes((left,0.07,0.3,0.25))
        self.ax.set_xlim(0,8)
        self.ax.set_ylim(-5,0)
        self.ax.set_aspect("equal")
        self.ax.axis("off")
        self.TEXT = []
        self.BELT = []
        for i,line in enumerate(lines):
            latex = "$\\displaystyle"+line+"$"
            text = self.ax.text(0,-2.5*(i+1)+1.0,latex,fontdict=self.mathdic)
            self.TEXT.append(text)
            cy = np.arange(-2.5*(i+1),-2.5*i+0.1,2.5)
            X,Y = np.meshgrid(self.cx,cy)
            C = np.append(0,np.ones(self.cn-1)).reshape((1,self.cn))
            belt = self.ax.pcolormesh(X,Y,C,cmap=self.cmap,zorder=10)
            self.BELT.append(belt)
    
    def kakikaki(self,k):
        n = (k-50)//70
        t = (k-50)%70
        if k==0:
            self.ax.set_zorder(20)
        if n<0:
            t -= 20
        if t<50:
            C = np.ones((1,self.cn))
            C[0,:2*t+1] = 0
            C[0,2*t+1:2*t+20] = np.arange(0.05,1.00,0.05)
            self.BELT[n+1].set_array(C)
        else:
            a = easing((t-50)/19)
            top = -(n+a)*2.5
            self.ax.set_ylim(top-5,top)
            self.TEXT[n].set_alpha(1-a)

#計算式オブジェクトの準備
lines1 = [
"=\\pi\\int_0^1\\big|(1-t\\,,0\\,)\\big|^2 dt",
"=\\pi\\int_0^1 (1-t)^2\\,dt",
"=\\pi\\Big[-\\!\\frac{1}{3}\\,(1-t)^3\\,\\Big]_0^1",
"=\\,\\frac{1}{3}\\,\\pi"
]
lines2 = [
"=\\pi\\int_{\\frac{1}{3}}^1\\big|(\\textstyle\\,\\frac{1-t}{2}\\,,\\frac{1-t}{2}\\,\\displaystyle)\\big|^2 dt",
"=\\pi\\int_{\\frac{1}{3}}^1 \\frac{1}{2}\\,(1-t)^2\\,dt",
"=\\pi\\Big[-\\!\\frac{1}{6}\\,(1-t)^3\\,\\Big]_{\\frac{1}{3}}^1",
"=\\,\\frac{4}{81}\\,\\pi"
]
lines3 = [
"=\\pi\\int_0^{\\frac{1}{3}}\\big|(1-2t\\,,\\,t\\,)\\big|^2 dt",
"=\\pi\\int_0^{\\frac{1}{3}} (1-4\\,t+5\\,t^2)\\,dt",
"=\\pi\\Big[\\;t-2\\,t^2+\\frac{5}{3}\\,t^3\\;\\Big]_0^{\\frac{1}{3}}",
"=\\,\\frac{14}{81}\\,\\pi"
]
kami1 = Kami(0.07,lines1)
kami2 = Kami(0.37,lines2)
kami3 = Kami(0.67,lines3)

#回転座標オブジェクト
class rot(np.ndarray):
    
    n = 61
    rotvec = np.array([0,0,2*np.pi])
    
    def __new__(cls,input_array):
        obj = np.asarray(input_array).view(cls)
        return obj
    
    def __init__(self,obj):
        self.Rvs = []
        for i in range(self.n):
            rvi = self.rotvec*i/(self.n-1)
            R = Rotation.from_rotvec(rvi).as_matrix()
            Rv = R @ obj
            self.Rvs.append(Rv)

#回転座標オブジェクトの準備
a = rot([0,0,1])
b = rot([1,0,0])
c = rot([0,1,0])
m = rot((a+c)/2)
g = rot((a+b+c)/3)

#回転面オブジェクト
class Kaitenmen:
    
    m = 30
    mendic = dict(alpha=0.2,clip_on=False)
    fc_cut = "darkcyan"
    
    def __init__(self,rot0,rot1,fc):
        self.n = len(rot0.Rvs)
        self.verts = []
        self.bands = []
        for i in range(1,self.n):
            vert = np.array([rot0.Rvs[i-1],rot1.Rvs[i-1],rot1.Rvs[i],rot0.Rvs[i]])
            self.verts.append(vert)
            band = Poly3DCollection([vert],fc=fc,**self.mendic)
            self.bands.append(band)
        self.vcuts = []
        for p in np.linspace(0,1,self.m):
            vcut = p*np.array(rot0.Rvs) + (1-p)*np.array(rot1.Rvs)
            self.vcuts.append(vcut)
        self.cut = Poly3DCollection([],fc=self.fc_cut,**self.mendic)
        ax1.add_collection3d(self.cut)
    
    def gururi(self,k):
        ax1.add_collection3d(self.bands[k])
        
    def cutter(self,k):
        self.cut.set_verts([self.vcuts[k]])
        
    def slider(self,a,slivec):
        slivec = np.array(slivec)
        for i,(vert,band) in enumerate(zip(self.verts,self.bands)):
            slider = np.tile(a*slivec,(4,1))
            new_vert = vert + slider
            band.set_verts([new_vert])
            if a==1:
                self.verts[i] = new_vert
        if a==1:
            for j,vcut in enumerate(self.vcuts):
                new_vcut = vcut + np.tile(slivec,(self.n,1)) 
                self.vcuts[j] = new_vcut 
                    
#回転面オブジェクトの準備             
kAB = Kaitenmen(a,b,"dodgerblue")
kAG = Kaitenmen(a,g,"mediumblue")
kGB = Kaitenmen(g,b,"royalblue")

#三角形オブジェクト
class Sankakukei(Poly3DCollection):
    
    mendic = dict(fc="orangered",clip_on=False)
    
    def __init__(self,rot0,rot1,rot2):
        super().__init__([[rot0,rot1,rot2]],**self.mendic)
        ax1.add_collection3d(self)
        n = len(rot0.Rvs)
        self.vs = []
        for i in range(1,n):
            v = [rot0.Rvs[i],rot1.Rvs[i],rot2.Rvs[i]]
            self.vs.append(v)
        
    def gururi(self,k):
        self.set_verts([self.vs[k]])

#三角形オブジェクトの準備    
AGB = Sankakukei(a,g,b)
AGM = Sankakukei(a,g,m)
BCM = Sankakukei(b,c,m)

#脇役テキストたちの準備
textdic = dict(color="0.2",ha="center",va="center",zorder=30,alpha=0)
t0 = fig.text(0.50,0.79,"$V=$",**textdic,fontsize=24)
t1 = fig.text(0.19,0.49,"$\\frac{1}{3}\\,\\pi$",**textdic,fontsize=18)
t2 = fig.text(0.39,0.49,"$-$",**textdic,fontsize=18)
t3 = fig.text(0.50,0.49,"$\\frac{4}{81}\\,\\pi$",**textdic,fontsize=18)
t4 = fig.text(0.61,0.49,"$-$",**textdic,fontsize=18)
t5 = fig.text(0.81,0.49,"$\\frac{14}{81}\\,\\pi$",**textdic,fontsize=18)
t6 = fig.text(0.50,0.22,"$=\\,\\frac{1}{9}\\,\\pi$",**textdic,fontsize=24)

#司会役
class Sikai:
    
    def __init__(self,clips):
        self.letters = []
        self.objects = []
        self.periods = []
        for clip in clips:
            self.letters.append(clip[0])
            self.objects.append(clip[1])
            self.periods.append(clip[2])
        endings = np.cumsum(self.periods)
        self.endings = np.append(0,endings)
        self.T = self.endings[-1]
    
    def section(self,i):
        p = np.argmin(~(i<self.endings[1:]))
        l = self.letters[p]
        x = self.objects[p]
        t = self.periods[p]
        k = i-self.endings[p]
        a = easing(k/(t-1))
        return l,x,k,a

#台本
clips = [
    ("P",[],10),
    ("A",[(AGM,1,0.5),(BCM,1,0)],30),
    ("G",[kAB,kAG,kGB,AGB,AGM],60),
    ("A",[(AGB,1,0),(AGM,0.5,0)],30),
    ("S",[(kAB,[0,0,4/3]),(kAG,[0,0,-1/3]),(kGB,[0,0,-4/3])],30),
    ("S",[(kAB,[1.4,-1.4,0]),(kGB,[-1.4,1.4,0])],30),
    ("S",[(kAB,[0,0,-4/3]),(kGB,[0,0,4/3])],30),
    ("C",kAB,30),
    ("K",kami1,220),
    ("A",[(kami1.TEXT[-1],1,0),(kami1.TEXT[-2],1,0),(t1,0,0.5)],30),
    ("C",kAG,30),
    ("K",kami2,220),
    ("A",[(kami2.TEXT[-1],1,0),(kami2.TEXT[-2],1,0),(t3,0,0.5)],30),
    ("C",kGB,30),
    ("K",kami3,220),
    ("A",[(kami3.TEXT[-1],1,0),(kami3.TEXT[-2],1,0),(t5,0,0.5)],30),
    ("A",[(t0,0,1)],20),
    ("A",[(t1,0.5,1),(t2,0,1),(t3,0.5,1),(t4,0,1),(t5,0.5,1)],20),
    ("A",[(t6,0,1)],20),
    ("P",[],20)
    ]

#司会役の手配
sk = Sikai(clips)

#上演
def update(i):
    
    l,x,k,a = sk.section(i)
    
    if l == "C":
        x.cutter(k)
    elif l == "K":
        x.kakikaki(k)
    elif l == "G":
        for obj in x:
            obj.gururi(k)
    elif l == "S":
        for obj,slivec in x:
            obj.slider(a,slivec)
    elif l == "A":
        for obj,a0,a1 in x:
            obj.set_alpha((1-a)*a0+a*a1)
    
mov = ani.FuncAnimation(fig,update,sk.T,interval=100)
plt.show()