"""Minimal material-point inference for original or expanded DENIM weights."""

import argparse
from pathlib import Path

import torch

from src.t2_graybox_discrete_energy import DENIM, rollout


ROOT = Path(__file__).resolve().parent
parser = argparse.ArgumentParser()
parser.add_argument(
    "--checkpoint",
    choices=("denim.pt", "denim-expanded.pt"),
    default="denim-expanded.pt",
)
args = parser.parse_args()
checkpoint = torch.load(ROOT / args.checkpoint, map_location="cpu", weights_only=False)
model = DENIM(channels=2)
model.load_state_dict(checkpoint["state_dict"])
model.eval()

points = 121
coordinate = torch.linspace(0.0, 1.0, points)
amplitude = torch.where(coordinate <= 0.5, 2.0 * coordinate, 2.0 * (1.0 - coordinate))
scalar = 0.008 * torch.sin(2.0 * torch.pi * coordinate) * amplitude
basis = torch.tensor((1.0, -0.5, -0.5, 0.0, 0.0, 0.0))
strain = (scalar[:, None] * basis).unsqueeze(0)

with torch.no_grad():
    result = rollout(
        strain,
        torch.tensor((190.0e9,)),
        torch.tensor((0.30,)),
        torch.tensor((280.0e6,)),
        model,
    )

print("checkpoint", args.checkpoint)
print("steps", points)
print("maximum_abs_stress_mpa", float(result["stress"].abs().max() / 1.0e6))
print("final_peeq", float(result["peeq"][0, -1]))
print("maximum_yield_residual_pa", float(result["yield_residual"].abs().max()))