# -*- coding: utf-8 -*-
"""Glaser-Armero S-shaped beam finite-rotation objectivity benchmark."""

import json
import math
import sys
from pathlib import Path

import numpy as np

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

from pytorsocae import TorsoCAESession


LENGTH = 1.0
HALF_HEIGHT = 0.1
WIDTH = 1.0
END_SHIFT = 2.0 * HALF_HEIGHT
SHEAR_MODULUS = 100.0
BULK_MODULUS = 116.66666666666667
YOUNGS_MODULUS = 9.0 * BULK_MODULUS * SHEAR_MODULUS / (
    3.0 * BULK_MODULUS + SHEAR_MODULUS
)
POISSON_RATIO = (3.0 * BULK_MODULUS - 2.0 * SHEAR_MODULUS) / (
    2.0 * (3.0 * BULK_MODULUS + SHEAR_MODULUS)
)
ANGLES_DEG = tuple(range(0, 91, 15))
MESH_DIVISIONS = (6, 1, 1)
FORCE_DRIFT_LIMIT = 1.0e-8
STRESS_DRIFT_LIMIT = 1.0e-8


def _prescribed_end_values(angle_deg: float):
    angle = math.radians(angle_deg)
    theta = f"({angle:.17g})*maximum(0.0, 2.0*t - 1.0)"
    shift = f"({END_SHIFT:.17g})*minimum(1.0, 2.0*t)"
    cosine = f"cos({theta})"
    sine = f"sin({theta})"
    left = [
        f"({cosine})*x - ({sine})*z - x",
        "0.0",
        f"({sine})*x + ({cosine})*z - z",
    ]
    right = [
        f"({cosine})*x - ({sine})*(z + ({shift})) - x",
        "0.0",
        f"({sine})*x + ({cosine})*(z + ({shift})) - z",
    ]
    return left, right


def _local_rotation(angle_deg: float) -> np.ndarray:
    theta = math.radians(angle_deg)
    cosine = math.cos(theta)
    sine = math.sin(theta)
    return np.array(
        [[cosine, 0.0, -sine], [0.0, 1.0, 0.0], [sine, 0.0, cosine]],
        dtype=float,
    )


def run_benchmark():
    session = TorsoCAESession()
    session.build_structured_mesh(
        mesh_id="mesh_0",
        name="Glaser-Armero S-Beam",
        body_id="body_0",
        bounds=(
            (0.0, LENGTH),
            (-0.5 * WIDTH, 0.5 * WIDTH),
            (-HALF_HEIGHT, HALF_HEIGHT),
        ),
        divisions=MESH_DIVISIONS,
        surface_tags={
            "xmin": 1,
            "xmax": 2,
            "ymin": 3,
            "ymax": 4,
            "zmin": 5,
            "zmax": 6,
        },
    )
    session.solid("Glaser-Armero S-Beam").material(
        E=YOUNGS_MODULUS, nu=POISSON_RATIO
    )
    session.surface(3, scope_id="mesh_0").bc(
        "fixed", frame="global", components=["y"]
    )
    session.surface(4, scope_id="mesh_0").bc(
        "fixed", frame="global", components=["y"]
    )
    session.set_physics(
        "structural", submodel="geometric_nonlinear", backend="dolfinx"
    )
    session.set_solver_options(
        algo="direct",
        precond="none",
        tol=1.0e-10,
        max_iter=500,
        num_steps=24,
        max_inner_iter=30,
        inner_tol=1.0e-9,
        max_load_bisections=8,
        n_cores=2,
        device="cpu",
    )

    rows = []
    final_state = {}
    for angle_deg in ANGLES_DEG:
        left, right = _prescribed_end_values(angle_deg)
        session.surface(1, scope_id="mesh_0").bc("fixed", values=left)
        session.surface(2, scope_id="mesh_0").bc("fixed", values=right)
        result = session.compute(mesh_ids=["mesh_0"])
        with np.load(result["npz_path"]) as data:
            reaction_index = int(np.flatnonzero(data["reaction_tags"] == 1)[0])
            global_force = np.asarray(
                data["reaction_forces"][reaction_index], dtype=float
            )
            if angle_deg == ANGLES_DEG[-1]:
                final_state = {
                    "coordinates": np.asarray(data["coordinates"], dtype=float).copy(),
                    "displacement": np.asarray(data["displacement"], dtype=float).copy(),
                }
        local_force = _local_rotation(angle_deg).T @ global_force
        rows.append(
            {
                "angle_deg": float(angle_deg),
                "global_reaction": global_force.tolist(),
                "local_reaction": local_force.tolist(),
                "reaction_magnitude": float(np.linalg.norm(global_force)),
                "max_von_mises": float(result["max_von_mises"]),
            }
        )

    reference_force = np.asarray(rows[0]["local_reaction"], dtype=float)
    reference_norm = max(float(np.linalg.norm(reference_force)), np.finfo(float).eps)
    reference_stress = max(abs(rows[0]["max_von_mises"]), np.finfo(float).eps)
    for row in rows:
        local_force = np.asarray(row["local_reaction"], dtype=float)
        row["local_force_drift"] = float(
            np.linalg.norm(local_force - reference_force) / reference_norm
        )
        row["stress_drift"] = float(
            abs(row["max_von_mises"] - rows[0]["max_von_mises"])
            / reference_stress
        )

    max_force_drift = max(row["local_force_drift"] for row in rows)
    max_stress_drift = max(row["stress_drift"] for row in rows)
    summary = {
        "case": "Glaser-Armero S-shaped beam finite-rotation objectivity",
        "citation": (
            "Glaser and Armero, Engineering Computations 14(7), 759-791 "
            "(1997), DOI 10.1108/02644409710188664"
        ),
        "scope": (
            "Rigid-rotation invariance of co-rotated reaction force and Cauchy "
            "stress for the TorsoCAE total-Lagrangian St. Venant-Kirchhoff model"
        ),
        "geometry": {
            "length": LENGTH,
            "height": 2.0 * HALF_HEIGHT,
            "width": WIDTH,
            "right_end_shift": END_SHIFT,
        },
        "material": {
            "shear_modulus": SHEAR_MODULUS,
            "bulk_modulus": BULK_MODULUS,
            "youngs_modulus": YOUNGS_MODULUS,
            "poisson_ratio": POISSON_RATIO,
        },
        "mesh": {
            "element": "Hex8",
            "divisions": list(MESH_DIVISIONS),
            "nodes": 28,
            "elements": 6,
        },
        "angles_deg": list(ANGLES_DEG),
        "results": rows,
        "max_local_force_drift": max_force_drift,
        "max_stress_drift": max_stress_drift,
        "force_drift_limit": FORCE_DRIFT_LIMIT,
        "stress_drift_limit": STRESS_DRIFT_LIMIT,
        "pass": (
            max_force_drift <= FORCE_DRIFT_LIMIT
            and max_stress_drift <= STRESS_DRIFT_LIMIT
        ),
    }
    return summary, final_state


if __name__ == "__main__":
    benchmark, _final_state = run_benchmark()
    print(json.dumps(benchmark, indent=2), flush=True)