import numpy as np
from matplotlib import pyplot as plt, animation
from mpl_toolkits.axes_grid1 import make_axes_locatable
plt.rcParams["figure.figsize"] = [5.00, 10.0]
plt.rcParams["pcolor.shading"] = 'auto'

def create_animation(t,x1,x2,tensor0,tensor1,tensor2,filename='output.gif'):
    n = len(t)
    fig, ax = plt.subplots(3)
    ax[0].set_title('x1 motion')
    ax[1].set_title('x2 motion')
    ax[2].set_title('x3 motion')
    pcolor_ax0 = ax[0].pcolor(x1, x2, tensor0[0,:,:], vmin=tensor0.min(), vmax=tensor0.max())
    pcolor_ax1 = ax[1].pcolor(x1, x2, tensor1[0,:,:], vmin=tensor1.min(), vmax=tensor1.max())
    pcolor_ax2 = ax[2].pcolor(x1, x2, tensor2[0,:,:], vmin=tensor2.min(), vmax=tensor2.max())
    div0 = make_axes_locatable(ax[0])
    div1 = make_axes_locatable(ax[1])
    div2 = make_axes_locatable(ax[2])
    cax0 = div0.append_axes('right', size='5%', pad=0.2)
    cax1 = div1.append_axes('right', size='5%', pad=0.2)
    cax2 = div2.append_axes('right', size='5%', pad=0.2)
    def update(i):
        pcolor_ax0.set_array(tensor0[i, :, :].flatten())
        pcolor_ax1.set_array(tensor1[i, :, :].flatten())
        pcolor_ax2.set_array(tensor2[i, :, :].flatten())
    fig.colorbar(pcolor_ax0, cax=cax0)
    fig.colorbar(pcolor_ax1, cax=cax1)
    fig.colorbar(pcolor_ax2, cax=cax2)
    anim = animation.FuncAnimation(fig, update, interval=1000*t[-1]/len(t), frames=n)
    anim.save(filename)

def g(x,r):
    return x/r

def uhat_prime(t,r):
    a = 5600 # m/s
    return np.exp(t-r/a-6)/(np.exp(t-r/a-6)+1)**2 # u_bar'

t = np.linspace(0,12,300) # t in [0,12] sec 
x1 = np.linspace(5000, 15000, 64) # x1 in [5km, 15km]
x2 = np.linspace(-5000, 5000, 64) # x2 in [-5km, 5km] 
x3 = 10000.0 # x3 = 10 km

u1 = np.zeros((len(t),len(x1),len(x2))) # (300, 64, 64)
u2 = np.zeros((t.shape[-1],x1.shape[-1],x2.shape[-1]))
u3 = np.zeros((t.shape[-1],x1.shape[-1],x2.shape[-1]))

for i, _t in enumerate(t):
    for j, _x1 in enumerate(x1):
        for k, _x2 in enumerate(x2):
            r = np.sqrt(_x1**2 + _x2**2 + x3**2)
            g1, g2, g3 = g(_x1,r), g(_x2,r), g(x3,r)
            u = uhat_prime(_t,r)
            u1[i,j,k] = g1**2 * g3 * u/r
            u2[i,j,k] = g1 * g2 * g3 * u/r
            u3[i,j,k] = g1 * g3**2 * u/r

create_animation(t,x1,x2,np.abs(u1),np.abs(u2),np.abs(u3))


#plt.figure()
#plt.plot(t,u1[:,31,31])
#plt.show()