"""Résolution d'une discrétisatin de l'équation de la chaleur.

Fait en classe le 7.3.2023.
"""

import numpy as np
import matplotlib.pyplot as plt

# paramètres du modèle
DIFFUSION = 1  # "alpha": paramètre de diffusion
T_FINAL = 10
L_BARRE = 20
AMPLITUDE_FLAMME = 1  # A dans l'énoncé
FLAMME_A = 1  # inverse de dispersion de la flamme en temps
FLAMME_B = 1  # inverse de dispersion de la flamme en espace
T_FLAMME = 3  # temps auquel la flamme est la plus intense

# paramètres de discrétisation
N_ETAPES = 500
DT = T_FINAL / N_ETAPES
N_INTERVALLES = 50
DX = L_BARRE / N_INTERVALLES
K = DIFFUSION*DT/DX**2

# paramètres visuels
T_PAUSE = 0.01  # en s
TEMP_MAX = 3


def flamme(x, t):
    """Modèle de la flamme avec une gaussienne en espace et temps.

    g(t,x)=A⋅exp(−a⋅(x−L/2)^2−b⋅(t−tc)^2).
    """
    return AMPLITUDE_FLAMME * np.exp(
        - FLAMME_A * (x - L_BARRE/2)**2
        - FLAMME_B * (t - T_FLAMME)**2
        )


def creer_graphique(positions, flamme, temperature):
    """Créer le graphique des températures au-dessus de celui de la flamme."""
    fig, axs = plt.subplots(2, 1)
    axs[0].set_title("Température de la barre")
    axs[0].set_ylabel("température")
    axs[0].set_ylim(0, TEMP_MAX)
    courbe_temperature, = axs[0].plot(positions, temperature, "b")
    axs[1].set_ylabel("chaleur externe")
    axs[1].set_xlabel("position")
    axs[1].set_ylim(0, AMPLITUDE_FLAMME)
    courbe_flamme, = axs[1].plot(positions, flamme, "r")
    return courbe_flamme, courbe_temperature
    
    
def evolution():
    """Faire évoluer le système et en montrer l'animation."""
    positions = np.linspace(0, L_BARRE, N_INTERVALLES + 1)
    chaleur = np.zeros(N_INTERVALLES + 1)
    temperature = np.zeros(N_INTERVALLES + 1)
    courbe_flamme, courbe_temperature = creer_graphique(positions,
                                                        chaleur, temperature)
    for temps in np.linspace(0, T_FINAL-DT, N_ETAPES):
        temp_nouvelle = np.zeros(N_INTERVALLES + 1)
        for pos in range(1, N_INTERVALLES-1):
            x = positions[pos]
            chaleur[pos] = flamme(x, temps)
            temp_nouvelle[pos] = (temperature[pos]
                                  + K*(temperature[pos+1] - 2*temperature[pos]
                                       + temperature[pos-1])
                                  + DT*chaleur[pos])
        temperature = temp_nouvelle
        courbe_flamme.set_ydata(chaleur)
        courbe_temperature.set_ydata(temperature)
        plt.pause(T_PAUSE)
            
        
# PROGRAMME PRINCIPAL
evolution()
plt.show()