"""TorsoCAE journal for the Scordelis-Lo cylindrical roof benchmark."""

from __future__ import annotations

import sys
from pathlib import Path

import numpy as np


ROOT = Path(__file__).resolve().parents[4]
SERVER_DIR = ROOT / "server"
if str(SERVER_DIR) not in sys.path:
    sys.path.insert(0, str(SERVER_DIR))

from pytorsocae import TorsoCAESession


RADIUS = 25.0
FULL_LENGTH = 50.0
HALF_ANGLE_DEG = 40.0
THICKNESS = 0.25
YOUNG = 4.32e8
POISSON = 0.0
DENSITY = 360.0
GRAVITY = 1.0

SHELL_REGION = 1
CROWN_SYMMETRY = 101
SPAN_SYMMETRY = 102
END_DIAPHRAGM = 103
FREE_EDGE = 104
MESH_ID = "scordelis_lo_roof"


def _node(i: int, j: int, n: int) -> int:
    return i * (n + 1) + j


def build_mesh(n: int) -> dict:
    """Build the regular quarter-roof Q4 mesh used by Ko et al. (2017)."""
    if n < 1:
        raise ValueError("Scordelis-Lo mesh resolution must be positive")

    half_span = 0.5 * FULL_LENGTH
    half_angle = np.deg2rad(HALF_ANGLE_DEG)
    points = np.empty(((n + 1) ** 2, 3), dtype=np.float64)
    for i in range(n + 1):
        x = half_span * i / n
        for j in range(n + 1):
            theta = half_angle * j / n
            points[_node(i, j, n)] = (
                x,
                RADIUS * np.sin(theta),
                RADIUS * np.cos(theta),
            )

    quads = np.empty((n * n, 4), dtype=np.int64)
    element = 0
    for i in range(n):
        for j in range(n):
            quads[element] = (
                _node(i, j, n),
                _node(i + 1, j, n),
                _node(i + 1, j + 1, n),
                _node(i, j + 1, n),
            )
            element += 1

    edges = []
    tags = []
    for i in range(n):
        edges.append((_node(i, 0, n), _node(i + 1, 0, n)))
        tags.append(CROWN_SYMMETRY)
        edges.append((_node(i, n, n), _node(i + 1, n, n)))
        tags.append(FREE_EDGE)
    for j in range(n):
        edges.append((_node(0, j, n), _node(0, j + 1, n)))
        tags.append(SPAN_SYMMETRY)
        edges.append((_node(n, j, n), _node(n, j + 1, n)))
        tags.append(END_DIAPHRAGM)

    return {
        "x": points,
        "topologies": {
            3: {
                "topology": quads,
                "cell_data": np.full(len(quads), SHELL_REGION, dtype=np.int32),
            },
            1: {
                "topology": np.asarray(edges, dtype=np.int64),
                "cell_data": np.asarray(tags, dtype=np.int32),
            },
        },
        "type_props": {
            3: {"dim": 2, "num_nodes": 4},
            1: {"dim": 1, "num_nodes": 2},
        },
        "gdim": 3,
    }


def configure_case(
    session: TorsoCAESession,
    *,
    n: int = 16,
    n_cores: int = 2,
) -> None:
    session.inline_mesh(
        build_mesh(n),
        mesh_id=MESH_ID,
        name=f"Scordelis-Lo Roof ({n} x {n})",
    )

    # With the cylinder axis along global x, rotation_1 is longitudinal
    # and rotation_2 is circumferential throughout this mesh.
    session.surface(CROWN_SYMMETRY, scope_id=MESH_ID).bc(
        "fixed",
        components=["y", "rotation_1"],
    )
    session.surface(SPAN_SYMMETRY, scope_id=MESH_ID).bc(
        "fixed",
        components=["x", "rotation_2"],
    )
    session.surface(END_DIAPHRAGM, scope_id=MESH_ID).bc(
        "fixed",
        components=["y", "z", "rotation_1"],
    )
    session.set_model_options(
        shell={
            "thickness": THICKNESS,
            "E": YOUNG,
            "nu": POISSON,
            "rho": DENSITY,
            "shear_correction": 5.0 / 6.0,
        },
        gravity=[0.0, 0.0, -GRAVITY],
    )
    session.set_physics(
        "structural",
        submodel="linear_elastic",
        backend="torso",
    )
    session.set_solver_options(
        algo="direct",
        precond="none",
        tol=1.0e-10,
        max_iter=500,
        n_cores=n_cores,
        device="cpu",
    )


def run_case(*, n: int = 16, n_cores: int = 2) -> dict:
    session = TorsoCAESession()
    configure_case(session, n=n, n_cores=n_cores)
    return session.compute(mesh_ids=[MESH_ID])


if __name__ == "__main__":
    result = run_case()
    print(
        "Scordelis-Lo roof:",
        f"max_displacement={result['max_displacement']:.6e}",
        f"max_von_mises={result['max_von_mises']:.6e}",
    )