#!/usr/bin/env python3
"""Precompute mpmath high-precision O3 vectors for engineering.tolerance.monte_carlo.

Independent of Node/V8 binary64 and of the IUT. Tabulated fields are analytic
σ / moments / stack_center only (O2-A identities). Seeded Monte Carlo
replay stays O2-B. NORTA/Sobol/Shapley are not tabulated.

80 dps, round-to-nearest binary64. Does not pin a gmpy2 backend.

Re-run: python3 scripts/lib/cvp/oracles/generate-monte-carlo-tolerance-o3.py
"""
from __future__ import annotations

import json
import math
import shutil
from pathlib import Path

import mpmath as mp

DPS = 80
GENERATOR_ID = "monte-carlo-tolerance-mpmath-o3"
GENERATOR_VERSION = "1.0.0"
SEED = "20260916.mc-o3"


def f64(x: mp.mpf) -> float:
    return float(x)


def error_moments(tp, tm, k, dist: str):
    tp = abs(mp.mpf(tp))
    tm = abs(mp.mpf(tm))
    if dist == "uniform":
        low, high = -tm, tp
        mean = (low + high) / 2
        std = (high - low) / mp.sqrt(12)
        return mean, std
    if dist == "triangular":
        a, c, b = -tm, mp.mpf(0), tp
        mean = (a + b + c) / 3
        var = (a * a + b * b + c * c - a * b - a * c - b * c) / 18
        if var < 0:
            var = mp.mpf(0)
        return mean, mp.sqrt(var)
    sp, sm = tp / k, tm / k
    mean = (sp - sm) / mp.sqrt(2 * mp.pi)
    e2 = mp.mpf("0.5") * (sp * sp + sm * sm)
    var = e2 - mean * mean
    if var < 0:
        var = mp.mpf(0)
    return mean, mp.sqrt(var)


def analytic(components: list[dict], *, k=3, rho_pairs=None) -> dict:
    k = mp.mpf(k)
    moments = [error_moments(c["tp"], c["tm"], k, c.get("dist", "normal")) for c in components]
    signed = [mp.mpf(c.get("c", 1)) * m[1] for c, m in zip(components, moments)]
    stack_nom = sum((mp.mpf(c.get("c", 1)) * mp.mpf(c.get("nom", 0)) for c in components), mp.mpf(0))
    stack_shift = sum((mp.mpf(c.get("c", 1)) * mp.mpf(c.get("shift", 0)) for c in components), mp.mpf(0))
    stack_dist_mean = sum((mp.mpf(c.get("c", 1)) * m[0] for c, m in zip(components, moments)), mp.mpf(0))
    stack_center = stack_nom + stack_shift + stack_dist_mean
    v = sum((s * s for s in signed), mp.mpf(0))
    for i, j, rho in rho_pairs or []:
        v += 2 * mp.mpf(rho) * signed[i] * signed[j]
    if v < 0 and v > -mp.mpf("1e-40"):
        v = mp.mpf(0)
    std_analytic = mp.sqrt(v)
    return {
        "std_analytic_f64": f64(std_analytic),
        "std_analytic_decimal": mp.nstr(std_analytic, 40, strip_zeros=False),
        "stack_nom_f64": f64(stack_nom),
        "stack_nom_decimal": mp.nstr(stack_nom, 40, strip_zeros=False),
        "stack_dist_mean_f64": f64(stack_dist_mean),
        "stack_dist_mean_decimal": mp.nstr(stack_dist_mean, 40, strip_zeros=False),
        "stack_center_f64": f64(stack_center),
        "stack_center_decimal": mp.nstr(stack_center, 40, strip_zeros=False),
    }


def row(vid: str, inputs: dict, extra: dict, *, kind: str) -> dict:
    return {"id": vid, "kind": kind, "inputs": inputs, **extra}


def vectors() -> list[dict]:
    k3 = 3
    return [
        row(
            "o3-mc",
            {"tol1": 0.1, "tol2": 0.2, "N": 2000, "seed": 42},
            analytic([{"tp": 0.1, "tm": 0.1}, {"tp": 0.2, "tm": 0.2}], k=k3),
            kind="normal-independent",
        ),
        row(
            "o3-corr",
            {
                "tolerances": [
                    {"name": "A", "half_tolerance": 0.3},
                    {"name": "B", "half_tolerance": 0.3},
                ],
                "correlations": [{"a": "A", "b": "B", "rho": 0.5}],
                "N": 2000,
                "seed": 3,
                "k_sigma": 3,
            },
            analytic(
                [{"tp": 0.3, "tm": 0.3}, {"tp": 0.3, "tm": 0.3}],
                k=3,
                rho_pairs=[(0, 1, "0.5")],
            ),
            kind="normal-correlated",
        ),
        row(
            "o3-gap",
            {
                "tolerances": [
                    {"name": "A", "nominal": 10, "half_tolerance": 0.1, "coefficient": 1},
                    {"name": "B", "nominal": 5, "half_tolerance": 0.2, "coefficient": -1},
                ],
                "N": 2000,
                "seed": 42,
                "LSL": 4.8,
                "USL": 5.2,
            },
            analytic(
                [{"tp": 0.1, "tm": 0.1, "nom": 10, "c": 1}, {"tp": 0.2, "tm": 0.2, "nom": 5, "c": -1}],
                k=k3,
            ),
            kind="gap-analytic",
        ),
        row(
            "o3-uniform",
            {"tolerances": [{"half_tolerance": 0.3, "distribution": "uniform"}], "N": 2000, "seed": 1},
            analytic([{"tp": 0.3, "tm": 0.3, "dist": "uniform"}], k=k3),
            kind="uniform",
        ),
        row(
            "o3-asym",
            {
                "tolerances": [{"name": "A", "tol_plus": 0.2, "tol_minus": 0.05, "distribution": "uniform"}],
                "N": 2000,
                "seed": 2,
            },
            analytic([{"tp": 0.2, "tm": 0.05, "dist": "uniform"}], k=k3),
            kind="asymmetric-uniform",
        ),
        row(
            "o3-tri",
            {"tolerances": [{"half_tolerance": 0.3, "distribution": "triangular"}], "N": 200, "seed": 1},
            analytic([{"tp": 0.3, "tm": 0.3, "dist": "triangular"}], k=k3),
            kind="triangular",
        ),
        row(
            "o3-pythag",
            {"tol1": 0.3, "tol2": 0.4, "k_sigma": 1, "N": 200, "seed": 1},
            analytic([{"tp": 0.3, "tm": 0.3}, {"tp": 0.4, "tm": 0.4}], k=1),
            kind="k1-pythagorean",
        ),
        row(
            "o3-one",
            {"tol1": 0.3, "N": 200, "seed": 1},
            analytic([{"tp": 0.3, "tm": 0.3}], k=k3),
            kind="single-normal",
        ),
    ]


def main() -> None:
    mp.mp.dps = DPS
    here = Path(__file__).resolve().parent
    table = {
        "family": "monte_carlo_tolerance",
        "generator_id": GENERATOR_ID,
        "generator_version": GENERATOR_VERSION,
        "seed": SEED,
        "mpmath_dps": DPS,
        "precision_bits": int(DPS * math.log2(10)),
        "library": f"mpmath {mp.__version__}",
        "notes": "Analytic σ/moments only; seeded MC remains O2-B. Independent of Node Math and the IUT.",
        "vectors": vectors(),
    }
    dest = here / "monte-carlo-tolerance-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 / "monte-carlo-tolerance-o3-tables.json")
    shutil.copy2(Path(__file__), public / "generate-monte-carlo-tolerance-o3.py")
    print(f"wrote {dest} ({len(table['vectors'])} vectors)")


if __name__ == "__main__":
    main()
