# -*- coding: utf-8 -*-
"""
Created on Thu Jun 29 12:19:27 2017

@author: Étienne Thibierge
"""

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

### Fermer toutes les figures avant de commencer
plt.close("all")

### Couleurs des courbes
colorlist = ['seagreen','firebrick','steelblue','orangered','darkorchid']


"""
RÉSOLUTION PHYSIQUE DE L'ÉQUATION DE LA CHALEUR
"""

### Paramètres à modifier en classe
Dth = .01    # diffusivité
DeltaT = 10 # écart de température


### Parametres physiques et numeriques
L_time=1       # durée totale
L_space=.1

N_time = 20000  # 20000
N_space = 100   # 100

# Increment spatial et temporel
dx=L_space/(N_space-1)
x = np.array([dx*i for i in range(N_space)])

dt=L_time/(N_time-1)


def laplacien(T,dx,Tinit,Tlast):
    n = len(T)
    lapl = np.zeros_like(T)
    for i in range(n):
        if i!=0 and i!=n-1:
            lapl[i] = (T[i+1] - 2*T[i] + T[i-1])/dx**2
        lapl[0] = (T[1] - 2*T[0] + Tinit)/dx**2
        lapl[n-1] = (Tlast - 2*T[n-1] + T[n-2])/dx**2
    return lapl

### Fonction qui fait la résolution
def resolution(Dth, DeltaT):
    print("Le temps caracteristique de diffusion est tdif = "+str(round(L_space**2/Dth))+" min")    
    
    print('Stabilité de la résolution :')
    print('\t D dt =', Dth*L_time/N_time)
    print('\t .5 dx**2=', .5*(L_space/N_space)**2)
    
    if Dth*L_time/N_time > .5*(L_space/N_space)**2:
        print('Résolution instable : il faut diminuer le pas de temps')
        return
    else:
        print('Résolution stable, je fais les calculs')
        
    # État initial
    Ti = np.array([20 for i in range(N_space)])
    Tg = Ti[0] + DeltaT
    Td = Ti[-1]
    T  = [Ti]
    
    # Initialisation de la couleur des courbes    
    img = 0
    
    plt.figure(figsize=(14,8))

    for i in range(N_time):
        delta_T = laplacien(Ti,dx,Tg,Td)
        Ti = Ti + dt * Dth * delta_T
#        for j in range(N_space):
#            Ti[j] = Ti[j] + dt*Dth*delta_T[j]
        T.append(Ti)
    
    # On ne trace que quelques instants
        if i==10 or i== 200 or i ==2000 or i==5000 or i==N_time-1 :
            plt.plot(x,T[i],label=r'$t$='+str(round(i*dt,3)),color=colorlist[img])
            img += 1
    
    plt.grid(True)
    plt.xlabel(r'$x$ (m)')
    plt.ylabel(r'Temp\'{e}rature (\si{\celsius})')
    plt.legend(fontsize='small')
    plt.ylim(Td-2,Tg+2)    

    return T


T = resolution(Dth, DeltaT)


"""
CRÉATION DE L'ANIMATION
"""

### Paramètres de la vidéo
duree = 15  # durée en secondes
ips   = 24  # images par seconde
nb_images = duree*ips   # nbre total d'image

### Création de la figure sur laquelle se fait l'animation
fig = plt.figure(figsize=(14,8))
ax = fig.add_subplot(111)

plt.grid(True)
plt.xlabel('x (m)')
plt.ylabel('T (C)')
plt.ylim(T[0][-1]-2,T[-1][0]+2)
plt.xlim(np.min(x),np.max(x))

### Mise en forme extérieure à la boucle pour effacement entre chaque image
line, = ax.plot([], [], '-', color='r')

def init():          						# Initialisation avec du vide
    line.set_data([], [])
    return line,

def animate(i):
    global T    
    
    # i est le numéro de l'image dans le film
    t = i/ips * L_time/duree
    j = int(t/dt)
    
    ### Température à l'instant t
    temperature = T[j]
    line.set_data(x, temperature)
    
    return line,

### Pour démarrer et arrêter l'animation à chaque clic
anim_running = True
def onClick(event):
    global anim_running
    if anim_running:
        film.event_source.stop()
        anim_running = False
    else:
        film.event_source.start()
        anim_running = True
fig.canvas.mpl_connect('button_press_event', onClick)

film = anim.FuncAnimation(fig, animate, frames=nb_images, interval=20)