"""Standalone CPU inference for the converted MagNet steady-state Transformers.

Run from an unpacked model repository. This example only loads ONNX/JSON/NPY
files and never unpickles the original PyTorch checkpoints.
"""

from __future__ import annotations

import argparse
import json
import math
from pathlib import Path

import numpy as np
import onnxruntime as ort

ROOT = Path(__file__).resolve().parent
MATERIALS = ("3C90", "3C94", "3E6", "3F4", "77", "78", "N27", "N30", "N49", "N87")


def _finite(value: float, name: str) -> float:
    number = float(value)
    if not math.isfinite(number):
        raise ValueError(f"{name} must be finite")
    return number


def make_waveform(kind: str, flux_t: float, duty: float = 0.5) -> np.ndarray:
    """Return 128 B samples in T, following the historical MagNet convention."""
    flux_t = _finite(flux_t, "flux_t")
    if flux_t <= 0:
        raise ValueError("flux_t must be a positive peak AC flux density in T")
    if kind == "sine":
        return (flux_t * np.sin(np.linspace(0, 2 * np.pi, 129)))[:-1]
    if kind != "triangle":
        raise ValueError("waveform must be sine or triangle")
    duty = _finite(duty, "duty")
    if not 0 < duty < 1:
        raise ValueError("triangle duty ratio must lie strictly between 0 and 1")
    b = np.interp(np.linspace(0, 1, 129), [0, duty, 1], [-flux_t, flux_t, -flux_t])[:-1]
    phase = 0.5 - duty / 2 if duty <= 0.5 else -duty / 2
    return np.roll(b, int(phase * 128))


def load_waveform_json(path: Path) -> np.ndarray:
    values = json.loads(path.read_text(encoding="utf-8"))
    if not isinstance(values, list) or len(values) != 128:
        raise ValueError("--b-json must contain a JSON array of exactly 128 B values in T")
    if any(isinstance(value, bool) or not isinstance(value, (float, int)) for value in values):
        raise ValueError("B values must all be JSON numbers")
    return np.asarray(values, dtype=np.float64)


def _normal(value: float, center: np.float32, scale: np.float32) -> np.float32:
    # Preserve the float32 arithmetic of the legacy PyTorch inference.
    return np.float32((np.float32(value) - center) / scale)


def measured_hull(
    material: str,
    frequency_hz: float,
    flux_t: float,
    bias_a_per_m: float,
    temperature_c: float,
    duty: float,
) -> bool:
    equations = np.load(ROOT / "domain" / f"hull_{material}.npy", allow_pickle=False)
    if (
        equations.ndim != 2
        or equations.shape[1] != 6
        or not np.issubdtype(equations.dtype, np.number)
    ):
        raise ValueError(f"invalid numeric measured-domain hull for {material}")
    point = np.array([frequency_hz, flux_t, bias_a_per_m, temperature_c, duty], dtype=np.float64)
    return bool(np.all(equations[:, :5] @ point + equations[:, 5] <= 1e-10))


def predict(
    material: str,
    b_waveform_t: np.ndarray,
    frequency_hz: float,
    temperature_c: float,
    bias_a_per_m: float,
    nominal_flux_t: float | None = None,
    duty: float | None = None,
) -> dict:
    """Predict the 128-point H waveform and signed loss density.

    nominal_flux_t and duty enable an ideal-waveform hull check. Omit both for
    an arbitrary B waveform, for which this historical hull is not applicable.
    """
    if material not in MATERIALS:
        raise ValueError(f"unknown material {material!r}; choose from {', '.join(MATERIALS)}")
    b = np.asarray(b_waveform_t, dtype=np.float64)
    if b.shape != (128,) or not np.isfinite(b).all() or np.max(b) == np.min(b):
        raise ValueError("B waveform must have 128 finite, varying samples in T")
    frequency_hz = _finite(frequency_hz, "frequency_hz")
    temperature_c = _finite(temperature_c, "temperature_c")
    bias_a_per_m = _finite(bias_a_per_m, "bias_a_per_m")
    if frequency_hz <= 0:
        raise ValueError("frequency_hz must be greater than zero")
    if (nominal_flux_t is None) != (duty is None):
        raise ValueError("nominal_flux_t and duty must both be given or both omitted")
    if nominal_flux_t is not None:
        nominal_flux_t = _finite(nominal_flux_t, "nominal_flux_t")
        duty = _finite(duty, "duty")
        if nominal_flux_t <= 0 or not 0 < duty < 1:
            raise ValueError("nominal_flux_t must be positive and duty strictly between 0 and 1")

    norm_values = json.loads((ROOT / material / "norm.json").read_text(encoding="utf-8"))
    if not isinstance(norm_values, list) or len(norm_values) != 10:
        raise ValueError("normalizer must contain ten values")
    norm = np.asarray(norm_values, dtype=np.float32)
    if not np.isfinite(norm).all() or np.any(norm[[1, 3, 5, 7]] == 0):
        raise ValueError("invalid normalization values")

    source = ((b.astype(np.float32) - norm[0]) / norm[1]).reshape(1, 128, 1)
    variables = np.array(
        [
            _normal(np.float32(np.log10(frequency_hz)), norm[2], norm[3]),
            _normal(temperature_c, norm[4], norm[5]),
            _normal(bias_a_per_m, norm[6], norm[7]),
        ],
        dtype=np.float32,
    ).reshape(1, 3)
    options = ort.SessionOptions()
    options.intra_op_num_threads = 1
    encoder = ort.InferenceSession(
        str(ROOT / material / "encoder.onnx"), options, providers=["CPUExecutionProvider"]
    )
    decoder = ort.InferenceSession(
        str(ROOT / material / "decoder.onnx"), options, providers=["CPUExecutionProvider"]
    )
    memory = encoder.run(["memory"], {"source": source, "variables": variables})[0]
    if memory.shape != (1, 128, 24):
        raise ValueError("unexpected encoder output shape")
    target = np.zeros((1, 129, 1), dtype=np.float32)
    target[0, 0, 0] = 0.1
    for step in range(1, 129):
        output = decoder.run(["output"], {"memory": memory, "target": target})[0]
        target[0, step, 0] = output[0, step - 1, 0]
    output = decoder.run(["output"], {"memory": memory, "target": target})[0]
    if output.shape != (1, 129, 1):
        raise ValueError("unexpected decoder output shape")
    h = (output[0, :128, 0] * norm[9] + norm[8]).astype(np.float32)
    loss = float(frequency_hz * np.trapezoid(h, b))
    if not np.isfinite(h).all() or not math.isfinite(loss):
        raise ValueError("model returned a non-finite prediction")

    in_hull = None
    if nominal_flux_t is not None:
        in_hull = measured_hull(
            material, frequency_hz, nominal_flux_t, bias_a_per_m, temperature_c, duty
        )
    warnings = []
    if in_hull is False:
        warnings.append("Operating point lies outside the historical measured-data convex hull.")
    if in_hull is None:
        warnings.append("Measured-domain hull is not applicable to an arbitrary waveform.")
    if loss < 0:
        warnings.append("Negative signed B-H loop integral; inspect waveform and operating point.")
    return {
        "model_id": "legacy-transformer-v1",
        "material_id": material,
        "frequency_hz": frequency_hz,
        "temperature_c": temperature_c,
        "bias_a_per_m": bias_a_per_m,
        "b_waveform_t": b.tolist(),
        "h_waveform_a_per_m": h.tolist(),
        "loss_w_per_m3": loss,
        "inside_measured_hull": in_hull,
        "units": {"flux_density": "T", "field_strength": "A/m", "frequency": "Hz", "loss": "W/m^3"},
        "warnings": warnings,
    }


def main() -> None:
    parser = argparse.ArgumentParser(description=__doc__)
    parser.add_argument("--material", required=True, choices=MATERIALS)
    parser.add_argument("--frequency-hz", type=float, required=True)
    parser.add_argument("--temperature-c", type=float, default=25.0)
    parser.add_argument("--bias-a-per-m", type=float, default=0.0)
    waveform = parser.add_mutually_exclusive_group()
    waveform.add_argument("--waveform", choices=("sine", "triangle"), default="sine")
    waveform.add_argument("--b-json", type=Path)
    parser.add_argument(
        "--flux-t", type=float, help="positive peak AC B in T for an ideal waveform"
    )
    parser.add_argument("--duty", type=float, default=0.5, help="triangle peak-time fraction (0,1)")
    args = parser.parse_args()
    if args.b_json is not None:
        if args.flux_t is not None:
            parser.error("--flux-t cannot be combined with --b-json")
        b = load_waveform_json(args.b_json)
        flux = duty = None
    else:
        if args.flux_t is None:
            parser.error("--flux-t is required for --waveform")
        b = make_waveform(args.waveform, args.flux_t, args.duty)
        flux = args.flux_t
        duty = args.duty if args.waveform == "triangle" else 0.5
    result = predict(
        args.material, b, args.frequency_hz, args.temperature_c, args.bias_a_per_m, flux, duty
    )
    print(json.dumps(result, indent=2))


if __name__ == "__main__":
    main()