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 Line3D
from matplotlib.patches import FancyArrowPatch
from mpl_toolkits.mplot3d.proj3d import proj_transform
from mpl_toolkits.mplot3d.art3d import Poly3DCollection

#紙の準備
fig = plt.figure()
fig.canvas.draw()
fig.suptitle("2007東北大数学　理系第５問",
             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',computed_zorder=False)
ax1.view_init(35,-55)
ax1.set_xlim(-0.9,1.5)
ax1.set_ylim(-1.2,1.2)
ax1.set_zlim(-1.2,1.2)
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")

#回転体のための関数その１
n = 61
def make_rotxyzs(xyz,rotvec):
    rotxyzs = []
    for i in range(n):
        rvi = rotvec*i/(n-1)
        R = Rotation.from_rotvec(rvi).as_matrix()
        rotxyz = R @ xyz
        rotxyzs.append(rotxyz)
    return rotxyzs

#回転体のための関数その２
col_r = (.5,1,1,.5)
col_u = (0,.5,.5,.1)
col_d = (0,0,1,.1)
def make_bands(rotxyzs,fc):
    bands = []
    for i,rotxyz in enumerate(rotxyzs[1:],1):
        verts = []
        for j in range(1,rotxyz.shape[1]):
            vert = [rotxyzs[i-1][:,j-1],rotxyzs[i-1][:,j],rotxyz[:,j],rotxyz[:,j-1]]
            verts.append(vert)
        band = Poly3DCollection(verts,fc=fc)
        bands.append(band)
    return bands

#回転体を準備（リング）
x1 = [1,1]
y1 = [0,0]
z1 = [1,2]
l = Line3D(x1,y1,z1,color="b")
ax1.add_line(l)
xyz1 = np.array([x1,y1,z1])
rotvec1  = np.array([0,0,2*np.pi])
rot1xyzs = make_rotxyzs(xyz1,rotvec1)
bands1   = make_bands(rot1xyzs,col_r)

#回転体を準備（タイヤ）
theta = np.linspace(0,2*np.pi,n)
x2  = np.cos(theta)
y2  = np.sin(theta)
z2u = 2*np.ones_like(theta)
z2d = 1*np.ones_like(theta)
xyz2u = np.stack([x2,y2,z2u])
xyz2d = np.stack([x2,y2,z2d])
rotvec2 = np.array([2*np.pi,0,0])
rot2xyzus = np.array(make_rotxyzs(xyz2u,rotvec2))
rot2xyzds = np.array(make_rotxyzs(xyz2d,rotvec2))
bands2u = make_bands(rot2xyzus[:,:,:(n+1)//2],col_u)
bands2d = make_bands(rot2xyzds[:,:,:(n+1)//2],col_d)

#リングを回すための関数
def rotate_ring(k):
    verts = []
    for i in range(n-1):
        vert = [rot2xyzds[k,:,i],rot2xyzus[k,:,i],rot2xyzus[k,:,i+1],rot2xyzds[k,:,i+1]]
        verts.append(vert)
    return verts
verts = rotate_ring(0)
ring = Poly3DCollection(verts,fc=col_r)

#輪切りを準備
x = np.linspace(-1,1,n)
r1 = np.sqrt(5-x**2)
r2 = np.sqrt(2-x**2)
cut_verts = []
for i in range(n):
    cx  = x[i]*np.ones_like(theta)
    cy1 = r1[i]*np.cos(theta)
    cz1 = r1[i]*np.sin(theta)
    cy2 = r2[i]*np.cos(-theta)
    cz2 = r2[i]*np.sin(-theta)
    vert = [list(zip(cx,cy1,cz1))+list(zip(cx,cy2,cz2))]
    cut_verts.append(vert)
cut = Poly3DCollection([],fc=col_r)
ax1.add_collection3d(cut)

#積分の計算を準備 
mathdic={
"size" : 16,
"usetex" : True,
"clip_on" : True,
}
lines = [
"=\\pi\\int_{-1}^1(\\,R^2-r^2\\,)\\,dx",
"=\\pi\\int_{-1}^1\\Biggl\\{(\\,y^2+2^2\\,)-(\\,y^2+1^2\\,)\\Biggr\\}\\,dx",
"=\\pi\\int_{-1}^1 3\\,dx",
"=\\pi\\,\\Biggl[\\,3x\\,\\Biggr]_{-1}^1",
"=\\pi\\times3\\times2",
"=6\\,\\pi"
]

#数式を書き込んでいきます
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+15)
    Times1.append(time1)
    X,Y = np.mgrid[-2:12: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)

#３D矢印を使うためのクラス
class Arrow3D(FancyArrowPatch):

    def __init__(self, x, y, z, dx, dy, dz, *args, **kwargs):
        super().__init__((0, 0), (0, 0), *args, **kwargs)
        self._xyz = (x, y, z)
        self._dxdydz = (dx, dy, dz)

    def draw(self, renderer):
        x1, y1, z1 = self._xyz
        dx, dy, dz = self._dxdydz
        x2, y2, z2 = (x1 + dx, y1 + dy, z1 + dz)
        xs, ys, zs = proj_transform((x1, x2), (y1, y2), (z1, z2), self.axes.M)
        self.set_positions((xs[0], ys[0]), (xs[1], ys[1]))
        super().draw(renderer)
        
    def do_3d_projection(self, renderer=None):
        x1, y1, z1 = self._xyz
        dx, dy, dz = self._dxdydz
        x2, y2, z2 = (x1 + dx, y1 + dy, z1 + dz)
        xs, ys, zs = proj_transform((x1, x2), (y1, y2), (z1, z2), self.axes.M)
        self.set_positions((xs[0], ys[0]), (xs[1], ys[1]))
        return np.min(zs)

#半径を示すための矢印を準備
arrow1 = Arrow3D(1,0,-.08,0,2,0,color="0.4")
arrow2 = Arrow3D(1,0,0.08,0,1,0,color="0.3")
arrows = [arrow1,arrow2]
for arrow in arrows:
    arrow.set_mutation_scale(10)
    arrow.set_arrowstyle("<|-|>")
    ax1.add_artist(arrow)
text1 = ax1.text(1,1.2,-.4,"$R$",color="0.4") 
text2 = ax1.text(1,0.2,0.4,"$r$",color="0.3")
text3 = fig.text(.295,.305,"$V$",fontdict=mathdic)
annotations = arrows+[text1,text2,text3]
for obj in annotations:
    obj.set_alpha(0)
    
#下部の解説枠を準備
waku = fig.text(0.5,0.25,"棒があります",
                size=14,color="steelblue",ha="center",fontfamily="HGSoeiKakupoptai",
                bbox=dict(boxstyle="round,pad=1.5",lw=1.5,fc="aliceblue",ec="steelblue"))

#立体は枠からはみ出して
for obj in [l,ring,cut]+bands1+bands2u+bands2d:
    obj.set_clip_on(False)

#回転その１
def rotation1(i):
    ax1.add_collection3d(bands1[i])
    l.set_data_3d(rot1xyzs[i+1])

#回転その２    
def rotation2(i):
    if i==0:
        for obj in bands1:
            obj.set_visible(False)
        ax1.add_collection3d(ring)
    ax1.add_collection3d(bands2u[i])
    ax1.add_collection3d(bands2d[i])
    verts = rotate_ring(i+1)
    ring.set_verts(verts)

#イージング関数
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))
        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,"R1":n-1,"A1":30,"R2":n-1,"A2":30,
           "S1":n,"A3":30,"L1":Times2[-1],"P2":40}
endings = np.cumsum(list(periods.values()))
endings = np.append(0,endings)

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

    if "R1"==key:
        rotation1(k)
    elif "A1"==key:
        l.set_alpha(1-a)
    elif "R2"==key:
        rotation2(k)
    elif "A2"==key:
        ring.set_alpha(col_r[3]*(1-a))
    elif "S1"==key:
        cut.set_verts(cut_verts[k])
    elif "A3"==key:
        for obj in annotations:
            obj.set_alpha(a)
        for obj in [waku,waku.get_bbox_patch()]:
            obj.set_alpha(1-a)
    elif "L1" == key:
        show_lines(k)
    
    if k==0:
        if "R1"==key:
            waku.set_text("棒を回します")
        elif "A1"==key:
            waku.set_text("リングができます")
        elif "R2"==key:
            waku.set_text("リングも回します")
        elif "A2"==key:
            waku.set_text("タイヤができました！")
        elif "S1"==key:
            waku.set_text("輪切りはこんな感じ♪")

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