import numpy as np
import matplotlib.pyplot as plt
import matplotlib.animation as ani
import matplotlib.gridspec as gridspec

#紙の準備
fig = plt.figure()
fig.canvas.draw()
fig.suptitle("2008東大数学　理系第３問",
             color="0.5",ha="right",x=0.95,y=0.95,zorder=-1)
gs = gridspec.GridSpec(4,6)
#見取図
ax1 = fig.add_subplot(gs[:3,:3],projection="3d")
ax1.view_init(45,-90)
ax1.axis("off")
#断面図
ax2 = fig.add_subplot(gs[:3,3:])
ax2.set_aspect("equal")
ax2.set_xlim(-1,1)
ax2.set_ylim(-0.9,1.1)
ax2.axis("off")
#数式の部分
ax3 = fig.add_subplot(gs[3,:])
ax3.set_xlim(0,1)
ax3.set_ylim(0,1)
ax3.axis("off")

#正八面体(骨格)の準備
X0 = np.array([np.cos(i*np.pi/3) for i in range(8)])/np.sqrt(3)
Y0 = np.array([np.sin(i*np.pi/3) for i in range(8)])/np.sqrt(3)
Z0 = np.array([-1+2*(i%2) for i in range(8)])/np.sqrt(6)
col_frame = "slateblue"
ax1.plot(X0[1:4],Y0[1:4],Z0[1:4],color=col_frame,linestyle="dashed")
ax1.plot(X0[3:],Y0[3:],Z0[3:],color=col_frame)
ax1.plot(X0[0:6:2],Y0[0:6:2],Z0[0:6:2],color=col_frame,linestyle="dashed")
ax1.plot(X0[4::2],Y0[4::2],Z0[4::2],color=col_frame)
ax1.plot(X0[1::2],Y0[1::2],Z0[1::2],color=col_frame)
ax2.plot(X0[1:],Y0[1:],color=col_frame)
ax2.plot(X0[0::2],Y0[0::2],color=col_frame,linestyle="dashed")
ax2.plot(X0[1::2],Y0[1::2],color=col_frame)

#変数の準備
NT = 21
T  = np.linspace(-1,1,NT)
Ntheta = 41
theta = np.linspace(0,2*np.pi,Ntheta)

#断面でカットする関数
def cut_surface(imgs,list_t):
    
    X = []
    Y = []
    Z = []
    COLOR = []
    
    for t in list_t:
        
        #断面の円
        r = np.sqrt(2*t**2/6+1)/2
        x = r*np.cos(theta)
        y = r*np.sin(theta)
        z = t*np.ones(Ntheta)/np.sqrt(6)
        
        #断面の六角形
        s  = (t+1)/2
        x0 = (1-s)*X0[[0,2,2,4,4,6]]+s*X0[[1,1,3,3,5,5]]
        y0 = (1-s)*Y0[[0,2,2,4,4,6]]+s*Y0[[1,1,3,3,5,5]]
        z0 = t*np.ones(6)/np.sqrt(6)
        
        color = (s/2,s,1)
        img1 =  [ax1.plot_trisurf(x,y,z,color=color,alpha=0.4)]
        img1 += ax1.plot(np.append(x0,x0[0]),np.append(y0,y0[0]),np.append(z0,z0[0]),color="orange",lw=2)
        img1 += [ax1.scatter(x0,y0,z0,color="darkorange")]
        img2 =  ax2.fill(x,y,color=color,alpha=0.5)
        img2 += ax2.fill(x0,y0,fc="w",ec="orange",lw=2,alpha=0.5)
        img1 += [ax2.scatter(x0,y0,color="darkorange",zorder=10)]
        text2 = [ax2.text(0,0.9,"$z={"+'{:+.1f}'.format(t)+"}\\,/_\\sqrt{6}}$",ha="center",va="top")]
        
        #1断面につき0.5秒
        for i in range(5):
            imgs.append(img1+img2+text2)
            
        X = np.hstack([X,x])
        Y = np.hstack([Y,y])
        Z = np.hstack([Z,z])
        COLOR.append(color)
    
    return X,Y,Z,COLOR

#行き１回目
imgs=[]
X,Y,Z,COLOR = cut_surface(imgs,T)

#回転体の表面を描画
for i in range(NT-1):
    for j in range(Ntheta-1):
        index = [i*Ntheta+j,i*Ntheta+j+1,(i+1)*Ntheta+j,(i+1)*Ntheta+j+1]
        ax1.plot_trisurf(X[index],Y[index],Z[index],
                          color=COLOR[i],alpha=0.3,shade=False)

#帰り１回目
cut_surface(imgs,-T)
#行き２回目
cut_surface(imgs,T)

#積分の式も表示
str_int = r"$\displaystyle V = \pi \int_{-\frac{1}{\sqrt{6}}}^{\frac{1}{\sqrt{6}}} \frac{1+2\,t^2}{4} \, dt = \frac{5\sqrt{6}}{54}\,\pi$"
ax3.text(0.5,0.5,str_int,
          usetex=True,fontsize=18,ha="center",va="center")
ax3.fill([0.42,0.58,0.58,0.42],[0,0,1,1],color="c",alpha=0.2)
ax3.text(0.5,1.1,"断面(円)の半径の２乗",ha="center",color="c")

mov=ani.ArtistAnimation(fig, imgs, 100)
plt.show()