# Solutore termico transitorio a elementi finiti sulla mesh shell della
# fascetta (vedi mesh.py per la geometria).
#
# Differenza rispetto al modello ai volumi finiti di simulate.py:
# - lì il dominio è la sezione x-z (spessore risolto, circonferenza y
#   collassata in un'attenuazione gaussiana del flusso e in un sink lineare
#   di tipo aletta);
# - qui il dominio è la superficie media del cilindro, con x e la coordinata
#   circonferenziale s = R·theta entrambe risolte, mentre lo spessore è
#   collassato (temperatura uniforme attraverso la parete).
#
# Il collasso in spessore è lecito: con spessore 0.18 mm il numero di Biot
# h·t/k vale ~1e-7 e il tempo di diffusione attraverso la parete t²/alpha è
# di pochi millisecondi, molto più rapido del transito delle sorgenti. Di
# conseguenza la skin depth non entra nel modello shell: conta solo il flusso
# totale assorbito per unità di superficie.
#
# Equazione risolta (per unità di superficie media):
#
#   rho·cp·t ∂T/∂t = k·t ∇²T + q(x, s, t) - (h_est + h_int)·(T - T_amb)
#
# con convezione aggiuntiva sui due bordi anulari x = 0 e x = L (coefficiente
# h_bordi su un'area pari a spessore × perimetro). L'impronta della sorgente
# è una gaussiana isotropa nel piano (x, s), con la distanza circonferenziale
# valutata sull'immagine più vicina perché la superficie è chiusa.
#
# Discretizzazione: elementi shell quadrangolari bilineari a 4 nodi. Gli
# elementi della mesh sono tutti rettangoli identici dx × ds, quindi le
# matrici di elemento sono calcolate una volta sola e assemblate in forma
# vettorizzata. Integrazione temporale con Eulero implicito: la matrice di
# sistema è costante e viene fattorizzata LU una volta per run.

import math

import numpy as np
import scipy.sparse as sp
from scipy.sparse.linalg import splu

from config import FEM
from materials import MATERIALI
from mesh import genera_mesh
# Helper della cinematica delle sorgenti condivisi con il modello ai volumi
# finiti: la corsa del gruppo di induttori deve essere identica nei due modelli.
from simulate import (
    _intervallo_attivo,
    _x_riferimento_finale_m,
    _x_riferimento_iniziale_m,
)

# Matrici di riferimento dell'elemento rettangolare bilineare a 4 nodi, con
# nodi locali in ordine antiorario: (0,0), (a,0), (a,b), (0,b).
#
# MASSA_RIF va moltiplicata per l'area a·b e dà l'integrale di N_i·N_j;
# RIGIDEZZA_X per b/a e RIGIDEZZA_S per a/b, e la loro somma dà l'integrale
# di grad(N_i)·grad(N_j).
MASSA_RIF = np.array(
    [[4.0, 2.0, 1.0, 2.0],
     [2.0, 4.0, 2.0, 1.0],
     [1.0, 2.0, 4.0, 2.0],
     [2.0, 1.0, 2.0, 4.0]]
) / 36.0

RIGIDEZZA_X = np.array(
    [[2.0, -2.0, -1.0, 1.0],
     [-2.0, 2.0, 1.0, -1.0],
     [-1.0, 1.0, 2.0, -2.0],
     [1.0, -1.0, -2.0, 2.0]]
) / 6.0

RIGIDEZZA_S = np.array(
    [[2.0, 1.0, -1.0, -2.0],
     [1.0, 2.0, -2.0, -1.0],
     [-1.0, -2.0, 2.0, 1.0],
     [-2.0, -1.0, 1.0, 2.0]]
) / 6.0


def _assembla(elementi: np.ndarray, matrice_elemento: np.ndarray, n_nodi: int):
    # Assembla una matrice globale sparsa a partire da un'unica matrice di
    # elemento 4×4, uguale per tutti gli elementi. I contributi ripetuti sullo
    # stesso (riga, colonna) vengono sommati dal formato COO.
    n_elementi = len(elementi)
    righe = np.broadcast_to(elementi[:, :, None], (n_elementi, 4, 4)).ravel()
    colonne = np.broadcast_to(elementi[:, None, :], (n_elementi, 4, 4)).ravel()
    valori = np.broadcast_to(matrice_elemento, (n_elementi, 4, 4)).ravel()
    return sp.coo_matrix(
        (valori, (righe, colonne)), shape=(n_nodi, n_nodi)
    ).tocsr()


def _massa_anello(indici: np.ndarray, lunghezza_segmento_m: float, n_nodi: int):
    # Matrice di massa 1D su un anello chiuso di nodi (integrale di N_i·N_j
    # lungo la circonferenza), usata per la convezione sui bordi x = 0 e x = L.
    n = len(indici)
    successivi = indici[(np.arange(n) + 1) % n]
    coefficiente = lunghezza_segmento_m / 6.0
    righe = np.concatenate([indici, indici, successivi, successivi])
    colonne = np.concatenate([indici, successivi, indici, successivi])
    valori = coefficiente * np.concatenate(
        [np.full(n, 2.0), np.full(n, 1.0), np.full(n, 1.0), np.full(n, 2.0)]
    )
    return sp.coo_matrix(
        (valori, (righe, colonne)), shape=(n_nodi, n_nodi)
    ).tocsr()


def prepara_stato_fem(
    fascetta: dict,
    aria: dict,
    sorgente: dict,
    sensore: dict,
    mesh_dati: dict | None = None,
    dt_s: float | None = None,
) -> dict:
    """Assembla le matrici FEM e fattorizza il sistema implicito.

    Restituisce lo stato costante per l'integrazione temporale: matrici,
    solutore LU, coordinate nodali, peso circonferenziale della sorgente e
    indice del nodo osservato dal sensore.
    """
    if mesh_dati is None:
        mesh_dati = genera_mesh(fascetta)
    if dt_s is None:
        dt_s = FEM["dt_s"]

    materiale = MATERIALI[fascetta["materiale"]]
    k = materiale["conducibilita_termica_W_mK"]
    rho = materiale["densita_kg_m3"]
    cp = materiale["calore_specifico_J_kgK"]

    spessore_m = mesh_dati["spessore_m"]
    raggio_m = mesh_dati["raggio_m"]
    circonferenza_m = 2.0 * math.pi * raggio_m
    n_x = mesh_dati["n_elementi_x"]
    n_theta = mesh_dati["n_elementi_circonferenza"]
    elementi = mesh_dati["elementi"]
    indice_nodo = mesh_dati["indice_nodo"]
    n_nodi = len(mesh_dati["nodi"])

    dx_m = mesh_dati["lunghezza_m"] / n_x
    ds_m = circonferenza_m / n_theta

    h_esterno = aria["h_esterno_W_m2K"]
    h_interno = aria["h_interno_W_m2K"]
    h_bordi = aria["h_bordi_W_m2K"]
    T_ambiente = aria["temperatura_ambiente_C"]

    # Matrice di massa di superficie: riusata per la capacità termica, per la
    # convezione sulle due facce e per il carico della sorgente.
    M_area = _assembla(elementi, MASSA_RIF * (dx_m * ds_m), n_nodi)
    K_cond = _assembla(
        elementi,
        k * spessore_m * (RIGIDEZZA_X * (ds_m / dx_m) + RIGIDEZZA_S * (dx_m / ds_m)),
        n_nodi,
    )

    # Convezione sui bordi anulari: l'area di scambio di ogni segmento è
    # spessore × lunghezza del segmento circonferenziale.
    H_bordi = h_bordi * spessore_m * (
        _massa_anello(indice_nodo[0, :], ds_m, n_nodi)
        + _massa_anello(indice_nodo[-1, :], ds_m, n_nodi)
    )

    C_su_dt = (rho * cp * spessore_m / dt_s) * M_area
    A = C_su_dt + K_cond + (h_esterno + h_interno) * M_area + H_bordi

    # Termine noto costante della convezione: (somma delle righe) × h × T_amb,
    # perché la somma delle righe della matrice di massa è l'integrale di N_i.
    area_nodale = np.asarray(M_area.sum(axis=1)).ravel()
    bordo_nodale = np.asarray(H_bordi.sum(axis=1)).ravel()
    carico_convezione = (h_esterno + h_interno) * area_nodale * T_ambiente
    carico_convezione += bordo_nodale * T_ambiente

    # Peso circonferenziale dell'impronta della sorgente: dipende solo
    # dall'offset y del percorso e dalla sigma dello spot, entrambi costanti
    # durante il run, quindi è valutato una volta sola qui. La distanza è
    # quella dell'immagine più vicina, perché la superficie è chiusa.
    sigma_m = sorgente["sigma_punto_m"]
    ds_sorgente = raggio_m * mesh_dati["theta_nodi_rad"] - sorgente["offset_y_percorso_m"]
    ds_sorgente -= circonferenza_m * np.round(ds_sorgente / circonferenza_m)
    peso_circonferenziale = np.exp(-0.5 * (ds_sorgente ** 2) / (sigma_m * sigma_m))

    # Nodo osservato dal sensore: x più vicino a x_mm sul piano theta = 0.
    x_sensore_m = sensore["x_mm"] / 1000.0
    i_sensore = int(np.argmin(np.abs(mesh_dati["x_nodi_m"] - x_sensore_m)))
    indice_sensore = int(indice_nodo[i_sensore, 0])

    return {
        "mesh": mesh_dati,
        "dt_s": dt_s,
        "M_area": M_area,
        "C_su_dt": C_su_dt,
        "carico_convezione": carico_convezione,
        "solutore": splu(sp.csc_matrix(A)),
        "x_nodi_m": mesh_dati["x_nodi_m"],
        "peso_circonferenziale": peso_circonferenziale,
        "forma_griglia": (n_x + 1, n_theta),
        "circonferenza_m": circonferenza_m,
        "area_nodale_m2": area_nodale,
        "T_ambiente_C": T_ambiente,
        "x_sensore_m": x_sensore_m,
        "indice_sensore": indice_sensore,
        "n_nodi": n_nodi,
    }


def flusso_nodale_W_m2(
    sorgente: dict,
    stato: dict,
    t_s: float,
) -> tuple[float, np.ndarray]:
    """Flusso termico assorbito nei nodi all'istante t_s.

    Restituisce (x_riferimento_m, q_nodale) dove x_riferimento_m è la
    posizione della sorgente di indice 0. Ogni sorgente attiva contribuisce
    con un'impronta gaussiana isotropa in (x, s).

    L'impronta gaussiana isotropa è separabile, q(x, s) = q_x(x) · q_s(s), e
    il fattore circonferenziale q_s è costante nel tempo (precalcolato in
    `prepara_stato_fem`). Qui si valuta quindi solo il fattore assiale sui
    n_x + 1 nodi distinti in x e si espande con un prodotto esterno, invece di
    valutare l'esponenziale su tutti i nodi della mesh.
    """
    x_sensore_m = stato["x_sensore_m"]
    x_rif_iniziale = _x_riferimento_iniziale_m(sorgente, x_sensore_m)
    x_rif_finale = _x_riferimento_finale_m(sorgente, x_sensore_m)
    x_riferimento = x_rif_iniziale + sorgente["velocita_m_s"] * t_s

    numero_sorgenti = sorgente.get("numero_sorgenti", 1)
    distanza = sorgente.get("distanza_sorgenti_m", 0.0)
    v = sorgente["velocita_m_s"]
    zero_dopo_fine = sorgente.get("zero_dopo_fine", True)
    sigma = sorgente["sigma_punto_m"]
    q_picco = sorgente["flusso_termico_picco_W_m2"] * sorgente["efficienza_riscaldamento"]

    peso_assiale = np.zeros(len(stato["x_nodi_m"]))
    for i in range(numero_sorgenti):
        x_i = x_riferimento + i * distanza

        if zero_dopo_fine:
            inizio_i = x_rif_iniziale + i * distanza
            fine_i = x_rif_finale + i * distanza
            if not _intervallo_attivo(inizio_i, fine_i, v, x_i):
                continue

        dx = stato["x_nodi_m"] - x_i
        peso_assiale += np.exp(-0.5 * (dx * dx) / (sigma * sigma))

    q_nodale = q_picco * np.outer(peso_assiale, stato["peso_circonferenziale"])
    return x_riferimento, q_nodale.ravel()


def passo_implicito_fem(stato: dict, T: np.ndarray, q_nodale: np.ndarray) -> np.ndarray:
    # Avanza il campo nodale di un passo dt: termine noto = capacità sul campo
    # precedente + carico della sorgente + carico di convezione, poi risoluzione
    # del sistema già fattorizzato.
    rhs = stato["C_su_dt"] @ T
    rhs += stato["M_area"] @ q_nodale
    rhs += stato["carico_convezione"]
    return stato["solutore"].solve(rhs)


def campo_iniziale(stato: dict) -> np.ndarray:
    return np.full(stato["n_nodi"], stato["T_ambiente_C"], dtype=float)


def simula_campo_fem(
    cfg_run: dict,
    durata_s: float,
    dt_frame_s: float,
    mesh_dati: dict | None = None,
) -> dict:
    """Integra il campo FEM fino a durata_s, salvando un campo ogni dt_frame_s."""
    fascetta = cfg_run["fascetta"]
    aria = cfg_run["aria"]
    sorgente = cfg_run["sorgente"]
    sensore = cfg_run["sensore"]

    stato = prepara_stato_fem(fascetta, aria, sorgente, sensore, mesh_dati=mesh_dati)
    dt = stato["dt_s"]
    indice_sensore = stato["indice_sensore"]

    T = campo_iniziale(stato)
    T_sensore = float(T[indice_sensore])
    tau_sensore = max(sensore["costante_tempo_s"], 1e-9)

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

    prossimo_frame_t = 0.0
    t = 0.0
    while t <= durata_s + 1e-12:
        x_rif, q_nodale = flusso_nodale_W_m2(sorgente, stato, t)
        T = passo_implicito_fem(stato, T, q_nodale)

        T_sensore += (T[indice_sensore] - T_sensore) * dt / tau_sensore

        if t + 1e-12 >= prossimo_frame_t:
            tempi.append(t)
            campi.append(T.copy())
            x_riferimenti.append(x_rif)
            T_vere.append(float(T[indice_sensore]))
            T_lette.append(T_sensore)
            prossimo_frame_t += dt_frame_s

        t += dt

    return {
        "stato": stato,
        "mesh": stato["mesh"],
        "tempi": np.array(tempi),
        "campi": campi,
        "x_riferimenti": np.array(x_riferimenti),
        "T_vere": np.array(T_vere),
        "T_lette": np.array(T_lette),
        "sorgente": sorgente,
        "T_ambiente_C": stato["T_ambiente_C"],
    }