# Animazione della superficie della fascetta durante il passaggio delle
# sorgenti.
#
# Riproduce la fisica di run_0001 (stesso seed di simulate.py) e mostra:
# - il profilo di flusso termico q(x) lungo la generatrice delle sorgenti;
# - il campo di temperatura sulla superficie cilindrica sviluppata in piano,
#   cioè T(x, arco) con l'arco centrato sulla generatrice del sensore;
# - il sensore infrarosso e la temperatura nel punto osservato.
#
# La mappa sviluppata è periodica in verticale: il bordo superiore e quello
# inferiore sono lo stesso punto della fascetta.

import random
from pathlib import Path

import matplotlib
import matplotlib.pyplot as plt
import numpy as np
from matplotlib.animation import FuncAnimation, PillowWriter
from matplotlib.collections import LineCollection

import sensore as sens
import sorgente as src
from config import SIMULAZIONE
from mesh import campo_su_griglia, chiudi_campo
from simulate import (
    campo_iniziale,
    configurazione_randomizzata,
    passo_termico,
    prepara_stato_termico,
)

# Istante di inizio dei fotogrammi mostrati (la simulazione parte comunque da 0).
T_INIZIO_ANIMAZIONE_S = 0.40

# Istante di fine dell'animazione.
T_FINE_ANIMAZIONE_S = 30

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

# Millisecondi tra i fotogrammi in riproduzione.
INTERVALLO_RIPRODUZIONE_MS = 30

# Disegna i bordi degli elementi shell sopra il campo di temperatura.
# Con la mesh di default (100 x 110, elementi di 1 x 2 mm) il reticolo è più
# fitto dei pixel disponibili e copre il campo: per ispezionare la mesh
# conviene attivarlo e insieme abbassare n_elementi_x e n_elementi_theta in
# config.py, per esempio a 25 x 28.
MOSTRA_MESH = False


def simula_campi(cfg_run: dict) -> dict:
    # Esegue la simulazione fino a T_FINE_ANIMAZIONE_S salvando, a ogni
    # fotogramma, campo di temperatura, profilo di flusso e stato del sensore.
    #
    # I campi sono restituiti come griglie (n_nodi_x, n_elementi_theta) con la
    # colonna dell'arco ruotata in modo da portare la generatrice del sensore
    # (theta = 0) al centro: è solo una scelta di visualizzazione.
    fascetta = cfg_run["fascetta"]
    aria = cfg_run["aria"]
    sorgente = cfg_run["sorgente"]
    sensore = cfg_run["sensore"]

    stato = prepara_stato_termico(fascetta, aria, sorgente, sensore)
    mesh = stato["mesh"]
    dt = stato["dt_s"]
    interpolatore = stato["interpolatore_sensore"]

    # Dopo np.roll(..., rotazione) la colonna theta = 0 finisce all'indice
    # "rotazione", quindi l'arco con segno della colonna m vale
    # (m - rotazione) * lato_arco.
    rotazione = mesh["n_elementi_theta"] // 2
    arco_centrato_m = (
        np.arange(mesh["n_elementi_theta"]) - rotazione
    ) * mesh["lato_arco_m"]

    T = campo_iniziale(stato)
    T_letta = sens.temperatura_osservata_C(interpolatore, T)

    tempi, campi, flussi, x_riferimenti = [], [], [], []
    T_vere, T_lette = [], []

    prossimo_frame_t = 0.0
    t = 0.0
    while t <= T_FINE_ANIMAZIONE_S + 1e-12:
        T, info = passo_termico(stato, T, t)
        T_vera = sens.temperatura_osservata_C(interpolatore, T)
        T_letta = sens.aggiorna_inerzia_C(
            T_letta, T_vera, dt, sensore["costante_tempo_s"]
        )

        if t + 1e-12 >= prossimo_frame_t:
            _, q_nodi = src.profilo_flusso_x_W_m2(
                sorgente, stato["x_sensore_m"], t, mesh["x_nodi_m"]
            )
            tempi.append(t)
            campi.append(np.roll(campo_su_griglia(mesh, T), rotazione, axis=1))
            flussi.append(q_nodi)
            x_riferimenti.append(info["x_riferimento_m"])
            T_vere.append(T_vera)
            T_lette.append(T_letta)
            prossimo_frame_t += DT_FRAME_S

        t += dt

    return {
        "tempi": np.array(tempi),
        "campi": campi,
        "flussi": flussi,
        "x_riferimenti": np.array(x_riferimenti),
        "T_vere": np.array(T_vere),
        "T_lette": np.array(T_lette),
        "mesh": mesh,
        "rotazione": rotazione,
        "x_nodi_mm": mesh["x_nodi_m"] * 1000.0,
        "arco_centrato_mm": arco_centrato_m * 1000.0,
        "arco_chiuso_mm": np.append(
            arco_centrato_m, arco_centrato_m[-1] + mesh["lato_arco_m"]
        ) * 1000.0,
        "lunghezza_mm": fascetta["lunghezza_mm"],
        "x_sensore_mm": sensore["x_mm"],
        "offset_sorgente_mm": sorgente["offset_y_percorso_m"] * 1000.0,
        "sorgente": sorgente,
        "T_ambiente_C": aria["temperatura_ambiente_C"],
    }


def main() -> None:
    rng = random.Random(SIMULAZIONE["seed"])
    cfg_run = configurazione_randomizzata(1, rng)
    dati = simula_campi(cfg_run)

    tempi = dati["tempi"]
    indice_inizio = int(np.searchsorted(tempi, T_INIZIO_ANIMAZIONE_S))
    n_frame = len(tempi) - indice_inizio

    lunghezza_mm = dati["lunghezza_mm"]
    x_vista_mm = (-10.0, lunghezza_mm + 10.0)
    arco_mm = dati["arco_chiuso_mm"]
    arco_min_mm, arco_max_mm = arco_mm[0], arco_mm[-1]

    # Il minimo evita un asse degenere se la finestra dell'animazione si
    # chiude prima che le sorgenti entrino nella fascetta.
    q_max_MW = max(max(q.max() for q in dati["flussi"]) / 1e6, 1e-3)
    T_max = max(c.max() for c in dati["campi"])

    sorgente = dati["sorgente"]
    numero_sorgenti = sorgente.get("numero_sorgenti", 1)
    distanza_mm = sorgente.get("distanza_sorgenti_m", 0.0) * 1000.0
    offset_mm = dati["offset_sorgente_mm"]

    fig, (ax_flusso, ax_superficie, ax_storia) = plt.subplots(
        3, 1, figsize=(10, 8), height_ratios=[1.0, 1.6, 1.2],
        gridspec_kw={"hspace": 0.45},
    )
    fig.suptitle("Superficie sviluppata della fascetta: sorgenti in transito e sensore")

    # Pannello 1: profilo di flusso lungo x e posizioni delle sorgenti.
    linea_flusso, = ax_flusso.plot([], [], color="tab:red")
    marker_sorgenti, = ax_flusso.plot(
        [], [], "v", color="tab:red", markersize=10, clip_on=False
    )
    ax_flusso.annotate(
        "verso di marcia",
        xy=(0.28, 0.85), xytext=(0.55, 0.85), xycoords="axes fraction",
        arrowprops={"arrowstyle": "->", "color": "gray"},
        color="gray", va="center",
    )
    ax_flusso.set_xlim(*x_vista_mm)
    ax_flusso.set_ylim(0.0, q_max_MW * 1.25)
    ax_flusso.set_ylabel("q(x) [MW/m²]")
    ax_flusso.grid(True, alpha=0.3)
    ax_flusso.set_xticklabels([])

    # Pannello 2: temperatura sulla superficie sviluppata. L'asse verticale è
    # l'arco circonferenziale, con 0 sulla generatrice del sensore.
    #
    # pcolormesh con shading "gouraud" interpola linearmente tra i valori
    # nodali: è esattamente ciò che fanno le funzioni di forma bilineari
    # dell'elemento Q4, quindi il disegno è la soluzione dell'elemento finito
    # e non un'interpolazione grafica arbitraria.
    Xg, Yg = np.meshgrid(dati["x_nodi_mm"], dati["arco_chiuso_mm"], indexing="ij")
    mappa = ax_superficie.pcolormesh(
        Xg, Yg,
        chiudi_campo(dati["campi"][indice_inizio]),
        shading="gouraud",
        cmap="inferno",
        vmin=dati["T_ambiente_C"],
        vmax=T_max,
    )
    if MOSTRA_MESH:
        # Bordi degli elementi shell: linee lungo x a ogni nodo
        # circonferenziale e viceversa.
        segmenti = [
            [(x, arco_min_mm), (x, arco_max_mm)] for x in dati["x_nodi_mm"]
        ] + [
            [(0.0, a), (lunghezza_mm, a)] for a in dati["arco_chiuso_mm"]
        ]
        ax_superficie.add_collection(
            LineCollection(segmenti, colors="white", linewidths=0.5, alpha=0.55)
        )
    ax_superficie.set_xlim(*x_vista_mm)
    ax_superficie.set_ylim(arco_min_mm, arco_max_mm)
    ax_superficie.set_ylabel("arco circonferenziale [mm]")
    ax_superficie.set_xlabel("x [mm]")
    ax_superficie.plot(
        [dati["x_sensore_mm"]], [0.0], "^", color="tab:blue", markersize=11
    )
    ax_superficie.text(
        dati["x_sensore_mm"] + 3, 0.0, "sensore IR", color="tab:blue", va="center"
    )
    ax_superficie.axhline(
        offset_mm, color="tab:cyan", linewidth=0.8, linestyle="--", alpha=0.7
    )
    # La colorbar è agganciata a tutti i pannelli per non restringere solo
    # quello della superficie, mantenendo allineati gli assi x.
    barra = fig.colorbar(
        mappa, ax=(ax_flusso, ax_superficie, ax_storia), pad=0.02, aspect=35
    )
    barra.set_label("T [°C]")

    # Pannello 3: temperatura nel punto osservato dal sensore.
    linea_vera, = ax_storia.plot([], [], label="T vera lato interno")
    linea_letta, = ax_storia.plot([], [], label="T sensore (con inerzia)")
    cursore = ax_storia.axvline(tempi[indice_inizio], color="gray", linewidth=0.8)
    ax_storia.set_xlim(0.0, T_FINE_ANIMAZIONE_S)
    ax_storia.set_ylim(15.0, 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)

    testo_tempo = ax_flusso.set_title(f"t = {tempi[indice_inizio]:.3f} s", loc="right")

    def aggiorna(frame: int):
        k = indice_inizio + frame
        t = tempi[k]

        linea_flusso.set_data(dati["x_nodi_mm"], dati["flussi"][k] / 1e6)

        x_sorgenti_mm = (
            dati["x_riferimenti"][k] * 1000.0
            + np.arange(numero_sorgenti) * distanza_mm
        )
        visibili = (x_sorgenti_mm >= x_vista_mm[0]) & (x_sorgenti_mm <= x_vista_mm[1])
        marker_sorgenti.set_data(
            x_sorgenti_mm[visibili],
            np.full(int(visibili.sum()), q_max_MW * 1.12),
        )

        mappa.set_array(chiudi_campo(dati["campi"][k]))

        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([t, t])
        testo_tempo.set_text(f"t = {t:.3f} s")

        return (
            linea_flusso, marker_sorgenti, mappa,
            linea_vera, linea_letta, cursore, testo_tempo,
        )

    animazione = FuncAnimation(
        fig, aggiorna, frames=n_frame, interval=INTERVALLO_RIPRODUZIONE_MS, blit=False
    )

    # Se il backend non è interattivo si salva una GIF invece di mostrare la finestra.
    if matplotlib.get_backend().lower() == "agg":
        percorso = Path("dataset") / "animazione_superficie.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()