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 Line3D,Poly3DCollection

#紙の準備
fig = plt.figure()
fig.suptitle("2023京大数学　理系第５問",color="0.7",ha="right",x=0.96,y=0.96)
ax1 = fig.add_subplot(111,projection='3d',computed_zorder=False)
ax1.set_xlim(-0.65,0.85)
ax1.set_ylim(-0.75,0.75)
ax1.set_zlim(-1.10,0.40)
ax1.set_box_aspect((1,1,1))
ax1.axis("off")
ax2 = fig.add_axes((0.35,0,0.6,0.3))
ax2.set_xlim(0,8)
ax2.set_ylim(-4,0)
ax2.axis("off")

#眺める角度の調整
view = [0,-90]
ax1.view_init(view[0],view[1])

#字幕スペースの準備
jimaku = plt.Rectangle((0,0),1,0.13,fc="darkslategrey",
                       zorder=20,transform=fig.transFigure)
fig.patches.append(jimaku)
kotobadic = dict(color="w",size=16,ha="center",va="center",
                 alpha=0,zorder=30,fontfamily="HGSoeiKakupoptai")
kotoba = fig.text(0.5,0.065,"",**kotobadic)

#らてふ
lines = [
"=\\pi\\int_0^1 y^2\\;dx",
"=\\pi\\int_0^1 \\left(\\,1-\\sqrt{x}\\;\\right)^4\\,dx",
"=\\pi\\int_0^1 \\left(\\,1-4\\,x^{\\frac{1}{2}}+6\\,x-4\\,x^{\\frac{3}{2}}+x^2\\;\\right)\\,dx",
"=\\pi\\,\\Bigg[\\,x-\\frac{8}{3}\\,x^{\\frac{3}{2}}+3\\,x^2-\\frac{8}{5}\\,x^{\\frac{5}{2}}+\\frac{1}{3}\\,x^3\\;\\Bigg]_0^1",
"=\\pi\\,\\Bigg(\\,1-\\frac{8}{3}+3-\\frac{8}{5}+\\frac{1}{3}\\;\\Bigg)",
"=\\,\\frac{1}{15}\\,\\pi"
]

#数式スペースの準備
mathdic = dict(size=16,usetex=True,clip_on=True)
colors = [(1,1,1,alpha) for alpha in np.arange(0,11)/10]
cmap = ListedColormap(colors)
txts,gyos,secs = [],[],[0]
for i,line in enumerate(lines):
    latex = "$\\displaystyle"+line+"$"
    text = ax2.text(0,-2*(i+1)+0.8,latex,fontdict=mathdic)
    txts.append(text)
    bbox = text.get_window_extent().transformed(ax2.transData.inverted())
    cx = np.arange(-2,bbox.x1+2,0.1)
    cy = np.arange(-2*(i+1),-2*i+0.1,2)
    X,Y = np.meshgrid(cx,cy)
    C = np.append(0,np.ones(len(cx)-2)).reshape((1,len(cx)-1))
    gyos.append(ax2.pcolormesh(X,Y,C,cmap=cmap,zorder=10))
    secs.append(secs[-1]+(len(cx)-20)//2)
    if i==0:
        secs.append(secs[-1])
    else:
        secs.append(secs[-1]+20)
sahen = fig.text(0.315,0.205,"$V$",**mathdic,zorder=30,alpha=0)
txts += [ax2.text(1.6,-11.2,"$\\displaystyle \\times\\;2$",**mathdic,alpha=0,zorder=20)]
txts += [ax2.text(2.5,-11.2,"$\\displaystyle =\\,\\frac{2}{15}\\,\\pi$",**mathdic,alpha=0,zorder=20)]

#平面図形の準備
nuri = Line3D([],[],[],c="darkkhaki")
sankaku = Poly3DCollection([],fc="khaki")
shahen = Line3D([],[],[],c="orange",lw=3,zorder=5)
himo = Line3D([0,0],[0,0],[1,0],c="darkorange",lw=5,solid_capstyle="round",clip_on=False,zorder=10)
heimen = [nuri,shahen,sankaku,himo]

#回すものの準備
m = 50
x = np.linspace(0,1,m+1)
y = np.zeros_like(x)
z = (1-np.sqrt(x))**2
sokumen = [np.stack([x,y,z]).T]
plate = Poly3DCollection([np.vstack([sokumen[0],[0,0,0]])],fc="darkkhaki",alpha=0)

#回転座標の準備
n = 60
for i in range(1,n+1):
    rotvec = i*np.array([0,0,-2*np.pi])/n
    R = Rotation.from_rotvec(rotvec)
    s = R.apply(sokumen[0])
    sokumen.append(s)

#回転面を作成
bands,bands2 = [],[]
for i in range(n):
    sikakus,sikakus2 = [],[]
    for j in range(m):
        sikaku = np.array([sokumen[i][j],sokumen[i][j+1],
                           sokumen[i+1][j+1],sokumen[i+1][j]])
        sikakus.append(sikaku)
        sikaku2 = np.array([1,1,-1])*sikaku
        sikakus2.append(sikaku2)
    band = Poly3DCollection(sikakus,fc="g",alpha=0.2,clip_on=False)
    bands.append(band)
    band2 = Poly3DCollection(sikakus2,fc="g",alpha=0,clip_on=False)
    bands2.append(band2)

#切断面の準備
cuts = []
for j in range(m+1):
    en = np.array([[],[],[]])
    for i in range(n+1):
        en = np.append(en,sokumen[i][m-j].reshape(3,1),axis=1)
    cuts.append(en)
danmen = Line3D([],[],[],c="darkorange",lw=3)

#お気に入りのイージング関数
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 Sikai:
    
    def __init__(self,clips):
        self.text = text
        self.texts = []
        self.verbs = []
        self.objects = []
        self.periods = []
        for clip in clips:
            self.texts.append(clip[0])
            self.verbs.append(clip[1:-1:2])
            self.objects.append(clip[2:-1:2])
            self.periods.append(clip[-1])
        endings = np.cumsum(self.periods)
        self.endings = np.append(0,endings)
        self.T = self.endings[-1]
    
    def work(self,i):
        p = np.argmin(~(i<self.endings[1:]))
        k = i-self.endings[p]
        l = self.periods[p]-1-k
        a = easing(k/(self.periods[p]-1))
        if k==0:
            kotoba.set_text(self.texts[p])
        elif k<6:
            kotoba.set_alpha(easing(k/5))
        elif l<6:
            kotoba.set_alpha(easing(l/5))
        waza = Waza(k,l,a)
        V = self.verbs[p]
        X = self.objects[p]
        for v,x in zip(V,X):
            getattr(waza,v)(x)

#使える技の数々
class Waza:
    
    def __init__(self,k,l,a):
        self.k = k
        self.l = l
        self.a = a
        
    #眺める角度を変える
    def shiten(self,x):
        ax1.view_init((1-self.a)*view[0]+self.a*x[0],
                      (1-self.a)*view[1]+self.a*x[1])
        if self.l==0:
            view[0]=x[0]
            view[1]=x[1]
        
    #通過領域を塗りつぶす
    def nuru(self,x):
        t = self.k/(self.k+self.l)
        oldnuri = nuri.get_data_3d()
        newnuri = np.hstack([oldnuri,[[t**2,0],[0,0],[(1-t)**2,1-t]]])
        nuri.set_data_3d(newnuri)
        shahen.set_data_3d([[0,t],[0,0],[1-t,0]])
        sankaku.set_verts([[(0,0,1-t),(0,0,0),(t,0,0)]])
        himo.set_data_3d([[0,0,t],[0,0,0],[1-t,0,0]])
    
    #回転させる
    def gururi(self,x):
        ax1.add_collection3d(bands[self.k])
        plate.set_verts([np.vstack([sokumen[self.k+1],[0,0,0]])])
        
    #断面で切る
    def wagiri(self,x):
        if self.k==0:
            ax1.add_line(danmen)
        elif self.l==0:
            danmen.set_visible(False)
        danmen.set_data_3d(cuts[self.k])
        for band in bands:
            band.set_facecolors(["g"]*self.l+["darkorange"]*self.k)
    
    #数式を書いていく
    def kakikaki(self,x):
        p = np.argmin(~(self.k<secs[1:]))
        q = self.k-secs[p]
        if p%2==0:
            C = gyos[p//2].get_array()
            C[2*q+1] = 0
            C[2*q+1:2*q+20] = np.arange(0.05,1.00,0.05)
            gyos[p//2].set_array(C)
        else:
            a = easing(q/19)
            top = -2*(p//2-1+a)
            ax2.set_ylim(top-4,top)
            txts[p//2-1].set_alpha(1-a)
    
    #何かを見せる
    def arawaru(self,x):
        if self.k==0:
            for obj in x:
                if type(obj)==Line3D:
                    ax1.add_line(obj)
                elif type(obj)==Poly3DCollection:
                    ax1.add_collection3d(obj)
        for obj in x:
            obj.set_alpha(self.a)
            
    #下側の回転面を見せる
    def arawaru2(self,x):
        if self.k==0:
            for obj in x:
                ax1.add_collection3d(obj)
        for obj in x:
            obj.set_alpha(0.2*self.a)
    
    #何かを消す
    def kieru(self,x):
        for obj in x:
            obj.set_alpha(1-self.a)

#台本
clips = [
    ("長さ$\\;\\mathbf{1}\\;$のひもがあります","arawaru",heimen,30),
    ("スルっと引っこぬきます","nuru",[],70),
    ("斜辺が通過した部分は","kieru",heimen,"arawaru",[plate],30),
    ("$\\mathbf{y\\;\\leqq\\;\\left(\\,1-\\sqrt{x}\\;\\right)\\,^2}$",30),
    ("これを","shiten",(30,-60),30),
    ("回します","gururi",[],n),
    ("円錐モドキが出来ました","kieru",[plate],30),
    ("輪ゴムを通すとこんな感じ","wagiri",[],m+1),
    ("積分したくねえ","kieru",[jimaku],"arawaru",[sahen],30),
    (None,"kakikaki",[],secs[-1]),
    ("おや！？","shiten",(0,-60),"arawaru",[jimaku],30),
    ("回転体のようすが・・・","arawaru2",bands2,"arawaru",[txts[-2]],30),
    ("そっちもアリなんかい！",15),
    ("というわけで",15),
    ("答えは","arawaru",[txts[-1]],30),
    ("おしまい！",40)
    ]

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

#上演
def update(i):
    sk.work(i)

mov = ani.FuncAnimation(fig,update,sk.T-10,interval=100)
plt.show()