import numpy as np
import matplotlib.pyplot as plt
import matplotlib.animation as ani
from matplotlib.colors import ListedColormap
from mpl_toolkits.mplot3d.art3d import Line3D
from mpl_toolkits.mplot3d.art3d import Poly3DCollection

#紙の準備
fig = plt.figure()
fig.canvas.draw()
fig.suptitle("2023東大数学　文系第４問",
             color="0.5",ha="right",x=0.95,y=0.95)
fig.subplots_adjust(left=0,right=1,wspace=0)
ax1 = fig.add_subplot(121,projection="3d",computed_zorder=False)
ax1.set_box_aspect((1,1,1))
ax1.set_xlim(-0.8,0.8)
ax1.set_ylim(-0.8,0.8)
ax1.set_zlim(-0.8,0.8)
ax1.axis("off")
ax2 = fig.add_subplot(122)
ax2.set_xlim( 0,8)
ax2.set_ylim(-8,0)
ax2.set_aspect("equal")
ax2.axis("off")

#球面
u = np.linspace(0,2*np.pi,100)
v = np.linspace(0,  np.pi,100)
x = np.outer(np.cos(u),np.sin(v))
y = np.outer(np.sin(u),np.sin(v))
z = np.outer(np.ones(np.size(u)),np.cos(v))
colors = [(0.6,0.9-0.2*k,0.8+0.2*k) for k in np.arange(0,100)/99]
cmap = ListedColormap(colors)
sphere = ax1.plot_surface(x,y,z,cmap=cmap,clip_on=False)

#数式の準備(区切り行はリスト、空白行はNoneで記述)
linedic = dict(va="baseline",math_fontfamily="cm",clip_on=True)
lines = [
(1.0,"$\\mathrm{AC}=\\mathrm{BC}=x$","とおく"),
(0.5,"△$\\mathrm{ABC}$","に余弦定理を適用すると"),
(0.0,"$1^2=x^2+x^2-2\\cdot x\\cdot x\\cdot\\cos\\angle\\mathrm{ACB}$"),
(1.5,"$1=2\\,x^2-2\\,x^2\\cdot\\dfrac{4}{5}$"),
(2.6,"$\\dfrac{2}{5}\\,x^2=1$"),
[3.0,"$x^2=\\dfrac{5}{2}$"],
None,
(0.5,"△$\\mathrm{ACM}$","に三平方の定理を適用して"),
(2.0,"$x^2=\\left(\\,\\dfrac{1}{2}\\right)^2+\\mathrm{CM}^2$"),
(2.5,"$\\dfrac{5}{2}=\\dfrac{1}{4}+\\mathrm{CM}^2$"),
(2.9,"$\\mathrm{CM}^2=\\dfrac{9}{4}$"),
[1.0,"CM>0より","$\\mathrm{CM}=\\dfrac{3}{2}$"],
None,
(0.5,"△$\\mathrm{OAB}$","は１辺が１の正三角形なので"),
(1.0,"$\\mathrm{OM}=\\sin{60\\degree}=\\frac{\\sqrt{3}}{2}$"),
None,
[1.0,"$\\angle\\mathrm{CMO}=\\theta$","とおく"],
(0.5,"△$\\mathrm{CMO}$","に余弦定理を適用すると"),
(0.0,"$1^2=\\left(\\,\\frac{3}{2}\\right)^2\\!+\\!\\left(\\,\\frac{\\sqrt{3}}{2}\\right)^2\\!-\\!2\\cdot\\frac{3}{2}\\cdot\\frac{\\sqrt{3}}{2}\\cdot\\cos\\theta$"),
(1.3,"$1^2=\\dfrac{9}{4}+\\dfrac{3}{4}-\\frac{3\\sqrt{3}}{2}\\,\\cos\\theta$"),
(2.2,"$\\frac{3\\sqrt{3}}{2}\\,\\cos\\theta=2$"),
[2.5,"$\\cos\\theta=\\frac{4}{3\\sqrt{3}}$"],
(0.1,"sinθ>0より","$\\sin\\theta=\\sqrt{1-\\cos^2\\theta}$"),
(3.4,"$=\\sqrt{1-\\frac{16}{27}}=\\frac{\\sqrt{11}}{3\\sqrt{3}}$"),
(0.5,"$\\sin2\\theta=2\\,\\sin\\theta\\,\\cos\\theta$","なので"),
[0.1,"$\\sin\\angle\\mathrm{CMD}=2\\cdot\\frac{4}{3\\sqrt{3}}\\cdot\\frac{\\sqrt{11}}{3\\sqrt{3}}=\\frac{8\\sqrt{11}}{27}$"],
None,
[0.1,"△$\\mathrm{CMD}=\\dfrac{1}{2}\\cdot\\mathrm{CM}\\cdot\\mathrm{DM}\\cdot\\sin\\angle\\mathrm{CMD}$"],
(1.7,"$=\\frac{1}{2}\\cdot\\frac{3}{2}\\cdot\\frac{3}{2}\\cdot\\frac{8\\sqrt{11}}{27}$"),
(1.7,"$=\\frac{\\sqrt{11}}{3}$"),
None,
(0.5,"求める体積は"),
(1.0,"$V=\\dfrac{1}{3}\\cdot\\Delta\\mathrm{CMD}\\cdot\\mathrm{AB}$"),
(1.35,"$=\\frac{1}{3}\\cdot\\frac{\\sqrt{11}}{3}\\cdot1$"),
(1.35,"$=\\frac{\\sqrt{11}}{9}$"),
]

#数式クラス
class MathLine:
    
    def __init__(self,lines):
        colors = [(1,1,1,alpha) for alpha in np.arange(0,11)/10]
        cmap = ListedColormap(colors)
        cx = np.arange(-2,10,0.1)
        self.cn = len(cx)-1
        self.BELT = []
        self.LEFT = []
        self.TIME = []
        self.GOES = [(0,0)]
        for i,line in enumerate(lines):
            if line is not None:
                tx = line[0]
                L = int(tx//0.2)
                self.LEFT.append(L)
                mojis = line[1:]
                for moji in mojis :
                    if moji[-1]=="$":
                        fs = 14
                    else:
                        fs = 11
                    text = ax2.text(tx,-i-0.65,moji,fontsize=fs,**linedic)
                    bbox = text.get_window_extent().transformed(ax2.transData.inverted())
                    tx = bbox.x1 + 0.2
                cy = np.arange(-i-1,-i+0.1,1)
                X,Y = np.meshgrid(cx,cy)
                C = np.ones((1,self.cn))
                C[0,0] = 0
                belt = ax2.pcolormesh(X,Y,C,cmap=cmap,zorder=10)
                self.BELT.append(belt)
                R = int(tx//0.2)+10
                T = R-L
                self.TIME.append(T)
            if type(line)==list:
                n = i - sum(l==None for l in lines[:i])
                self.GOES.append((i,n))
            
    def appear_line(self,nth,k):
        C = np.ones((1,self.cn))
        left = 2*(self.LEFT[nth-1]+k)
        C[0,:left+1] = 0
        C[0,left+1:left+20] = np.arange(0.05,1.00,0.05)
        self.BELT[nth-1].set_array(C)
    
    def go_down(self,nth,a):
        i1 = self.GOES[nth-1][0]
        i2 = self.GOES[nth][0]
        n1 = self.GOES[nth-1][1]
        n2 = self.GOES[nth][1]
        down = i2-i1
        top = -i1-a*down
        ax2.set_ylim(top-8,top)
        for n in range(n1,n2):
            self.BELT[n].set_alpha(a)

#数式インスタンス
ml = MathLine(lines)

#点の座標
o = np.array([0,0,0])
a = np.array([np.sqrt(3)/2,-1/2,0])
b = np.array([np.sqrt(3)/2, 1/2,0])
c = np.array([-np.sqrt(3)/6,0, np.sqrt(33)/6])
d = np.array([-np.sqrt(3)/6,0,-np.sqrt(33)/6])
m = (a+b)/2

#三角形
mendic = dict(fc="peru")
ACM = Poly3DCollection([[a,c,m]],**mendic)
BCM = Poly3DCollection([[b,c,m]],**mendic)
OAB = Poly3DCollection([[o,a,b]],**mendic)
OCM = Poly3DCollection([[o,c,m]],**mendic)
CDM = Poly3DCollection([[c,d,m]],**mendic)

#線分
sendic = dict(c="lightslategrey",lw=1)
CA = Line3D(*zip(c,a),**sendic)
CB = Line3D(*zip(c,b),**sendic)
AB = Line3D(*zip(a,b),**sendic)
CD = Line3D(*zip(c,d),**sendic)
DA = Line3D(*zip(d,a),**sendic)
DB = Line3D(*zip(d,b),**sendic,linestyle="--")
CM = Line3D(*zip(c,m),**sendic)
DM = Line3D(*zip(d,m),**sendic)
OM = Line3D(*zip(o,m),**sendic,linestyle="--")
OC = Line3D(*zip(o,c),**sendic,linestyle="--")

#テキストたち(右)
textdic = dict(color="0.3",ha="center",va="center")
tA = ax1.text(*1.10*a,"A",**textdic)
tB = ax1.text(*1.10*b,"B",**textdic)
tC = ax1.text(*1.12*c,"C",**textdic)
tD = ax1.text(*1.12*d,"D",**textdic)
tM = ax1.text(*1.14*m,"M",**textdic)
tO = ax1.text(-0.12,0,0,"O",**textdic)
tmaru = [ax1.text(*(c+a)/2,"◯",**textdic)]
tmaru+= [ax1.text(*(c+b)/2,"◯",**textdic)]
tmaru+= [ax1.text(*(d+a)/2,"◯",**textdic)]
tmaru+= [ax1.text(*(d+b)/2,"◯",**textdic)]

#テキストたち(左)
textdic = dict(color="0.2",ha="center",fontsize=12)
tf1 = fig.text(0.26,0.88,"半径１の球面",**textdic)
tf2 = fig.text(0.26,0.80,"球面上の４点 A,B,C,D",**textdic)
tf3a= fig.text(0.14,0.20,"$\\mathrm{AB=1,}$",**textdic)
tf3b= fig.text(0.31,0.20,"$\\mathrm{AC=BC,\\,AD=BD}$",**textdic)
tf3c= fig.text(0.31,0.20,"$\\mathrm{CM=DM=\\dfrac{3}{2}}$",**textdic)
tf4a= fig.text(0.26,0.12,"$\\cos\\angle\\mathrm{ACB}=\\cos\\angle\\mathrm{ADB}=\\dfrac{4}{5}$",**textdic)
tf4b= fig.text(0.32,0.11,"$\\sin\\angle\\mathrm{CMD}=\\frac{8\\sqrt{11}}{27}$",**textdic)

#長さ表記を準備するための関数(立体版)
pdic = dict(color="0.6",linewidth=1,zorder=-10,clip_on=False)
tdic = dict(color="0.3",size=11,ha="center",va="center")
def Length(edge0,edge1,bump,latex):
    e0 = np.array([edge0]).T
    e1 = np.array([edge1]).T
    N = int(np.linalg.norm(e1-e0)*100)+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
    parts  = ax1.plot(*p[:,:N//2-10],**pdic)
    parts += ax1.plot(*p[:,N//2+11:],**pdic)
    parts +=[ax1.text(*p[:,N//2],latex,**tdic)]
    return parts

#長さ表記たち
lAB = Length(a,b,[0.2,0,0],"$1$")
lAM = Length(a,m,[0.2,0,0],"$\\dfrac{1}{2}$")
lAC = Length(a,c,[-.04,-.16,0],"$x$")
lBC = Length(b,c,[-.04,0.16,0],"$x$")
lCM = Length(c,m,[.14,0,0.14],"$\\dfrac{3}{2}$")
lDM = Length(d,m,[.14,0,-.14],"$\\dfrac{3}{2}$")
lOC = Length(o,c,[-.16,0,-.04],"$1$")
lOM = Length(o,m,[0,0,-0.2],"$\\frac{\\sqrt{3}}{2}$")
lAB2= Length(a,b,[.28,0,0],"$1$")

#角度表記を準備するための関数(立体版)
def Angle(p0,p1,p2):
    a = p0-p1
    b = p2-p1
    n = np.cross(a,b)
    n = n/np.linalg.norm(n)
    theta = np.arccos(np.dot(a,b)/np.linalg.norm(a)/np.linalg.norm(b))
    N = int(theta/0.1)
    verts = [p1]
    r = 0.3*a/np.linalg.norm(a)
    for t in np.linspace(0,theta,N):
        verts += [p1+np.cos(t)*r+np.sin(t)*np.cross(n,r)]
    return Poly3DCollection([verts],fc="orangered")

#角度表記たち    
aACB = Angle(a,c,b)
aADB = Angle(a,d,b)
aCMO = Angle(c,m,o)
aDMO = Angle(d,m,o)

#出現消滅の台本
appears = [
    [sphere,tf1],
    [CA,CB,AB,CD,DA,DB,tA,tB,tC,tD,tf2],
    [tf3a,tf3b]+tmaru+lAB,
    [tf4a,aACB,aADB],
    [ACM,BCM],
    lAC+lBC,
    [CM,tM]+lAM,
    [OAB,CD,DA,DB,OM,tO,tD,tf3c],
    [OCM,OC,CM]+lCM+lOC+lOM,
    [aCMO],
    [CD,DM,tD],
    [aDMO],
    [CDM,tf4b]+lCM+lDM,
    [CA,CB,AB,DA,DB,tA,tB]+lAB2,
    ]
disappears = [
    [],
    [],
    [],
    [],
    [CD,DA,DB,tD,aADB]+tmaru,
    [],
    [BCM,aACB]+lAB+lBC,
    [ACM,CM,tf3b,tf4a]+lAC+lAM,
    [OAB,CA,CB,AB,CD,DA,DB,tA,tB,tD],
    [],
    [OCM]+lCM+lOC+lOM,
    [],
    [OC,OM,tO],
    [aCMO,aDMO]+lCM+lDM,
    ]

#出現消滅クラス
class AppearDisappear:
    
    def __init__(self,Appears,Disappears):
        self.Appears = Appears
        self.Disappears = Disappears
        for objs in Appears:
            for obj in objs:
                obj.set_alpha(0)
                if type(obj) == Poly3DCollection:
                    ax1.add_collection3d(obj)
                elif type(obj) == Line3D:
                    ax1.add_line(obj)
    
    def do(self,nth,a):
        for obj in self.Appears[nth-1]:
            if type(obj) == Poly3DCollection:
                alpha = 0.2*a
            else:
                alpha = a
            obj.set_alpha(alpha)
        for obj in self.Disappears[nth-1]:
            if type(obj) == Poly3DCollection:
                alpha = 0.2*(1-a)
            else:
                alpha = 1-a
            obj.set_alpha(alpha)

#出現消滅インスタンス
ad = AppearDisappear(appears,disappears)

#視点変更クラス
elevs = [30,45,30,0,30]
azims = [-75,0,-75,-90,-75]
class ChangeView:
    
    def __init__(self,Elevs,Azims):
        self.Elevs = Elevs
        self.Azims = Azims
        ax1.view_init(self.Elevs[0],self.Azims[0])
        
    def do(self,nth,a):
        elev0 = self.Elevs[nth-1]
        elev1 = self.Elevs[nth]
        azim0 = self.Azims[nth-1]
        azim1 = self.Azims[nth]
        ax1.view_init(elev0+a*(elev1-elev0),azim0+a*(azim1-azim0))

#視点変更インスタンス
cv = ChangeView(elevs,azims)

#全体台本の準備
clips = ["A"]*4+["CA"]+["LA"]+["L"]*5+\
        ["GA"]+["L"]*5+\
        ["CGA"]+["L"]*2+\
        ["CA"]+["LA"]+["L"]*2+\
        ["G"]+["L"]*5+\
        ["GA"]+["LA"]+["L"]+\
        ["GA"]+["L"]*3+\
        ["CGA"]+["L"]*4+["E"]

#司会クラス
class TimeLine:
    
    def __init__(self,clips):
        self.clips = clips
        self.nths = []
        self.period = []
        for i,clip in enumerate(clips):
            ndic = {}
            for letter in clip:
                ndic[letter] = sum(letter in c for c in clips[:i])+1
            self.nths.append(ndic)
            if "L" in clip:
                t = ml.TIME[ndic["L"]-1]
            else:
                t = 30
            self.period.append(t)
        endings = np.cumsum(self.period)
        self.endings = np.append(0,endings)
        self.T = self.endings[-1]
        
    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 section(self,i):
        p = np.argmin(~(i<self.endings[1:]))
        clip = self.clips[p]
        ndic = self.nths[p]
        k = i-self.endings[p]
        t = self.period[p]
        a = self.easing(k/(t-1))
        return clip,ndic,k,a

#司会インスタンス
tl = TimeLine(clips)

#上演
def update(i):
    
    clip,ndic,k,a = tl.section(i)
    
    if "C" in clip:
        cv.do(ndic["C"],a)
    if "L" in clip:
        ml.appear_line(ndic["L"],k)
    if "G" in clip:
        ml.go_down(ndic["G"],a)
    if "A" in clip:
        ad.do(ndic["A"],a)

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