# Assemblaggio delle matrici globali della shell termica.
#
# Tutti gli elementi della mesh cilindrica sono lo stesso rettangolo, quindi
# ogni matrice elementare si calcola una volta e si replica su tutti gli
# elementi. L'assemblaggio è la somma dei contributi sui gradi di libertà
# condivisi: due elementi confinanti condividono i due nodi del bordo comune,
# ed è questo che genera la conduzione tra elementi adiacenti. Non serve
# nessun accoppiamento aggiuntivo tra elementi dello stesso materiale.

import numpy as np
import scipy.sparse as sp

import elementi_shell as el


def _assembla_matrice(matrice_elemento: np.ndarray, connettivita: np.ndarray, n_nodi: int):
    # Replica una matrice elementare costante su tutti gli elementi e somma i
    # contributi sui gradi di libertà condivisi.
    n_elementi, n_locali = connettivita.shape
    righe = np.repeat(connettivita, n_locali, axis=1).ravel()
    colonne = np.tile(connettivita, (1, n_locali)).ravel()
    dati = np.tile(matrice_elemento.ravel(), n_elementi)
    return sp.coo_matrix(
        (dati, (righe, colonne)), shape=(n_nodi, n_nodi)
    ).tocsr()


def _assembla_vettore(vettore_elemento: np.ndarray, connettivita: np.ndarray, n_nodi: int):
    dati = np.tile(vettore_elemento, connettivita.shape[0])
    return np.bincount(connettivita.ravel(), weights=dati, minlength=n_nodi)


def assembla_capacita(mesh: dict, materiale: dict, spessore_m: float):
    # C = integrale di rho * cp * spessore * N^T N dA.
    #
    # Matrice consistente: con integrazione temporale implicita conserva
    # meglio l'energia rispetto alla versione concentrata sulla diagonale.
    capacita_superficiale = (
        materiale["densita_kg_m3"] * materiale["calore_specifico_J_kgK"] * spessore_m
    )
    M_e = el.matrice_massa_elemento(mesh["lato_x_m"], mesh["lato_arco_m"])
    return _assembla_matrice(
        capacita_superficiale * M_e, mesh["connettivita"], mesh["n_nodi"]
    )


def assembla_conduzione(mesh: dict, materiale: dict, spessore_m: float):
    # K_cond = integrale di B^T D B * spessore dA, con D isotropa.
    K_e = el.matrice_conduzione_elemento(
        mesh["lato_x_m"],
        mesh["lato_arco_m"],
        materiale["conducibilita_termica_W_mK"],
        spessore_m,
    )
    return _assembla_matrice(K_e, mesh["connettivita"], mesh["n_nodi"])


def assembla_convezione(mesh: dict, aria: dict, spessore_m: float):
    # Convezione sulle facce e sui bordi assiali.
    #
    # Nelle shell la faccia esterna e quella interna sono entrambe superfici
    # fisiche esposte all'aria, quindi ogni elemento scambia su tutta la sua
    # area con il coefficiente combinato h_esterno + h_interno. I bordi
    # assiali (x = 0 e x = lunghezza) espongono invece solo lo spessore della
    # lamiera: il loro contributo è proporzionale a h_bordi * spessore ed è
    # marginale rispetto a quello delle facce.
    T_ambiente = aria["temperatura_ambiente_C"]
    h_facce = aria["h_esterno_W_m2K"] + aria["h_interno_W_m2K"]

    M_e = el.matrice_massa_elemento(mesh["lato_x_m"], mesh["lato_arco_m"])
    f_e = el.vettore_carico_uniforme_elemento(mesh["lato_x_m"], mesh["lato_arco_m"])

    K_conv = _assembla_matrice(h_facce * M_e, mesh["connettivita"], mesh["n_nodi"])
    f_ambiente = _assembla_vettore(
        h_facce * T_ambiente * f_e, mesh["connettivita"], mesh["n_nodi"]
    )

    h_bordi_efficace = aria["h_bordi_W_m2K"] * spessore_m
    M_bordo = el.matrice_massa_segmento(mesh["lato_arco_m"])
    f_bordo = el.vettore_carico_segmento(mesh["lato_arco_m"])

    K_conv = K_conv + _assembla_matrice(
        h_bordi_efficace * M_bordo, mesh["segmenti_bordo"], mesh["n_nodi"]
    )
    f_ambiente = f_ambiente + _assembla_vettore(
        h_bordi_efficace * T_ambiente * f_bordo,
        mesh["segmenti_bordo"],
        mesh["n_nodi"],
    )

    return K_conv, f_ambiente


def prepara_assemblatore_sorgente(mesh: dict) -> dict:
    # Precalcola quanto serve per assemblare il vettore della sorgente a ogni
    # passo temporale: le funzioni di forma nei punti di Gauss pesate per il
    # peso di quadratura e lo jacobiano, e le coordinate dei punti di Gauss
    # separate per direzione.
    det_j = mesh["lato_x_m"] * mesh["lato_arco_m"] / 4.0
    pesi = el.pesi_ai_punti_gauss()
    N_gauss = el.matrice_forme_ai_punti_gauss()
    offset_x, offset_arco = el.coordinate_locali_punti_gauss(
        mesh["lato_x_m"], mesh["lato_arco_m"]
    )

    # Coordinate assolute dei punti di Gauss, per direzione: (n_elementi_x, 2)
    # e (n_elementi_theta, 2).
    x_gauss = mesh["x_nodi_m"][: mesh["n_elementi_x"], None] + offset_x[None, :]
    arco_gauss = mesh["arco_nodi_m"][:, None] + offset_arco[None, :]

    return {
        "peso_forme": (pesi * det_j)[:, None] * N_gauss,
        "connettivita": mesh["connettivita"],
        "n_nodi": mesh["n_nodi"],
        "n_elementi_x": mesh["n_elementi_x"],
        "n_elementi_theta": mesh["n_elementi_theta"],
        "x_gauss_m": x_gauss,
        "arco_gauss_m": arco_gauss,
    }


def assembla_sorgente(
    assemblatore: dict, flusso_x_W_m2: np.ndarray, fattore_arco: np.ndarray
) -> np.ndarray:
    # f_src = integrale di N^T q'' dA, con q'' fattorizzato come
    # q''(x, arco) = flusso_x(x) * fattore_arco(arco).
    #
    # La fattorizzazione non è un'approssimazione: la gaussiana della
    # sorgente è separabile in x e in arco, e tutte le sorgenti del gruppo
    # condividono la stessa posizione circonferenziale, quindi il fattore
    # circonferenziale è comune e costante nel tempo.
    #
    # flusso_x_W_m2 ha forma (n_elementi_x, 2) e fattore_arco
    # (n_elementi_theta, 2): sono i valori nei punti di Gauss di ciascuna
    # direzione.
    #
    # Per gran parte di un run le sorgenti sono fuori dalla fascetta e il
    # fattore assiale è identicamente nullo: in quel caso il vettore è nullo
    # e non serve percorrere la mesh.
    if not flusso_x_W_m2.any():
        return np.zeros(assemblatore["n_nodi"])

    q_gauss = (
        flusso_x_W_m2[:, None, :, None] * fattore_arco[None, :, None, :]
    ).reshape(-1, 4)
    contributi = q_gauss @ assemblatore["peso_forme"]
    return np.bincount(
        assemblatore["connettivita"].ravel(),
        weights=contributi.ravel(),
        minlength=assemblatore["n_nodi"],
    )