# Animazione 3D isometrica della fascetta con il campo di temperatura calcolato
# a elementi finiti sulla mesh shell (fem.py).
#
# La coordinata circonferenziale è una direzione risolta del modello FEM,
# quindi la temperatura dipinta su ogni elemento è quella effettivamente
# calcolata, senza ricostruzioni.
#
# Ogni elemento shell è disegnato come faccia piana con colore pari alla media
# dei suoi quattro valori nodali. L'ombreggiatura è calcolata a mano dalla
# normale radiale dell'elemento: Poly3DCollection ombreggia solo alla
# creazione, mentre qui i colori delle facce cambiano a ogni fotogramma.

from pathlib import Path

import matplotlib
import matplotlib.pyplot as plt
import numpy as np
from matplotlib import cm
from matplotlib.animation import FuncAnimation, PillowWriter
from mpl_toolkits.mplot3d.art3d import Poly3DCollection

from config import USCITA
from fem import simula_campo_fem
from mesh import genera_mesh, riepilogo_mesh

# Istante del primo fotogramma mostrato: i primi istanti sono ancora uniformi
# alla temperatura ambiente.
T_INIZIO_ANIMAZIONE_S = 0.40

# Tempo simulato tra un fotogramma e il successivo.
DT_FRAME_S = 0.15

# Millisecondi tra i fotogrammi in riproduzione.
INTERVALLO_RIPRODUZIONE_MS = 60

# Direzione della luce (x = asse del cilindro, y-z = piano della sezione) e
# quota di luce ambiente, per dare volume alla superficie.
DIREZIONE_LUCE = np.array([-0.25, 0.45, 0.85])
LUCE_AMBIENTE = 0.45

# Se True disegna il reticolo degli elementi sopra il campo. Con la mesh fitta
# i bordi degli elementi sul lato nascosto rendono la superficie confusa,
# quindi per default i bordi prendono il colore della faccia: così le facce
# adiacenti combaciano senza lasciare fessure di antialiasing.
MOSTRA_BORDI_ELEMENTI = False


def _illuminazione(mesh_dati: dict) -> np.ndarray:
    # Fattore moltiplicativo di luminosità per elemento, dalla normale
    # radiale uscente valutata nel baricentro dell'elemento.
    facce = mesh_dati["nodi"][mesh_dati["elementi"]]
    baricentri = facce.mean(axis=1)
    normali = baricentri.copy()
    # La normale della superficie cilindrica è radiale: nessuna componente
    # lungo l'asse x.
    normali[:, 0] = 0.0
    normali /= np.linalg.norm(normali, axis=1, keepdims=True)

    luce = DIREZIONE_LUCE / np.linalg.norm(DIREZIONE_LUCE)
    diffusa = np.clip(normali @ luce, 0.0, None)
    return LUCE_AMBIENTE + (1.0 - LUCE_AMBIENTE) * diffusa


def main() -> None:
    mesh_dati = genera_mesh()
    print(riepilogo_mesh(mesh_dati))

    print("Integrazione FEM in corso...")
    dati = simula_campo_fem(dt_frame_s=DT_FRAME_S, mesh_dati=mesh_dati)

    tempi = dati["tempi"]
    indice_inizio = int(np.searchsorted(tempi, T_INIZIO_ANIMAZIONE_S))
    indici_frame = list(range(indice_inizio, len(tempi)))

    T_ambiente = dati["T_ambiente_C"]
    T_max = max(campo.max() for campo in dati["campi"])
    print(
        f"T massima nodale: {T_max:.1f} °C — "
        f"T massima nel punto del sensore: {dati['T_vere'].max():.1f} °C\n"
        f"Skin depth (diagnostica): {1000.0 * dati['skin_depth_m']:.3f} mm"
    )

    elementi = mesh_dati["elementi"]
    facce = mesh_dati["nodi"][elementi]
    illuminazione = _illuminazione(mesh_dati)[:, None]

    raggio_m = mesh_dati["raggio_m"]
    lunghezza_m = mesh_dati["lunghezza_m"]

    sorgente = dati["sorgente"]
    numero_sorgenti = sorgente.get("numero_sorgenti", 1)
    distanza_m = sorgente.get("distanza_sorgenti_m", 0.0)
    theta_sorgente = sorgente["offset_y_percorso_m"] / raggio_m

    norm = matplotlib.colors.Normalize(vmin=T_ambiente, vmax=T_max)
    cmap = matplotlib.colormaps["inferno"]

    fig = plt.figure(figsize=(9, 8.5))
    griglia = fig.add_gridspec(2, 1, height_ratios=[2.6, 1.0], hspace=0.05)
    ax = fig.add_subplot(griglia[0], projection="3d")
    ax_storia = fig.add_subplot(griglia[1])
    ax.view_init(elev=35.264, azim=45)
    # zoom > 1 riempie il riquadro: con set_axis_off gli assi 3D lascerebbero
    # altrimenti molto margine vuoto sopra il pannello della storia.
    ax.set_box_aspect((lunghezza_m, 2 * raggio_m, 2 * raggio_m), zoom=1.12)
    ax.set_axis_off()
    ax.set_xlim(0.0, lunghezza_m)
    ax.set_ylim(-raggio_m, raggio_m)
    ax.set_zlim(-raggio_m, raggio_m)

    collezione = Poly3DCollection(facce, linewidths=0.25, shade=False)
    ax.add_collection3d(collezione)

    marker_sorgenti = ax.plot(
        [], [], [], "o", color="cyan", markersize=6, zorder=10
    )[0]

    mappabile = cm.ScalarMappable(cmap=cmap, norm=norm)
    mappabile.set_array([])
    # La colorbar è agganciata a entrambi i pannelli per non restringere solo
    # quello 3D: così i due riquadri restano della stessa larghezza.
    barra = fig.colorbar(mappabile, ax=(ax, ax_storia), shrink=0.6, pad=0.02, aspect=30)
    barra.set_label("T [°C]")

    titolo = ax.set_title("")

    # Pannello inferiore: storia della temperatura nel punto osservato dal
    # sensore. T_vere è il valore nodale calcolato dal FEM, T_lette lo stesso
    # segnale degradato dal sensore reale (inerzia, rumore, quantizzazione).
    # A parte il rumore le due curve sono quasi sovrapposte: la prima è
    # tracciata spessa e trasparente perché la seconda resti leggibile sopra
    # di essa.
    linea_vera, = ax_storia.plot(
        [], [], color="tab:blue", linewidth=3.0, alpha=0.4,
        label="T nodo del sensore (FEM)",
    )
    linea_letta, = ax_storia.plot(
        [], [], color="tab:red", linewidth=1.0,
        label="T misurata dal sensore (inerzia + rumore + quantizzazione)",
    )
    cursore = ax_storia.axvline(tempi[indice_inizio], color="gray", linewidth=0.8)
    ax_storia.set_xlim(0.0, tempi[-1])
    ax_storia.set_ylim(
        T_ambiente - 0.05 * (T_max - T_ambiente),
        max(dati["T_vere"].max(), dati["T_lette"].max()) * 1.08,
    )
    ax_storia.set_xlabel("Tempo [s]")
    ax_storia.set_ylabel("T [°C]")
    ax_storia.legend(loc="upper left")
    ax_storia.grid(True, alpha=0.3)

    def disegna_frame(k: int):
        T_elementi = dati["campi"][k][elementi].mean(axis=1)
        colori = cmap(norm(T_elementi))
        # Ombreggiatura moltiplicativa sui soli canali RGB.
        colori[:, :3] *= illuminazione
        collezione.set_facecolor(colori)
        collezione.set_edgecolor("#40404060" if MOSTRA_BORDI_ELEMENTI else colori)

        x_sorgenti_m = dati["x_riferimenti"][k] + np.arange(numero_sorgenti) * distanza_m
        visibili = (x_sorgenti_m >= 0.0) & (x_sorgenti_m <= lunghezza_m)
        n_visibili = int(visibili.sum())
        marker_sorgenti.set_data_3d(
            x_sorgenti_m[visibili],
            np.full(n_visibili, raggio_m * np.sin(theta_sorgente) * 1.06),
            np.full(n_visibili, raggio_m * np.cos(theta_sorgente) * 1.06),
        )

        linea_vera.set_data(tempi[: k + 1], dati["T_vere"][: k + 1])
        linea_letta.set_data(tempi[: k + 1], dati["T_lette"][: k + 1])
        cursore.set_xdata([tempi[k], tempi[k]])

        titolo.set_text(
            f"FEM shell — t = {tempi[k]:.2f} s — "
            f"T sensore = {dati['T_lette'][k]:.1f} °C"
        )
        return collezione, marker_sorgenti, linea_vera, linea_letta, cursore, titolo

    animazione = FuncAnimation(
        fig, disegna_frame, frames=indici_frame,
        interval=INTERVALLO_RIPRODUZIONE_MS, blit=False,
    )

    if matplotlib.get_backend().lower() == "agg":
        cartella = Path(USCITA["cartella"])
        cartella.mkdir(parents=True, exist_ok=True)
        percorso = cartella / "animazione_fem.gif"
        animazione.save(
            percorso, writer=PillowWriter(fps=1000 // INTERVALLO_RIPRODUZIONE_MS)
        )
        print(f"Backend non interattivo: animazione salvata in {percorso}")
        return

    plt.show()


if __name__ == "__main__":
    main()