import numpy as np
import matplotlib.pyplot as plt
import matplotlib.animation as ani
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.canvas.draw()
fig.suptitle("2012九大数学　理系第１問",
             color="0.7",ha="right",x=0.96,y=0.96)
gs = GridSpec(2,4,height_ratios=[3,2])
ax1 = fig.add_subplot(gs[0,:],projection='3d')
ax1.view_init(90,-90)
ax1.set_xlim(-1.2,1.8)
ax1.set_ylim(-1.5,1.5)
ax1.set_zlim(-1.5,1.5)
ax1.set_box_aspect((1,1,1))
ax1.axis("off")
ax2 = fig.add_subplot(gs[1,1:])
ax2.set_xlim(0,8)
ax2.set_ylim(-3,0)
ax2.set_aspect("equal")
ax2.axis("off")

#円を準備する関数
def draw_circle(theta):
    x = 2*np.cos(theta)
    y = 2*np.sin(theta)+1
    z = np.zeros_like(x)
    xyz = np.stack([x,y,z])
    vert = list(zip(x,y,z))
    c = Poly3DCollection([vert],fc='darkorange')
    ax1.add_collection3d(c)
    return c,xyz

#円を準備
n = 31
theta1 = np.linspace(-np.pi/2,3*np.pi/2,n+15)
theta2 = np.linspace(-np.pi/6,7*np.pi/6,n)
c1,_   = draw_circle(theta1)
c2,xyz = draw_circle(theta2)

#回転体を準備
m = 61
phi = np.linspace(0,2*np.pi,m)
rotxyzs = []
for i in range(m):
    rotvec = np.array([phi[i],0,0])
    rot = Rotation.from_rotvec(rotvec)
    rotxyz = rot.as_matrix() @ xyz
    rotxyzs.append(rotxyz)
bands = []
for i in range(1,m):
    verts = []
    for j in range(1,n):
        vert = [rotxyzs[i-1][:,j-1],rotxyzs[i-1][:,j],rotxyzs[i][:,j],rotxyzs[i][:,j-1]]
        verts.append(vert)
    band = Poly3DCollection(verts,ec=(1,0,0,0.1),fc=(1,0,0,0.1))
    bands.append(band)

#輪切りを準備
h = 80
x = np.linspace(2,-2,h)
y1 =  np.sqrt(4-x**2)+1
y2 = -np.sqrt(4-x**2)+1
cut_verts = []
for i in range(h):
    cx = x[i]*np.ones_like(phi)
    if y2[i]>0:
        cy1 = y1[i]*np.cos(phi)
        cz1 = y1[i]*np.sin(phi)
        cy2 = y2[i]*np.cos(-phi)
        cz2 = y2[i]*np.sin(-phi)
        vert = list(zip(cx,cy1,cz1))+list(zip(cx,cy2,cz2))
    else:  
        cy = y1[i]*np.cos(phi)
        cz = y1[i]*np.sin(phi)
        vert = list(zip(cx,cy,cz))
    cut_verts.append(vert)
c3 = Poly3DCollection([(0,0,0)],fc='orangered',alpha=0.8)
ax1.add_collection3d(c3)

#フォント辞書 
mathdic={
"size" : 16,
"usetex" : True,
"clip_on" : True,
}

#原稿を準備
lines = [
"=\\pi\\int_{y\\,\\geq\\,0}\\,y^2\\,dx",
"=\\pi\\int_{\\frac{7}{6}\\pi}^{-\\frac{\\pi}{6}}\\!y^2\\,\\frac{dx}{d\\theta}\\,d\\theta",
"=\\pi\\int_{\\frac{7}{6}\\pi}^{-\\frac{\\pi}{6}}(2\\sin\\theta+1)^2(2\\cos\\theta)^{\\prime}\\,d\\theta",
"=\\pi\\int_{\\frac{7}{6}\\pi}^{-\\frac{\\pi}{6}}(2\\sin\\theta+1)^2(-2\\sin\\theta)\\,d\\theta",
"=2\\pi\\int_{-\\frac{\\pi}{6}}^{\\frac{7}{6}\\pi}(2\\sin\\theta+1)^2\\sin\\theta\\,d\\theta",
"=2\\pi\\int_{-\\frac{\\pi}{6}}^{\\frac{7}{6}\\pi}(4\\sin^3\\theta+4\\sin^2\\theta+\\sin\\theta)\\,d\\theta",
"=2\\pi\\int_{-\\frac{\\pi}{6}}^{\\frac{7}{6}\\pi}(\\scriptstyle3\\sin\\theta-\\sin3\\theta+2-2\\cos2\\theta+\\sin\\theta\\displaystyle)\\,d\\theta",
"=2\\pi\\int_{-\\frac{\\pi}{6}}^{\\frac{7}{6}\\pi}(\\scriptstyle4\\sin\\theta-\\sin3\\theta-2\\cos2\\theta+2\\displaystyle)\\,d\\theta",
"=2\\pi\\Biggl[-4\\cos\\theta\\!+\\!\\frac{\\cos3\\theta}{3}\\!-\\!\\sin2\\theta\\!+\\!2\\theta\\Biggr]_{-\\frac{\\pi}{6}}^{\\frac{7}{6}\\pi}",
"=2\\pi\\Biggl\\{\\scriptstyle-4(-\\frac{\\sqrt{3}}{2}-\\frac{\\sqrt{3}}{2})+0-(\\frac{\\sqrt{3}}{2}+\\frac{\\sqrt{3}}{2})+2(\\frac{7}{6}\\pi+\\frac{\\pi}{6})\\Biggr\\}",
"=2\\pi\\left(3\\sqrt{3}+\\frac{8}{3}\\,\\pi\\right)",
"=6\\sqrt{3}\\,\\pi+\\frac{16}{3}\\,\\pi^2",
]

#数式を書き込んでいきます
colors = [(1,1,1,alpha) for alpha in np.arange(0,11)/10]
cmap   = ListedColormap(colors)
Texts,Times1,Pcms,Times2 = [],[],[],[]
for i,line in enumerate(lines):
    latex = "$\\displaystyle"+line+"$"
    text = ax2.text(0,-1.5*(i+1)+0.63,latex,fontdict=mathdic)
    Texts.append(text)
    bbox = text.get_window_extent().transformed(ax2.transData.inverted())
    time1 = int(bbox.x1*5+30)
    Times1.append(time1)
    X,Y = np.mgrid[-2:14:0.1,-1.5*(i+1):-1.5*i+0.1:1.5]
    C = np.ones((len(X)-1,1))
    C[0] = 0
    pcm = ax2.pcolormesh(X,Y,C,cmap=cmap,zorder=10)
    Pcms.append(pcm)
Times2 = np.insert(Times1,range(2,len(Times1)),20)
Times2 = np.cumsum(Times2)

#その他もろもろを準備
jiku,= ax1.plot([-20,20],[0,0],[0,0],color="lightsteelblue")
waku = fig.text(0.5,0.2,"円 $\\mathbf{x^2+(y-1)^2=4}$ を\n$\\mathbf{x}\\,$軸まわりに回転させます！",
                size=12,color="steelblue",ha="center",linespacing=2,fontfamily="HGSoeiKakupoptai",
                bbox=dict(boxstyle="round,pad=1.5",lw=1.5,fc="aliceblue",ec="steelblue"))
L    = fig.text(.295,.305,"$V$",fontdict=mathdic,alpha=0.0)
CW   = ax2.text(0.68,-0.9,"↻",size=24,rotation=180,fontfamily='Meiryo',clip_on=True)
last = fig.text(0.68,0.16,"おしまい",size=16,color="0.4",fontfamily="HGSoeiKakupoptai",alpha=0.0)

#立体ははみ出して
for obj in bands+[jiku,c2,c3]:
    obj.set_clip_on(False)

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

#単一行を徐々に見せていく関数
def appear_line(i,pcm):
    C = np.ones(len(pcm.get_array()))
    C[:2*i+1] = 0
    C[2*i+1:2*i+20] = np.arange(0.05,1.00,0.05)
    pcm.set_array(C)

#複数行を徐々に見せていく関数
def show_lines(i):
    n = np.argmin(~(i<Times2))
    Times3 = np.append(0,Times2)
    k = i-Times3[n]
    if n==0:
        appear_line(k,Pcms[0])
    elif n%2 == 1:
        appear_line(k,Pcms[n//2+1])
    else:
        alpha = (k+1)/20
        Texts[n//2-1].set_alpha(easing(1-1.5*alpha))
        if n==2:
            CW.set_alpha(easing(1-1.5*alpha))
        top = -1.5*((n//2-1)+easing(alpha))
        ax2.set_ylim(top-3,top)

#セクションを判定する関数
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 = {"P1":20,"A1":30,"A2":30,
           "Rt":m-1,"A3":30,"Sl":h,
           "A4":30,"L1":Times2[-1],
           "E":20,"P2":10}
endings = np.cumsum(list(periods.values()))
endings = np.append(0,endings)

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

    if "A1"==key:
        c1.set_alpha(1-a)
    elif "A2"==key:
        ax1.view_init(90-60*a,-90+30*a)
    elif "Rt" == key:
        ax1.add_collection3d(bands[k])
        vert = rotxyzs[k+1].T
        c2.set_verts([vert])
    elif "A3" == key:
        for obj in [c2,jiku]:
            obj.set_alpha(1-a)
    elif "Sl" == key:
        if k==0:
            waku.set_text("輪切りにすると\nこんな感じ♪")
        c3.set_verts([cut_verts[k]])
    elif "A4" == key:
        for obj in [waku,waku.get_bbox_patch()]:
            obj.set_alpha(1-a)
        L.set_alpha(a)
    elif "L1" == key:
        show_lines(k)
    elif "E" == key:
        last.set_alpha(a)

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