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

## question 2
def cinetique(Tmax, n, alpha, beta, a0, b0, c0):
    t, a, b, c = [0], [a0], [b0], [c0]
    dt = Tmax/n
    for k in range(0,n-1):
        t.append(t[k] + dt)
        a.append(a[k] - alpha*a[k]**2*dt)
        b.append(b[k] + (alpha*a[k]**2-beta*b[k]**2)*dt)
        c.append(c[k] + beta*b[k]**2*dt)
    return t, a, b, c

## question 3
alpha, beta = 10,1
Tmax = 30
n = 1000
a0, b0, c0 = 1, 0, 0

t, a, b, c = cinetique(Tmax, n, alpha, beta, a0, b0, c0)

plt.plot(t,a,label="[A]")
plt.plot(t,b,label="[B]")
plt.plot(t,c,label="[C]")
plt.legend(loc = "best")
plt.show()

## question 4
def u_initial(Tcool,Thot,dx,dy,nx,ny,r0,r1):
    u0 = Tcool * np.ones((nx,ny))
    xm, ym = nx*dx//2, ny*dy//2
    for i in range(nx):
        for j in range(ny):
            x,y = i*dx, j*dy
            if r0**2 <= (x-xm)**2 + (y-ym)**2 <= r1**2:
                u0[i,j] = Thot
    return u0

# def u_initial(Tcool,Thot,dx,dy,nx,ny,r0,r1):
#     u0 = ... * np.ones((..., ...))
#     xm, ym = nx*dx//2, ny*dy//2
#     for i in range():
#         for j in range(ny):
#             x,y = ..., ...
#             if ... <= (x-xm)**2 + (y-ym)**2 <= ...:
#                 u0[i,j] = ...
#     return u0

# test
u0 = u_initial(300,700,0.1,0.1,100,100,4,5)
plt.imshow(u0)
plt.show()

## question 5
def iteration(u,D,dt,dx,dy):
    u2 = np.copy(u)
    nx, ny = np.shape(u)
    for i in range(1,nx-1):
        for j in range(1,ny-1):
            ux = (u[i+1,j] - 2*u[i,j] + u[i-1,j])/((dx)**2)
            uy = (u[i,j+1] - 2*u[i,j] + u[i,j-1])/((dy)**2)
            u2[i,j] = u[i,j] + D*dt*(ux + uy)
    return u2


## Question 6 - Test (cellule à exécuter d'un bloc)

# INITIALISATION

# longueur et largeur de la plaque (mm)
L = l = 10
# pas d'espace dans les directions x et y (mm)
dx = dy = 0.1
dx2, dy2 = dx*dx, dy*dy

# coefficient de diffusion thermique de l'acier (mm^2.s^(-1))
D = 4.

# températures initiales de la plaque et de l'anneau (K)
Tcool, Thot = 300, 1900

# taille des matrices de température
nx, ny = int(L/dx), int(l/dy)

# pas de temps
dt = dx2 * dy2 / (2 * D * (dx2 + dy2))
#print("Delta_t = ", dt)

# diamètres intérieur et extérieur de l'anneau (mm)
r0, r1 = 1.5, 4

# matrice des température à t = 0
u0 = u_initial(Tcool,Thot,dx,dy,nx,ny,r0,r1)
u = np.copy(u0)

# nombre d'instant
n = 1000
print("t_max = ", n*dt)
# ANIMATION

# affichage de la matrice initiale
fig = plt.figure()
ax = fig.add_subplot(1,1,1)
im = ax.imshow(u0, cmap=plt.get_cmap('magma'), vmin=Tcool, vmax=Thot,
               interpolation='bicubic')
ax.set_axis_off()
ax.set_title('0.0 ms')

# ajout de la barre de couleur à droite
fig.subplots_adjust(right=0.9)
cbar_ax = fig.add_axes([0.85, 0.15, 0.03, 0.7])
cbar_ax.set_xlabel('$T$ / K', labelpad=20)
fig.colorbar(im, cax=cbar_ax)

def animate(k):
    """Advance the simulation and animation by one time step."""
    global u0, u
    u0, u = u, iteration(u0,D,dt,dx,dy)
    u0, u = u, iteration(u0,D,dt,dx,dy)
    ax.set_title('{:.1f} ms'.format(2*k*dt*1000))
    im.set_data(u)

interval = 10
ani = animation.FuncAnimation(fig, animate, frames=n, repeat=False,
                              interval=interval)

#sauvegarde de l'animation en gif
writer = animation.PillowWriter(fps=250, bitrate=1800)
#ani.save('equation_chaleur.gif', writer=writer)
plt.show()