import numpy as np
import matplotlib.pyplot as plt
import matplotlib.animation as ani
from mpl_toolkits.mplot3d import proj3d
from scipy.spatial.transform import Rotation
from mpl_toolkits.mplot3d.art3d import Poly3DCollection

#紙の準備
fig = plt.figure()
fig.suptitle("2015京大数学　文系第４問",
             color="0.7",ha="right",x=0.96,y=0.08)
fig.subplots_adjust(left=0.27,right=0.72)
ax=fig.add_subplot(projection='3d')
ax.set_xlim(-1,1)
ax.set_ylim(-1,1)
ax.set_zlim(0,2)
ax.set_box_aspect((1,1,1))
ax.view_init(10,-75)
ax.axis("off")

#まんまるの球
u = np.linspace(0,np.pi,25)
v = np.linspace(0,2*np.pi,49)
Sx = np.outer(np.sin(u),np.sin(v))
Sy = np.outer(np.sin(u),np.cos(v))
Sz = np.outer(np.cos(u),np.ones_like(v))+1
sphere = ax.plot_surface(Sx,Sy,Sz,
                         color="cyan",alpha=0.2)

#光が当たる円
N = 100
theta = np.linspace(-np.pi,np.pi,N)
rot = Rotation.from_rotvec(np.array([0,np.pi/4,0]))
C = np.array([np.cos(theta)/np.sqrt(2)-1/np.sqrt(2),
              np.sin(theta)/np.sqrt(2),
              np.zeros_like(theta)])
q = rot.as_matrix() @ C + np.array([np.ones_like(theta),
                                    np.zeros_like(theta),
                                    np.ones_like(theta)])
circle = ax.plot_trisurf(q[0,:],q[1,:],q[2,:],
                         color="k",alpha=0)

#地面にできる影
r = np.zeros((3,N-2))
for i in range(1,N-1):
    beta = 2/(2-q[2,i])
    r[0,i-1] = beta*(q[0,i]-1)+1
    r[1,i-1] = beta* q[1,i]
shade = ax.plot_trisurf(r[0,:],r[1,:],r[2,:],
                        clip_on=False,color="k",alpha=0)

#球を包みこむ円錐
polys = []
for i in range(N-3):
    poly = [(1,0,2),
            (r[0,i],r[1,i],r[2,i]),
            (r[0,i+1],r[1,i+1],r[2,i+1])]
    polys.append(poly)
lastpoly = [(1,0,2),
            (r[0,-1],r[1,-1],r[2,-1]),
            (r[0,0],r[1,0],r[2,0])]
polys.append(lastpoly)
cone = Poly3DCollection(polys,clip_on=False,facecolor="yellow",alpha=0) 
ax.add_collection3d(cone)

#x軸とy軸を引きます
ax.plot([-2.5,2.5],[0,0],[0,0],clip_on=False,color="0.9")
ax.plot([0,0],[-3.5,3.5],[0,0],clip_on=False,color="0.9")
xend = [( 2.5,-0.07,0),( 2.5,0.07,0),( 2.5+0.14*np.sin(np.pi/3),0,0)]
yend = [(-0.07,-3.5,0),(0.07,-3.5,0),(0,-3.5-0.14*np.sin(np.pi/3),0)]
ends = Poly3DCollection([xend,yend],clip_on=False,facecolor="0.9")
ax.add_collection3d(ends)
ax.text(2.75,0,0,"$x$",color="0.7",ha="center",va="center")
#見栄えがよくないので、y軸の先端だけ動かします
ytx,yty,_ = proj3d.proj_transform(0,-4.0,0,ax.get_proj())
ytxt = ax.text2D(ytx,yty,"$y$",color="0.7",ha="center",va="center")
#答えとなる領域の数式
ftxt = fig.text(0.1,0.2,"$x \\leq -\\frac{\\:1\\:}{\\:4\\:}\\,y^2+1$",
                color="0.4",fontsize=20,usetex=True,alpha=0)

#動く点たち
PR,    = ax.plot([],[],[],clip_on=False,color="khaki")
histR, = ax.plot([],[],[],clip_on=False,color="0.7")
histQ, = ax.plot([],[],[],clip_on=False,color="steelblue")
R,     = ax.plot([],[],[],clip_on=False,marker="o",color="lightslategrey")
Q,     = ax.plot([],[],[],marker="o",color="dodgerblue")
P,     = ax.plot( 1, 0, 2,marker="o",color="darkorange",markersize=10,zorder=10)

#お気に入りのイージング関数
def easing(x):
    if x<0.5:
        return 2*x**2
    else:
        return 1-(-2*x+2)**2/2

#軌跡を描くための関数
def trajectory(hist,vec):
    x = np.append(hist.get_data_3d()[0],vec[0])
    y = np.append(hist.get_data_3d()[1],vec[1])
    z = np.append(hist.get_data_3d()[2],vec[2])
    hist.set_data_3d(x,y,z)

#台本
def update(i):
    
    #地面に線をひきます
    if i < N-2:
        PR.set_data_3d([1,r[0,i]],[0,r[1,i]],[2,r[2,i]])
        trajectory(histR,r[:,i])
        trajectory(histQ,q[:,i+1])
        R.set_data_3d(r[0,i],r[1,i],r[2,i])
        Q.set_data_3d(q[0,i+1],q[1,i+1],q[2,i+1]) 
    elif i < 100:
        pass
    
    #光をあてます
    elif i < 200:
        k = easing((i-100)/99)
        for obj in [PR,R,Q]:
            obj.set_alpha(1-k)
        for obj in [circle,shade,cone]:
            obj.set_alpha(0.2*k)
    
    #上から見下ろします
    elif i < 300:
        k = easing((i-200)/99)
        ax.view_init(10+80*k,-75+75*k)
        ytx,yty,_ = proj3d.proj_transform(0,-4+0.25*k,0,ax.get_proj())
        ytxt.set_position((ytx,yty))
    
    #数式をお見せします
    elif i < 380:
        k = easing((i-300)/79)
        for obj in [histQ,P]:
            obj.set_alpha(1-k)
        for obj in [sphere,circle,cone]:
            obj.set_alpha(0.2-0.2*k)
    elif i < 420:
        k = easing((i-380)/39)
        ftxt.set_alpha(k)

#上演
mov = ani.FuncAnimation(fig,update,450,interval=100)
plt.show()