"""Minimal offline inference for the material-conditioned DENIM v2 bundle."""

from __future__ import annotations

import argparse
import json
import sys
from pathlib import Path

import torch
from safetensors.torch import load_file


def run(bundle: Path) -> dict[str, float]:
    repository_root = bundle.parent
    if (repository_root / "src").is_dir():
        sys.path.insert(0, str(repository_root))
    from src.t2_denim_conditional import ConditionalDENIM, material_descriptor, rollout

    model = ConditionalDENIM()
    model.load_state_dict(load_file(bundle / "weights.safetensors"))
    model.eval()
    parameters = {
        "material_model": "chaboche_combined",
        "young_pa": 190.0e9,
        "poisson": 0.30,
        "yield_stress_pa": 280.0e6,
        "backstress_c1_pa": 28.0e9,
        "backstress_gamma1": 80.0,
        "backstress_c2_pa": 10.0e9,
        "backstress_gamma2": 12.0,
        "isotropic_saturation_pa": 70.0e6,
        "isotropic_rate": 8.0,
    }
    strain = torch.zeros((1, 121, 6))
    strain[0, :, 0] = torch.linspace(0.0, 0.01, 121)
    with torch.no_grad():
        result = rollout(
            strain,
            torch.tensor([parameters["young_pa"]]),
            torch.tensor([parameters["poisson"]]),
            torch.tensor([parameters["yield_stress_pa"]]),
            material_descriptor(parameters)[None],
            model,
        )
    summary = {
        "maximum_absolute_stress_mpa": float(result["stress"].abs().max() / 1.0e6),
        "final_peeq": float(result["peeq"][0, -1]),
        "maximum_yield_residual_pa": float(result["yield_residual"].abs().max()),
    }
    print(json.dumps(summary, indent=2))
    return summary


if __name__ == "__main__":
    parser = argparse.ArgumentParser()
    parser.add_argument("--bundle", type=Path, default=Path(__file__).resolve().parent)
    run(parser.parse_args().bundle)