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

'''
座標のセッティング
'''
def ax_settings(ax,left,right,bottom,top):
    ax.set_xlim(left,right)
    ax.set_ylim(bottom,top)
    ax.set_xticks(np.arange(left,right+1,1))
    ax.set_yticks(np.arange(bottom,top+1,1))
    ax.axhline(0,color='0.2')
    ax.axvline(0,color='0.2')
    ax.set_aspect("equal")
    ax.grid()

'''
放物線
'''
def f(x):
    y = x**2 + a*x + b
    return y

'''
点(a,b)の経路
'''
def path(T,i):
    t = T[i]
    if t<2*np.sqrt(2):
        t = t/np.sqrt(2)
        a = t
        b = t-2
    elif t<2*np.sqrt(2)+4:
        t = t-2*np.sqrt(2)
        a = 2-t
        b = 0
    else:
        t = (t-(2*np.sqrt(2)+4))/np.sqrt(2)
        a = t-2
        b = -(t-2)-2
    return a,b

'''
図の準備
'''
fig = plt.figure()
fig.suptitle("2021東大数学　理系第１問 文系第３問",x=0.5, y=0.9)
ax1 = fig.add_subplot(121)
ax1.set_title("xy平面")
ax2 = fig.add_subplot(122)
ax2.set_title("ab平面")
ax_settings(ax1,-4,4,-4,4)
ax_settings(ax2,-4,4,-4,4)

'''
変数の準備
'''
x = np.arange(-4,4,0.1)
T = np.arange(0,4+4*np.sqrt(2),0.1)

'''
動画を作成する
'''
imgs = []
for i in range(len(T)):
    
    img1 = ax1.plot()
    img2 = ax2.plot()
    
    '''
    過去の軌跡
    '''
    for j in range(i):
        a,b = path(T,j)
        img1 += ax1.plot(x,f(x),color=(0,0,0,0.2))
        img2 += ax2.plot(a,b,marker=".",color=(1,0,0,0.2))
        
    '''
    現在地
    '''
    a,b = path(T,i)
    img1 += ax1.plot(x,f(x),color="k")
    img2 += ax2.plot(a,b,marker="o",color=(1,0,0))
    
    '''
    パラメータの表示
    '''
    img2 +=[ax2.text(1.2,3.2,"a= "+"{:4.1f}".format(a))]
    img2 +=[ax2.text(1.2,2.2,"b= "+"{:4.1f}".format(b))]
        
    img = img1 + img2
    imgs.append(img)
    
mov=ani.ArtistAnimation(fig, imgs, 50)
plt.show()