"""
TorsoCAE Validation — geometric_nonlinear case 2: Cook's membrane.

Classic tapered, skewed cantilever panel (Cook, 1974) under a distributed
shear load on the short free edge — the standard nonlinear/mixed-FEM mesh-
distortion-sensitivity benchmark. Geometry (in consistent length units):
vertices (0,0), (48,44), (48,60), (0,44), extruded to unit thickness. The
left edge (x=0, y in [0,44]) is clamped; the right edge (x=48, y in
[44,60]) carries a uniform shear traction.

Reference-number caveat (documented up front, not just in the README):
the widely cited "converged tip deflection ~23.9" figure in the FEM
literature is for the *linear elastic*, 2D plane-stress version of this
problem (E=1, nu=1/3, total load=1) — it does not apply as an external
ground truth here for two independent reasons: (1) we solve the
`geometric_nonlinear` (Total-Lagrangian, large-rotation) submodel, and at
the classical unit load level the linear-theory deflection is already
~40-50% of the panel span, i.e. the classical load magnitude sits far
outside the small-strain regime the linear reference assumes; (2) this is
a 3D solid (free z-faces) rather than an idealized 2D plane-stress
element, which shifts the converged value even in the small-deformation
limit. We therefore validate via h-refinement mesh-convergence (primary
criterion, as explicitly permitted by the validation plan when no
authoritative reference exists for the parameterization used) rather than
against the 23.9 number. The load magnitude below is our own consistent,
documented parameterization.

Load-magnitude note (found during validation, not originally anticipated):
this specific geometry has a genuinely sharp acute corner at the clamped
end (~42.5 deg wedge angle), where local strain concentrates far faster
than the panel's overall deflection grows. The compressible Saint
Venant-Kirchhoff model this submodel uses is only valid for large
ROTATIONS with small-to-moderate local STRAIN; above an applied traction
of ~700 Pa on this geometry, local strain at that corner exceeds SVK's
valid range and its stress response becomes non-monotonic, which no amount
of load-step refinement can converge around (verified: bisecting the load
increment down by 1000x still fails at the same absolute traction level,
on both the coarsest and finest mesh -- ruling out mesh distortion/element
inversion as the cause). We therefore load well below that threshold
(500 Pa maximum, ~30% safety margin) -- small enough to stay in SVK's
valid regime, but this necessarily means the nonlinear correction here is
itself small (~1% vs. linear theory at this load) since a *meaningfully*
large nonlinear-differentiating deflection on this specific sharp-cornered
geometry is not reachable within the model's valid range. This case's
purpose is therefore narrowed to mesh-distortion-sensitivity convergence
of the solver (its documented benchmark role) rather than demonstrating a
large geometric-nonlinear effect -- that is already covered, at a
meaningful 44.8-degree rotation matching an exact reference to 0.95%, by
case_01 (elastica). See also: a general load-bisection retry fallback
(matching the existing hyperelastic/elastoplastic drivers) was added to
`_solve_structural_geometric_nonlinear` in dolfinx_backend.py during this
investigation -- a real backend gap (this submodel previously had no retry
at all and crashed outright with `residual=nan` on any Newton failure).
It benefits any geometric_nonlinear case, though it cannot rescue a solve
that is genuinely outside SVK's valid strain range, as above.

Geometry construction note: TorsoCAE's csg_extrude() profile builder fits
a periodic spline through the given points (not a polyline), so sharp
corners get rounded/overshot unless very densely sampled. Instead we build
the exact polygon via box + two EXACT half-plane cuts: take the full
[0,48]x[0,60] rectangle and boolean-subtract two large rotated boxes whose
flat faces are pinned through the box's own corners (0,0)/(0,44) and
(0,44)/(0,60) via csg_transform("rotate", ...) about those points. This is
exact (no spline approximation) and — verified empirically — leaves the
clamped (x=0) and loaded (x=48) faces as separate, stable surface tags
(1 and 6 respectively) across mesh resolutions.

Run: HWLOC_COMPONENTS=-gl python3 validation/structural/geometric_nonlinear/case_02_cooks_membrane/run_validation.py
"""
import json
import math
import sys
import time
from pathlib import Path

import numpy as np

sys.path.insert(0, str(Path(__file__).resolve().parents[4] / "server"))
from pytorsocae import TorsoCAESession

CASE_DIR = Path(__file__).resolve().parent

# ── Geometry (Cook 1974 classic profile, consistent length units) ──────────
LX, H0, H1, T = 48.0, 44.0, 60.0, 1.0
THETA_BOTTOM_DEG = math.degrees(math.atan2(H0, LX))          # ~42.51 deg
THETA_TOP_DEG    = math.degrees(math.atan2(H1 - H0, LX))     # ~18.43 deg
R = 200.0  # oversized cutter box extent, well beyond the panel

# ── Material / load (own parameterization; see module docstring) ──────────
E, NU = 1.0e6, 0.3                     # Pa, -  (soft solid: keeps forces modest)
TRACTION_Y = 5.0e2                     # Pa, uniform shear on the loaded edge (see docstring load-magnitude note)
LOADED_EDGE_AREA = (H1 - H0) * T       # m^2

MESH_RESOLUTIONS = [1.5, 1.0, 0.7, 0.5]   # size_factor, coarse -> fine
NUM_LOAD_STEPS = 20
MAX_LOAD_BISECTIONS = 8

TIP_POINT = np.array([LX, 0.5 * (H0 + H1), 0.5 * T])   # midpoint of loaded edge


def build_geometry(session: TorsoCAESession) -> None:
    session.csg_make_shape("box", {"dx": LX, "dy": H1, "dz": T, "cx": 0, "cy": 0, "cz": 0},
                            name="CookBox", body_id="body_main")
    session.csg_make_shape("box", {"dx": R, "dy": R, "dz": T + 4, "cx": 0, "cy": -R, "cz": -2},
                            name="BottomCutter", body_id="body_bc")
    session.csg_transform("body_bc", "rotate",
                           {"ax": 0, "ay": 0, "az": 1, "angle_deg": THETA_BOTTOM_DEG, "cx": 0, "cy": 0, "cz": 0})
    session.csg_make_shape("box", {"dx": R, "dy": R, "dz": T + 4, "cx": 0, "cy": H0, "cz": -2},
                            name="TopCutter", body_id="body_tc")
    session.csg_transform("body_tc", "rotate",
                           {"ax": 0, "ay": 0, "az": 1, "angle_deg": THETA_TOP_DEG, "cx": 0, "cy": H0, "cz": 0})
    session.csg_boolean("cut", "body_main", "body_bc", name="Cook1", body_id="body_c1")
    session.csg_boolean("cut", "body_c1", "body_tc", name="Cook", body_id="body_cook")
    session.csg_select("body_cook")


def run_case(size_factor: float) -> dict:
    session = TorsoCAESession()
    build_geometry(session)
    session.mesh(mesh_id="mesh_0", algo_id="delaunay", size_factor=size_factor,
                 size_min=0, size_max=0, order=2, dim=3)
    session.solid("Cook").material(E=E, nu=NU)
    session.surface(1, scope_id="body_cook").bc("fixed")
    session.surface(6, scope_id="body_cook").bc("traction", values=[0, TRACTION_Y, 0])
    session.set_physics("structural", submodel="geometric_nonlinear", backend="dolfinx")
    session.set_solver_options(algo="direct", precond="none", tol=1e-08, max_iter=2000,
                                num_steps=NUM_LOAD_STEPS, max_inner_iter=25, inner_tol=1e-07,
                                max_load_bisections=MAX_LOAD_BISECTIONS,
                                n_cores=2, device="cpu")
    t0 = time.time()
    result = session.compute(mesh_ids=["mesh_0"])
    elapsed = time.time() - t0

    npz = np.load(result["npz_path"])
    coords = npz["coordinates"]
    disp = npz["displacement"]
    tip_idx = np.argmin(np.sum((coords - TIP_POINT) ** 2, axis=1))
    u_tip = disp[tip_idx]
    n_dofs = int(coords.shape[0])
    return {
        "size_factor": size_factor,
        "n_dofs": n_dofs,
        "max_displacement_m": result.get("max_displacement"),
        "max_von_mises_pa": result.get("max_von_mises"),
        "tip_node_ref_xyz": coords[tip_idx].tolist(),
        "tip_disp_xyz": u_tip.tolist(),
        "tip_dy": float(u_tip[1]),
        "elapsed_s": elapsed,
    }


def richardson_extrapolate(rows: list[dict]) -> dict | None:
    """Simple 3-point Richardson extrapolation on the 3 finest tip_dy values,
    assuming geometric mesh refinement and monotone convergence O(h^p)."""
    if len(rows) < 3:
        return None
    y = [r["tip_dy"] for r in rows[-3:]]
    y1, y2, y3 = y
    denom = (y3 - y2) - (y2 - y1)
    if abs(denom) < 1e-12:
        return None
    # Richardson extrapolation assuming constant refinement ratio between levels
    y_inf = y3 - (y3 - y2) ** 2 / denom
    return {"tip_dy_extrapolated": y_inf}


def main():
    rows = [run_case(sf) for sf in MESH_RESOLUTIONS]

    finest, prev = rows[-1], rows[-2]
    self_convergence_pct = abs(finest["tip_dy"] - prev["tip_dy"]) / abs(finest["tip_dy"]) * 100.0
    richardson = richardson_extrapolate(rows)

    summary = {
        "case": "geometric_nonlinear / Cook's membrane",
        "citation": "Cook, R.D. (1974); geometry per the widely used Cook 1974 tapered-panel benchmark",
        "reference_caveat": (
            "The classical ~23.9 converged tip deflection figure applies to the LINEAR "
            "elastic 2D plane-stress version at unit load (E=1, nu=1/3, load=1); it is not "
            "a valid external ground truth for this geometric_nonlinear 3D-solid run at our "
            "load level (see module docstring). Validated via h-refinement mesh convergence "
            "instead, per the validation plan's documented fallback for this case."
        ),
        "geometry": {"vertices": [[0, 0], [LX, H0], [LX, H1], [0, H0]], "thickness_m": T},
        "material": {"E_Pa": E, "nu": NU},
        "load": {"traction_Pa": TRACTION_Y, "loaded_edge_area_m2": LOADED_EDGE_AREA,
                  "resultant_N": TRACTION_Y * LOADED_EDGE_AREA},
        "tip_point_definition": "midpoint of the loaded (x=48) edge, y-component of displacement",
        "mesh_convergence": rows,
        "richardson_extrapolation": richardson,
        "self_convergence_pct": self_convergence_pct,
        "pass": self_convergence_pct < 5.0,
        "pass_criterion": "relative change in tip_dy between the two finest meshes < 5% (mesh-convergence trend; no independent reference at this load level, see reference_caveat)",
    }
    (CASE_DIR / "results.json").write_text(json.dumps(summary, indent=2))
    print(json.dumps(summary, indent=2), flush=True)
    plot_results(rows, richardson)


def plot_results(rows, richardson):
    import matplotlib
    matplotlib.use("Agg")
    import matplotlib.pyplot as plt

    fig, axes = plt.subplots(1, 2, figsize=(11, 4.4), dpi=150)

    sf = [r["size_factor"] for r in rows]
    dy = [r["tip_dy"] for r in rows]
    ax = axes[0]
    ax.plot(sf, dy, "o-", color="#2563eb", label="TorsoCAE (dolfinx, P2 tet)", linewidth=2, markersize=7)
    if richardson:
        ax.axhline(richardson["tip_dy_extrapolated"], color="#16a34a", linestyle="--", linewidth=1.5,
                    label=f"Richardson extrapolation ({richardson['tip_dy_extrapolated']:.3f})")
    ax.set_xlabel("Mesh size_factor (coarser -> finer, right to left)")
    ax.set_ylabel("Loaded-edge midpoint tip deflection, y (units)")
    ax.set_title("Cook's membrane: mesh convergence (geometric_nonlinear)")
    ax.invert_xaxis()
    ax.legend(loc="lower left", fontsize=8)
    ax.grid(alpha=0.3)

    ax = axes[1]
    verts = [[0, 0], [LX, H0], [LX, H1], [0, H0], [0, 0]]
    vx = [v[0] for v in verts]
    vy = [v[1] for v in verts]
    ax.plot(vx, vy, "-", color="#94a3b8", linewidth=1.5, label="Undeformed panel")
    finest = rows[-1]
    tip_ref = finest["tip_node_ref_xyz"]
    tip_disp = finest["tip_disp_xyz"]
    ax.plot([tip_ref[0]], [tip_ref[1]], "s", color="#94a3b8", markersize=6)
    ax.plot([tip_ref[0] + tip_disp[0]], [tip_ref[1] + tip_disp[1]], "o", color="#2563eb",
             markersize=10, label="TorsoCAE tip (finest mesh)")
    ax.annotate("", xy=(tip_ref[0] + tip_disp[0], tip_ref[1] + tip_disp[1]), xytext=(tip_ref[0], tip_ref[1]),
                arrowprops=dict(arrowstyle="->", color="#2563eb"))
    ax.set_xlabel("x (units)")
    ax.set_ylabel("y (units)")
    ax.set_title("Loaded-edge tip displacement (finest mesh)")
    ax.axis("equal")
    ax.legend(loc="upper left", fontsize=8)
    ax.grid(alpha=0.3)

    fig.tight_layout()
    plots_dir = CASE_DIR / "plots"
    plots_dir.mkdir(exist_ok=True)
    fig.savefig(plots_dir / "convergence.png")


if __name__ == "__main__":
    main()