#!/usr/bin/env python3
"""Export the RSL-RL FireAnt actor from a training checkpoint to ONNX."""

from __future__ import annotations

import argparse
from pathlib import Path

import numpy as np
import onnx
import onnxruntime as ort
import torch


PROJECT_ROOT = Path(__file__).resolve().parents[1]
DEFAULT_CHECKPOINT = (
    PROJECT_ROOT.parent
    / "rl_training"
    / "logs"
    / "rsl_rl"
    / "deeprobotics_m20_support"
    / "2026-09-12_17-52-15_stage2_multiterrain"
    / "model_39899.pt"
)
DEFAULT_OUTPUT = PROJECT_ROOT / "policy" / "fire_ant_policy.onnx"


def build_actor() -> torch.nn.Module:
    return torch.nn.Sequential(
        torch.nn.Linear(63, 512),
        torch.nn.ELU(),
        torch.nn.Linear(512, 256),
        torch.nn.ELU(),
        torch.nn.Linear(256, 128),
        torch.nn.ELU(),
        torch.nn.Linear(128, 8),
    )


def parse_args() -> argparse.Namespace:
    parser = argparse.ArgumentParser(description=__doc__)
    parser.add_argument("--checkpoint", type=Path, default=DEFAULT_CHECKPOINT)
    parser.add_argument("--output", type=Path, default=DEFAULT_OUTPUT)
    return parser.parse_args()


def main() -> None:
    args = parse_args()
    if not args.checkpoint.is_file():
        raise FileNotFoundError(args.checkpoint)

    checkpoint = torch.load(args.checkpoint, map_location="cpu", weights_only=False)
    state_dict = checkpoint["model_state_dict"]
    actor_state = {
        name.removeprefix("actor."): tensor
        for name, tensor in state_dict.items()
        if name.startswith("actor.")
    }
    actor = build_actor().eval()
    actor.load_state_dict(actor_state, strict=True)

    args.output.parent.mkdir(parents=True, exist_ok=True)
    sample = torch.zeros(1, 63, dtype=torch.float32)
    with torch.no_grad():
        torch_output = actor(sample).cpu().numpy()
    torch.onnx.export(
        actor,
        sample,
        args.output,
        input_names=["obs"],
        output_names=["actions"],
        dynamic_axes={"obs": {0: "batch"}, "actions": {0: "batch"}},
        opset_version=17,
        do_constant_folding=True,
        dynamo=False,
    )

    model = onnx.load(args.output)
    onnx.checker.check_model(model)
    session = ort.InferenceSession(str(args.output), providers=["CPUExecutionProvider"])
    ort_output = session.run(["actions"], {"obs": sample.numpy()})[0]
    max_error = float(np.max(np.abs(torch_output - ort_output)))
    if max_error > 1e-5:
        raise RuntimeError(f"ONNX validation failed: max error {max_error:.3e}")
    print(f"Exported {args.checkpoint}")
    print(f"ONNX: {args.output} (input 63, output 8, max error {max_error:.3e})")


if __name__ == "__main__":
    main()