#!/usr/bin/env python3
"""mpmath O3 for mechanical.statics.plastic_moment. Mp = sigma_y * Z."""
from __future__ import annotations
import json, math, shutil
from pathlib import Path
import mpmath as mp


def pick(inputs, *keys):
    for key in keys:
        if key in inputs and inputs[key] not in (None, ""):
            return mp.mpf(str(inputs[key]))
    return None


def row(vid, inputs, kind):
    sy = pick(inputs, "sigma_y", "Fy", "yield_stress", "yield")
    Z = pick(inputs, "Z", "Zx", "plastic_modulus")
    Mp = sy * Z
    return {
        "id": vid,
        "kind": kind,
        "inputs": inputs,
        "Mp_f64": float(Mp),
        "Mp_decimal": mp.nstr(Mp, 40, strip_zeros=False),
        "sigma_y_f64": float(sy),
        "sigma_y_decimal": mp.nstr(sy, 40, strip_zeros=False),
        "Z_f64": float(Z),
        "Z_decimal": mp.nstr(Z, 40, strip_zeros=False),
    }


def main():
    mp.mp.dps = 80
    here = Path(__file__).resolve().parent
    vectors = [
        row("o3-unit", {"sigma_y": "250", "Z": "18"}, "unit"),
        row("o3-SI", {"sigma_y": "250e6", "Z": "8e-6"}, "SI"),
        row("o3-alias", {"Fy": "300", "plastic_modulus": "12"}, "alias"),
        row("o3-awkward", {"sigma_y": "255.5", "Z": "17.25"}, "awkward"),
        row("o3-Zx", {"sigma_y": "400", "Zx": "10"}, "Zx"),
        row("o3-small", {"sigma_y": "0.5", "Z": "0.02"}, "small"),
        row("o3-large", {"sigma_y": "1e9", "Z": "1e-3"}, "large"),
        row("o3-mild", {"sigma_y": "36", "Z": "100"}, "mild"),
    ]
    table = {
        "family": "plastic_moment",
        "generator_id": "plastic-moment-mpmath-o3",
        "generator_version": "1.0.0",
        "seed": "20260927.plastic-moment-o3",
        "mpmath_dps": 80,
        "precision_bits": int(80 * math.log2(10)),
        "library": f"mpmath {mp.__version__}",
        "notes": "Mp = sigma_y Z. Not elastic S, not transverse shear.",
        "vectors": vectors,
    }
    dest = here / "plastic-moment-o3-tables.json"
    dest.write_text(json.dumps(table, indent=2) + "\n", encoding="utf-8")
    public = here.parents[3] / "public/developers/cvp/reproduce"
    public.mkdir(parents=True, exist_ok=True)
    shutil.copy2(dest, public / "plastic-moment-o3-tables.json")
    shutil.copy2(Path(__file__), public / "generate-plastic-moment-o3.py")
    print(f"wrote {dest} ({len(vectors)} vectors)")


if __name__ == "__main__":
    main()