#!/usr/bin/env python3
"""Analytic membrane reference -- what the reported endpoint actually contains.

This script needs no FEM run and no project tree. It works out, in closed form,
what the published proxy is made of, and it changes how two of the paper's
headline sentences should be written.

Finding 1 -- the prestress enters the REPORTED number as an additive constant
-----------------------------------------------------------------------------
The solver and the post-processor treat the initial stress differently:

    equilibrium (nonlinear_solver.py, _region_forms)
        PK1 += (F - I) . S0          -> exactly zero at u = 0
    reported proxy (_write_outer_surface_resultants)
        sigma += F . S0 . F^T / J    -> equals S0 at u = 0

Both are defensible on their own. Together they mean the reference configuration
is in equilibrium only because a constant reference-state term was subtracted
from the residual, i.e. the model silently assumes a fictitious balancing load
that holds the prestressed reference in equilibrium and is then carried, dead,
through the whole deformation. That assumption is never stated in the manuscript
and it must be, because a prestressed closed shell cannot be in equilibrium on
its own without a balancing core pressure.

The consequence for the number: the proxy multiplies a tangential principal
Cauchy stress by the CURRENT thickness, and the current thickness is recovered
as lambda_n = J / J_s (Nanson). For the prestress part those two factors cancel:

    N_pre = 0.5 * tr_tangential(F S0 F^T / J) * t0 * (J / J_s)
          = S0_t * t0 * (lambda_1/lambda_2 + lambda_2/lambda_1) / 2

which is exactly S0_t * t0 for equibiaxial in-plane stretch and never less
(AM-GM). With prestress_pa = 1000 * 0.05 / 0.12 and t0 = 0.12 um that constant
is 0.05 mN/m -- the assumed pretension, unchanged by any amount of shrinkage.

Finding 2 -- so the elastic membrane resultant reverses sign
------------------------------------------------------------
Subtracting the constant from the published endpoints:

    Control       0.066064 - 0.05 = +0.016064 mN/m
    Hyperosmotic  0.038092 - 0.05 = -0.011908 mN/m

The elastic part does not merely relax, it goes into net compression. That is
independently corroborated by Fig. S4d, where the hyperosmotic compression area
fraction is 62.6 %. Three things follow:

  * The ABSOLUTE difference (-0.027972 mN/m) is essentially free of the direct
    prestress contribution, because the constant cancels. An earlier reading of
    ours -- that the pretension assumption drives the difference -- was wrong in
    that direct sense and is withdrawn.
  * The PERCENT change is a different story. -42.34 % is computed against a
    denominator that is 76 % additive assumption. The underlying elastic change
    is about -174 %. Moving the pretension across its own UQ range (0.025 to
    0.1 mN/m) rescales the reported percentage while barely touching the
    mechanics, which is precisely why the UQ found a Spearman rho of +0.29 for
    pretension. Report the absolute difference as primary.
  * A shell in net elastic compression is a shell that would wrinkle. The
    stability check (05_run_stability_check.py) is therefore not a formality.

Finding 3 -- closed-form membrane resultant, which makes the direction analytic
-------------------------------------------------------------------------------
For an incompressible-in-plane Neo-Hookean membrane with isotropic growth g,
in-plane stretch l and thickness stretch l_t, imposing plane stress sigma_33 = 0
on the implemented strain energy gives

    N_elastic = (mu * t0 / g^2) * (1 - (l_t / l)^2)

N vanishes exactly at the natural state l = l_t = g, is positive when the
membrane is stretched beyond its natural configuration and negative when it is
compressed below it. Monotonicity of N in l is checked numerically below. This
turns "all 64 LHS points and all 77 boundary points unload" from an empirical
robustness claim into a one-line consequence of the constitutive model, which is
what the manuscript should say instead of spending a main figure on it.

Usage:
    python3 17_analytic_membrane_reference.py [--pretension-mn-per-m 0.05]
                                              [--shell-thickness-um 0.12]
                                              [--control 0.066064]
                                              [--hyperosmotic 0.038092]
"""

from __future__ import annotations

import argparse
import json

import numpy as np


def prestress_mean_resultant(deformation, reference_normal, prestress_pa, thickness_um):
    """Mean tangential prestress resultant in mN/m, exactly as the code forms it."""
    normal = np.asarray(reference_normal, dtype=float)
    normal = normal / np.linalg.norm(normal)
    initial_stress = prestress_pa * (np.eye(3) - np.outer(normal, normal))

    jacobian = np.linalg.det(deformation)
    cauchy = deformation @ initial_stress @ deformation.T / jacobian

    # Nanson: current area vector, its norm is the area stretch J_s.
    cofactor = jacobian * np.linalg.inv(deformation).T
    area_vector = cofactor @ normal
    area_stretch = np.linalg.norm(area_vector)
    current_normal = area_vector / area_stretch
    thickness_stretch = jacobian / area_stretch

    tangential = np.eye(3) - np.outer(current_normal, current_normal)
    # Trace of the tangential block is the sum of the two tangential principals,
    # and trace is basis independent, so the mean resultant is exactly additive
    # in (elastic + prestress) even though the principal directions are not.
    trace_tangential = np.trace(tangential @ cauchy @ tangential)
    return 0.5 * trace_tangential * thickness_um * thickness_stretch * 1.0e-3


def elastic_membrane_resultant(stretch, growth, shear_pa, poisson, thickness_um):
    """Closed-form plane-stress membrane resultant, mN/m.

    Solves sigma_33 = 0 for the thickness stretch, then evaluates
    N = (mu t0 / g^2) (1 - (l_t/l)^2).
    """
    lame_lambda = 2.0 * shear_pa * poisson / (1.0 - 2.0 * poisson)

    def normal_stress(thickness_stretch):
        elastic_jacobian = stretch**2 * thickness_stretch / growth**3
        return (
            shear_pa * thickness_stretch**2 / growth**2
            + lame_lambda * np.log(elastic_jacobian)
            - shear_pa
        )

    low, high = 1.0e-6, 10.0
    for _ in range(200):
        mid = 0.5 * (low + high)
        if normal_stress(mid) > 0.0:
            high = mid
        else:
            low = mid
    thickness_stretch = 0.5 * (low + high)
    resultant_pa_um = (
        shear_pa * thickness_um / growth**2 * (1.0 - (thickness_stretch / stretch) ** 2)
    )
    return resultant_pa_um * 1.0e-3, thickness_stretch


def main() -> int:
    parser = argparse.ArgumentParser(description=__doc__)
    parser.add_argument("--pretension-mn-per-m", type=float, default=0.05)
    parser.add_argument("--shell-thickness-um", type=float, default=0.12)
    parser.add_argument("--control", type=float, default=0.066064)
    parser.add_argument("--hyperosmotic", type=float, default=0.038092)
    parser.add_argument("--shell-young-pa", type=float, default=2000.0)
    parser.add_argument("--shell-poisson", type=float, default=0.35)
    parser.add_argument("--shell-coupling", type=float, default=0.35)
    parser.add_argument("--natural-volume-ratio", type=float, default=0.725517)
    parser.add_argument("--seed", type=int, default=0)
    parser.add_argument("--json-out", default=None)
    args = parser.parse_args()

    prestress_pa = 1000.0 * args.pretension_mn_per_m / args.shell_thickness_um
    report: dict = {"prestress_pa": prestress_pa}

    print("=" * 72)
    print("1. Prestress contribution to the reported mean resultant")
    print("=" * 72)
    print(f"prestress_pa = 1000 * {args.pretension_mn_per_m} / {args.shell_thickness_um}"
          f" = {prestress_pa:.4f} Pa")
    print(f"predicted invariant contribution = S0_t * t0 = "
          f"{prestress_pa * args.shell_thickness_um * 1e-3:.6f} mN/m\n")

    cases = [
        ("undeformed", np.eye(3)),
        ("equibiaxial contraction 0.90", np.diag([0.90, 0.90, 0.90])),
        ("equibiaxial contraction 0.75", np.diag([0.75, 0.75, 0.75])),
        ("stretched 1.15", np.diag([1.15, 1.15, 1.15])),
        ("in-plane anisotropic 1.2/0.8", np.diag([1.2, 0.8, 1.0])),
    ]
    normal = np.array([0.0, 0.0, 1.0])
    invariant = []
    for label, deformation in cases:
        value = prestress_mean_resultant(
            deformation, normal, prestress_pa, args.shell_thickness_um
        )
        invariant.append((label, value))
        print(f"  {label:<30} N_pre = {value:.6f} mN/m")

    rng = np.random.default_rng(args.seed)
    randoms = []
    while len(randoms) < 200:
        deformation = np.eye(3) + 0.3 * rng.standard_normal((3, 3))
        if np.linalg.det(deformation) < 0.05:
            continue
        direction = rng.standard_normal(3)
        randoms.append(
            prestress_mean_resultant(
                deformation, direction, prestress_pa, args.shell_thickness_um
            )
        )
    randoms = np.asarray(randoms)
    print(f"\n  200 random deformations: min {randoms.min():.6f}, "
          f"median {np.median(randoms):.6f}, max {randoms.max():.6f} mN/m")
    print("  Never below the equibiaxial value -- the excess is pure in-plane")
    print("  anisotropy (AM-GM), not a response to volume change.")
    report["prestress_invariant_cases"] = dict(invariant)
    report["prestress_random_min"] = float(randoms.min())

    print()
    print("=" * 72)
    print("2. Implied elastic decomposition of the published endpoints")
    print("=" * 72)
    offset = args.pretension_mn_per_m
    control_elastic = args.control - offset
    hyper_elastic = args.hyperosmotic - offset
    total_difference = args.hyperosmotic - args.control
    print(f"  Control       {args.control:.6f} - {offset:.4f} = {control_elastic:+.6f} mN/m")
    print(f"  Hyperosmotic  {args.hyperosmotic:.6f} - {offset:.4f} = {hyper_elastic:+.6f} mN/m")
    print(f"  reported percent change   {100 * total_difference / args.control:+.2f}%")
    print(f"  elastic  percent change   {100 * (hyper_elastic - control_elastic) / control_elastic:+.2f}%")
    print(f"  absolute difference       {total_difference:+.6f} mN/m (offset cancels)")
    print(f"  offset share of Control   {100 * offset / args.control:.0f}%")
    if control_elastic > 0.0 > hyper_elastic:
        print("\n  The elastic membrane resultant REVERSES SIGN: the hyperosmotic")
        print("  shell is in net elastic compression. Consistent with the 62.6 %")
        print("  compression area in Fig. S4d. Wrinkling is unmodelled, so the")
        print("  tangent stability check is required before this is reportable.")
    report["decomposition"] = {
        "control_elastic_mn_per_m": control_elastic,
        "hyperosmotic_elastic_mn_per_m": hyper_elastic,
        "absolute_difference_mn_per_m": total_difference,
        "reported_percent": 100 * total_difference / args.control,
        "elastic_percent": 100 * (hyper_elastic - control_elastic) / control_elastic,
        "offset_share_of_control_percent": 100 * offset / args.control,
        "sign_reversal": bool(control_elastic > 0.0 > hyper_elastic),
    }

    print()
    print("=" * 72)
    print("3. Closed-form elastic membrane resultant and its monotonicity")
    print("=" * 72)
    shear = args.shell_young_pa / (2.0 * (1.0 + args.shell_poisson))
    growth = args.natural_volume_ratio ** (args.shell_coupling / 3.0)
    print(f"  shell mu = {shear:.1f} Pa, shell growth stretch g = {growth:.6f}")
    print(f"  ({args.natural_volume_ratio}^({args.shell_coupling}/3), i.e. the shell's own natural shrinkage)\n")
    print("   stretch l   thickness l_t    N_elastic (mN/m)")
    stretches = np.linspace(0.80, 1.10, 13)
    values = []
    for stretch in stretches:
        resultant, thickness_stretch = elastic_membrane_resultant(
            stretch, growth, shear, args.shell_poisson, args.shell_thickness_um
        )
        values.append(resultant)
        print(f"   {stretch:8.4f}   {thickness_stretch:10.6f}    {resultant:+.6f}")
    values = np.asarray(values)
    crossing = float(np.interp(0.0, values, stretches))
    print(f"\n  N = 0 at l = {crossing:.6f}; the shell's natural stretch is g = {growth:.6f}")
    print("  (they coincide, as they must: zero elastic resultant at the natural state)")
    increasing = bool(np.all(np.diff(values) > 0))
    print(f"\n  strictly increasing in l over the sampled range: {increasing}")
    if increasing:
        print("  => dN/dl > 0. Isotropic natural shrinkage lowers l, therefore it")
        print("     lowers N. The direction of the published result is a property")
        print("     of the constitutive model, not an empirical finding, and no")
        print("     point of the parameter box can reverse it. State it in one")
        print("     line and reassign Main Fig. 3 to magnitude and decomposition.")
    report["monotonic_in_stretch"] = increasing
    report["natural_stretch"] = growth

    if args.json_out:
        with open(args.json_out, "w", encoding="utf-8") as handle:
            json.dump(report, handle, indent=2)
        print(f"\nWrote {args.json_out}")
    return 0


if __name__ == "__main__":
    raise SystemExit(main())
