import numpy as np
import matplotlib.pyplot as plt
import matplotlib.animation as ani
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)
fig.subplots_adjust(left=0,right=0.985,bottom=0.18)
ax = fig.add_subplot(projection='3d',computed_zorder=False)
ax.view0 = (90,-90)
ax.view_init(*ax.view0)
ax.set_box_aspect((1,1,1))
ax.set_xlim(-1,1)
ax.set_ylim(-1,1)
ax.set_zlim(0,2)
ax.axis("off")

#折れる面クラス
class Men:
    
    rot1 = -np.arccos(-1/np.sqrt(3))*np.array([1/2,np.sqrt(3)/2,0])
    fix1 = np.array([-3/4,np.sqrt(3)/4,0])
    
    def __init__(self,verts,fc=(0.7,0.3,0,0.2),ec=(0.5,0.5,0.5)):
        self.verts = []
        for vert in verts:
            self.verts.append(vert+(0,))
        self.poly  = Poly3DCollection([self.verts],clip_on=False,fc=fc,ec=ec)
        ax.add_collection3d(self.poly)
        self.R1verts = self.fold(self.verts,self.rot1,self.fix1)
        
    def fold(self,verts,rotvec,fixpoint):
        R = Rotation.from_rotvec(rotvec)
        Rverts = []
        for i,vert in enumerate(verts):
            Rvert = R.apply(np.array(vert)-fixpoint)+fixpoint
            Rverts.append(Rvert)
        return Rverts
    
    def fold1(self,a):
        Rverts = self.fold(self.verts,a*self.rot1,self.fix1)
        self.poly.set_verts([Rverts])
        
    def fold2(self,a,p2):
        f1 = self.R1verts[0]
        f2 = self.R1verts[len(self.verts)-4]
        fix = (f1+f2)/2
        p1 = self.R1verts[-2]-fix
        p2 = p2-fix
        theta = np.arccos(np.dot(p1,p2)/np.linalg.norm(p1)/np.linalg.norm(p2))
        rot = theta*(f1-f2)/np.linalg.norm(f1-f2)
        Rverts = self.fold(self.R1verts,a*rot,fix)
        self.poly.set_verts([Rverts])

#面たち、その１
hex0 = Men([(
    np.cos(i*np.pi/3),
    np.sin(i*np.pi/3))
    for i in range(6)],fc=(0,0.7,0.3,0.2),ec=None)
tri1 = Men([(
    np.cos(i*np.pi/2)/2,
    np.sin(i*np.pi/2)/2+np.sqrt(3)/2)
    for i in range(3)])
tri2 = Men([(
    np.cos((i/2-2/3)*np.pi)/2+3/4,
    np.sin((i/2-2/3)*np.pi)/2-np.sqrt(3)/4)
    for i in range(3)])
tri3 = Men([(
    np.cos((i/2+2/3)*np.pi)/2-3/4,
    np.sin((i/2+2/3)*np.pi)/2-np.sqrt(3)/4)
    for i in range(3)])
pnt1 = Men([(
    np.cos(i*np.pi/3),
    np.sin(i*np.pi/3))
    for i in range(4,6)]+[(
    np.cos((-i/2)*np.pi),
    np.sin((-i/2)*np.pi)-(np.sqrt(3)+1)/2)
    for i in range(3)])
pnt2 = Men([(
    np.cos(i*np.pi/3),
    np.sin(i*np.pi/3))
    for i in range(2,4)]+[(
    np.cos((-i/2+4/3)*np.pi)-(3+np.sqrt(3))/4,
    np.sin((-i/2+4/3)*np.pi)+(np.sqrt(3)+1)/4)
    for i in range(3)])
pnt3 = Men([(
    np.cos(i*np.pi/3),
    np.sin(i*np.pi/3))
    for i in range(0,2)]+[(
    np.cos((-i/2+2/3)*np.pi)+(3+np.sqrt(3))/4,
    np.sin((-i/2+2/3)*np.pi)+(np.sqrt(3)+1)/4)
    for i in range(3)])

#面たち、その２
P = np.array([(
    np.cos((i/2-1/6)*np.pi)-(3+np.sqrt(3))/4,
    np.sin((i/2-1/6)*np.pi)+(np.sqrt(3)+1)/4,
    np.sqrt(2)*(1-i//4)) for i in range(5)])
pnt456 = Poly3DCollection([
    [hex0.R1verts[5],hex0.R1verts[0],P[1],P[0],P[3]],
    [hex0.R1verts[1],hex0.R1verts[2],P[4],P[0],P[1]],
    [hex0.R1verts[3],hex0.R1verts[4],P[3],P[0],P[4]]
    ],fc=(1,1,1,0),clip_on=False)
ax.add_collection3d(pnt456)

#文字たち
tdic = dict(ha="center",va="center",color="0.3",
            size=16,fontfamily="HGSoeiKakupoptai")
txt1 = fig.text(-1,0.28,"立方体の",**tdic)
txt2 = fig.text(-1,0.20,"ジャスト",**tdic)
txt3 = fig.text(-1,0.12,"半分や！",**tdic)

#司会役クラス
class Sikai:
    
    def easing(self,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 __init__(self,clips):
        self.verbs = []
        self.objects = []
        self.periods = []
        for clip in clips:
            self.verbs.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:]))
        v = self.verbs[p]
        x = self.objects[p]
        a = self.easing((i-self.endings[p])/(self.periods[p]-1))
        return v,x,a

#台本
clips = [
    ("P",[],20),
    ("V",(35,-90),30),
    ("F1",[],40),
    ("V",(35,-105),15),
    ("F2",(tri3,pnt2.verts[2]),15),
    ("V",(35,-195),30),
    ("F2",(pnt1,pnt2.verts[3]),30),
    ("V",(35,-210),15),
    ("F2",(tri2,P[2]),15),
    ("V",(35,-225),15),
    ("F2",(pnt3,pnt2.verts[3]),30),
    ("V",(35,-315),30),
    ("F2",(tri1,pnt2.verts[4]),15),
    ("V",(35,-405),30),
    ("A1",[],40),
    ("A2",[],10),
    ("A3",(txt1,-0.55),10),
    ("A3",(txt2,-0.50),10),
    ("A3",(txt3,-0.45),10),
    ("P",[],20)]

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

#上演
def update(i):
    
    v,x,a = sk.section(i)
    
    if v == "V":
        elev = a*x[0] + (1-a)*ax.view0[0]
        azim = a*x[1] + (1-a)*ax.view0[1]
        ax.view_init(elev,azim)
        if a==1:
            ax.view0 = x
            
    elif v == "F1":
        for obj in [hex0,tri1,tri2,tri3,pnt1,pnt3]:
            obj.fold1(a)
        new_x = -a*(3+np.sqrt(3))/4
        new_y =  a*(np.sqrt(3)+1)/4
        ax.set_xlim(new_x-1,new_x+1)
        ax.set_ylim(new_y-1,new_y+1)
        
    elif v == "F2":
        x[0].fold2(a,x[1])
        
    elif v == "A1":
        pnt456.set_facecolor((0.3,0.7,0,0.2*a))
        pnt456.set_edgecolor((0.5,0.5,0.5,0.5*a))
        
    elif v == "A2":
        ax.set_zlim(-0.8*a,2-0.8*a)
        
    elif v == "A3":
        x[0].set_x(x[1]+a)

#上演
mov = ani.FuncAnimation(fig,update,sk.T,interval=100)
plt.show()