# Elemento shell termico quadrilatero a 4 nodi (Q4) con una temperatura per
# nodo, cioè temperatura uniforme nello spessore.
#
# L'elemento vive nel piano (x, arco) della superficie sviluppata. Poiché la
# mesh cilindrica è strutturata e uniforme, tutti gli elementi sono lo stesso
# rettangolo lato_x × lato_arco: le matrici elementari si calcolano una volta
# sola e si replicano in fase di assemblaggio.
#
# Le funzioni di forma bilineari sono
#
#     N_a(xi, eta) = (1 + xi_a * xi) * (1 + eta_a * eta) / 4
#
# con (xi_a, eta_a) i vertici dell'elemento di riferimento. La mappa verso
# l'elemento fisico è affine, quindi lo jacobiano è costante:
# det(J) = lato_x * lato_arco / 4.
#
# Integrazione con quadratura di Gauss 2x2, esatta per i polinomi che
# compaiono nelle matrici di massa e conduzione dell'elemento Q4.

import numpy as np


# Vertici dell'elemento di riferimento, nello stesso ordine dei nodi locali
# usati dalla connettività in mesh.py.
NODI_RIFERIMENTO = np.array([(-1.0, -1.0), (1.0, -1.0), (1.0, 1.0), (-1.0, 1.0)])

# Quadratura di Gauss a 2 punti per direzione.
PUNTI_GAUSS_1D = np.array([-1.0 / np.sqrt(3.0), 1.0 / np.sqrt(3.0)])
PESI_GAUSS_1D = np.array([1.0, 1.0])


def funzioni_forma(xi: float, eta: float) -> np.ndarray:
    xi_a, eta_a = NODI_RIFERIMENTO[:, 0], NODI_RIFERIMENTO[:, 1]
    return 0.25 * (1.0 + xi_a * xi) * (1.0 + eta_a * eta)


def derivate_funzioni_forma(xi: float, eta: float) -> tuple[np.ndarray, np.ndarray]:
    xi_a, eta_a = NODI_RIFERIMENTO[:, 0], NODI_RIFERIMENTO[:, 1]
    dN_dxi = 0.25 * xi_a * (1.0 + eta_a * eta)
    dN_deta = 0.25 * eta_a * (1.0 + xi_a * xi)
    return dN_dxi, dN_deta


def _punti_quadratura_2d():
    # Punti (xi, eta) e pesi della quadratura 2x2, nell'ordine
    # g = p * 2 + q con xi = punto p (direzione x) ed eta = punto q
    # (direzione circonferenziale). Lo stesso ordine è atteso dai valori del
    # flusso passati a vettore_sorgente_elementi().
    for p, xi in enumerate(PUNTI_GAUSS_1D):
        for q, eta in enumerate(PUNTI_GAUSS_1D):
            yield xi, eta, PESI_GAUSS_1D[p] * PESI_GAUSS_1D[q]


def matrice_forme_ai_punti_gauss() -> np.ndarray:
    # Matrice (4 punti di Gauss, 4 nodi) con le funzioni di forma valutate
    # nei punti di quadratura.
    return np.array([funzioni_forma(xi, eta) for xi, eta, _ in _punti_quadratura_2d()])


def pesi_ai_punti_gauss() -> np.ndarray:
    # Pesi di quadratura nei 4 punti, nello stesso ordine.
    return np.array([peso for _, _, peso in _punti_quadratura_2d()])


def coordinate_locali_punti_gauss(lato_x_m: float, lato_arco_m: float):
    # Offset dei punti di Gauss rispetto al vertice di riferimento
    # dell'elemento, separati per direzione: 2 valori lungo x e 2 lungo
    # l'arco. La separazione è possibile perché l'elemento è un rettangolo.
    offset_x = lato_x_m * (1.0 + PUNTI_GAUSS_1D) / 2.0
    offset_arco = lato_arco_m * (1.0 + PUNTI_GAUSS_1D) / 2.0
    return offset_x, offset_arco


def matrice_massa_elemento(lato_x_m: float, lato_arco_m: float) -> np.ndarray:
    # Integrale di N^T N sull'area dell'elemento. Matrice consistente (non
    # concentrata): è la base sia della capacità termica sia della convezione
    # sulle facce.
    det_j = lato_x_m * lato_arco_m / 4.0
    M = np.zeros((4, 4))
    for xi, eta, peso in _punti_quadratura_2d():
        N = funzioni_forma(xi, eta)
        M += peso * det_j * np.outer(N, N)
    return M


def matrice_conduzione_elemento(
    lato_x_m: float, lato_arco_m: float, conducibilita_W_mK: float, spessore_m: float
) -> np.ndarray:
    # Integrale di B^T D B * spessore sull'area, con D = k * identità
    # (materiale isotropo). Rappresenta la conduzione tangenziale, cioè sia
    # quella assiale sia quella circonferenziale: sul cilindro sviluppato le
    # due direzioni sono ortogonali e la derivata circonferenziale rispetto
    # all'arco equivale a (1/R) d/dtheta.
    det_j = lato_x_m * lato_arco_m / 4.0
    K = np.zeros((4, 4))
    for xi, eta, peso in _punti_quadratura_2d():
        dN_dxi, dN_deta = derivate_funzioni_forma(xi, eta)
        # Jacobiano diagonale e costante: la derivata fisica è quella
        # naturale riscalata dal semilato dell'elemento.
        B = np.stack([dN_dxi * (2.0 / lato_x_m), dN_deta * (2.0 / lato_arco_m)])
        K += peso * det_j * conducibilita_W_mK * spessore_m * (B.T @ B)
    return K


def vettore_carico_uniforme_elemento(lato_x_m: float, lato_arco_m: float) -> np.ndarray:
    # Integrale di N^T sull'area: distribuisce ai nodi un carico superficiale
    # uniforme unitario.
    det_j = lato_x_m * lato_arco_m / 4.0
    f = np.zeros(4)
    for xi, eta, peso in _punti_quadratura_2d():
        f += peso * det_j * funzioni_forma(xi, eta)
    return f


def _funzioni_forma_segmento(xi: float) -> np.ndarray:
    return np.array([0.5 * (1.0 - xi), 0.5 * (1.0 + xi)])


def matrice_massa_segmento(lunghezza_m: float) -> np.ndarray:
    # Integrale di N^T N su un segmento a 2 nodi, per i bordi assiali.
    det_j = lunghezza_m / 2.0
    M = np.zeros((2, 2))
    for xi, peso in zip(PUNTI_GAUSS_1D, PESI_GAUSS_1D):
        N = _funzioni_forma_segmento(xi)
        M += peso * det_j * np.outer(N, N)
    return M


def vettore_carico_segmento(lunghezza_m: float) -> np.ndarray:
    # Integrale di N^T su un segmento a 2 nodi.
    det_j = lunghezza_m / 2.0
    f = np.zeros(2)
    for xi, peso in zip(PUNTI_GAUSS_1D, PESI_GAUSS_1D):
        f += peso * det_j * _funzioni_forma_segmento(xi)
    return f