# Generazione della mesh a elementi shell quadrilateri sulla superficie
# cilindrica della fascetta.
#
# La superficie è parametrizzata da (x, theta): x è la coordinata assiale,
# theta quella circonferenziale. Al posto di theta si usa quasi ovunque
# l'ascissa curvilinea arco = raggio_medio * theta, così la shell diventa un
# rettangolo lunghezza × circonferenza e gli elementi sono tutti identici.
#
# I nodi formano una griglia (n_elementi_x + 1) × n_elementi_theta: in
# direzione circonferenziale non c'è un nodo finale distinto perché theta = 0
# e theta = 2*pi sono lo stesso nodo. La periodicità è quindi strutturale,
# nasce dalla connettività e non da vincoli imposti a posteriori.
#
# Ordine dei nodi locali di ogni elemento (antiorario nel piano x-arco):
#
#     4 ---- 3      arco
#     |      |       ^
#     |      |       |
#     1 ---- 2       +---> x

import numpy as np


def costruisci_mesh_cilindrica(
    lunghezza_m: float,
    raggio_medio_m: float,
    n_elementi_x: int,
    n_elementi_theta: int,
) -> dict:
    # Costruisce nodi e connettività della shell cilindrica.
    n_nodi_x = n_elementi_x + 1
    n_nodi = n_nodi_x * n_elementi_theta
    n_elementi = n_elementi_x * n_elementi_theta

    lato_x_m = lunghezza_m / n_elementi_x
    circonferenza_m = 2.0 * np.pi * raggio_medio_m
    lato_arco_m = circonferenza_m / n_elementi_theta

    x_nodi_m = np.arange(n_nodi_x) * lato_x_m
    theta_nodi_rad = np.arange(n_elementi_theta) * (2.0 * np.pi / n_elementi_theta)
    arco_nodi_m = theta_nodi_rad * raggio_medio_m

    # Connettività: elemento (i, j) collega i nodi (i, j), (i+1, j),
    # (i+1, j+1), (i, j+1), con j+1 riavvolto sulla circonferenza.
    i_elem = np.repeat(np.arange(n_elementi_x), n_elementi_theta)
    j_elem = np.tile(np.arange(n_elementi_theta), n_elementi_x)
    j_succ = (j_elem + 1) % n_elementi_theta

    def indice(i, j):
        return i * n_elementi_theta + j

    connettivita = np.stack(
        [
            indice(i_elem, j_elem),
            indice(i_elem + 1, j_elem),
            indice(i_elem + 1, j_succ),
            indice(i_elem, j_succ),
        ],
        axis=1,
    )

    # Segmenti dei due bordi assiali (x = 0 e x = lunghezza): anelli chiusi di
    # n_elementi_theta segmenti, usati per la convezione sullo spessore.
    j = np.arange(n_elementi_theta)
    j_dopo = (j + 1) % n_elementi_theta
    segmenti_bordo = np.concatenate(
        [
            np.stack([indice(0, j), indice(0, j_dopo)], axis=1),
            np.stack(
                [indice(n_elementi_x, j), indice(n_elementi_x, j_dopo)], axis=1
            ),
        ]
    )

    return {
        "n_elementi_x": n_elementi_x,
        "n_elementi_theta": n_elementi_theta,
        "n_elementi": n_elementi,
        "n_nodi_x": n_nodi_x,
        "n_nodi": n_nodi,
        "lunghezza_m": lunghezza_m,
        "raggio_medio_m": raggio_medio_m,
        "circonferenza_m": circonferenza_m,
        "lato_x_m": lato_x_m,
        "lato_arco_m": lato_arco_m,
        "x_nodi_m": x_nodi_m,
        "theta_nodi_rad": theta_nodi_rad,
        "arco_nodi_m": arco_nodi_m,
        "connettivita": connettivita,
        "segmenti_bordo": segmenti_bordo,
        "area_totale_m2": lunghezza_m * circonferenza_m,
    }


def campo_su_griglia(mesh: dict, T: np.ndarray) -> np.ndarray:
    # Rimappa il vettore nodale sulla griglia (n_nodi_x, n_elementi_theta),
    # comoda per le mappe sviluppate e per le superfici 3D.
    return T.reshape(mesh["n_nodi_x"], mesh["n_elementi_theta"])


def coordinate_3d(mesh: dict, chiudi_circonferenza: bool = True):
    # Coordinate cartesiane dei nodi per la vista 3D: x lungo l'asse del
    # cilindro, y e z sulla sezione circolare.
    #
    # Con chiudi_circonferenza si ripete la prima colonna in coda, così la
    # superficie disegnata non mostra una fessura in theta = 0. È solo una
    # necessità di disegno: il nodo ripetuto non è un grado di libertà.
    theta = mesh["theta_nodi_rad"]
    if chiudi_circonferenza:
        theta = np.append(theta, 2.0 * np.pi)
    X, Theta = np.meshgrid(mesh["x_nodi_m"], theta, indexing="ij")
    R = mesh["raggio_medio_m"]
    return X, R * np.sin(Theta), R * np.cos(Theta)


def chiudi_campo(campo: np.ndarray) -> np.ndarray:
    # Ripete la prima colonna circonferenziale in coda, in accordo con
    # coordinate_3d(chiudi_circonferenza=True).
    return np.concatenate([campo, campo[:, :1]], axis=1)