#!/usr/bin/env python3
"""Export a DreamWaQ rsl_rl checkpoint (ActorCritic_DWAQ) to ONNX for MuJoCo.

rsl_rl has NO state preprocessor, so no normalization is needed (unlike SKRL).
The exported model uses the MEAN CENet code (deterministic deploy).

ONNX: inputs observations(1,45) + obs_history(1,5,45) -> actions(1,12)
(matches scripts/dreamwaq_sim2sim_mujoco.py interface)

Usage:
    uv run scripts/export_dreamwaq_rsl_onnx.py            # auto-find latest
    uv run scripts/export_dreamwaq_rsl_onnx.py --checkpoint runs/.../model_700.pt
"""
import argparse, glob, os, sys
import torch
import torch.nn as tnn

sys.path.insert(0, os.path.dirname(os.path.dirname(os.path.abspath(__file__))))
from motrix_rl.dwaq_rsl.actor_critic_dwaq import ActorCritic_DWAQ

PROJECT = os.path.dirname(os.path.dirname(os.path.abspath(__file__)))


class DwaqInfer(tnn.Module):
    """Deterministic inference: CENet mean code + RunningStats normalize + actor."""
    def __init__(self, ac: ActorCritic_DWAQ):
        super().__init__()
        self.encoder = ac.encoder
        self.encode_mean_vel = ac.encode_mean_vel
        self.encode_mean_latent = ac.encode_mean_latent
        self.actor = ac.actor
        # Bake RunningStats into the exported model
        if hasattr(ac, 'actor_normalizer') and ac.actor_normalizer.count > 10:
            self.register_buffer("norm_mean", ac.actor_normalizer.mean.clone())
            self.register_buffer("norm_std", ac.actor_normalizer.var.sqrt().clone())
            self.has_norm = True
        else:
            self.has_norm = False

    def forward(self, obs, obs_history):
        h = self.encoder(obs_history.reshape(obs_history.shape[0], -1))  # (B,225)->(B,64)
        vel = self.encode_mean_vel(h)        # (B,3)  mean velocity estimate
        latent = self.encode_mean_latent(h)  # (B,16) mean latent
        code = torch.cat([vel, latent], dim=-1)   # (B,19) = [vel, latent]
        x = torch.cat([code, obs], dim=-1)        # (B,64) = [code, obs]
        if self.has_norm:
            x = torch.clamp((x - self.norm_mean) / (self.norm_std + 1e-8), -5.0, 5.0)
        return self.actor(x)                      # (B,12)


def main():
    p = argparse.ArgumentParser()
    p.add_argument("--checkpoint", default=None)
    p.add_argument("--output", default=os.path.join(PROJECT, "exports_go1_dreamwaq", "policy.onnx"))
    p.add_argument("--num-obs", type=int, default=45)
    p.add_argument("--num-priv", type=int, default=235)
    p.add_argument("--num-hist", type=int, default=5)
    p.add_argument("--num-act", type=int, default=12)
    args = p.parse_args()

    if args.checkpoint is None:
        runs = sorted(glob.glob(os.path.join(PROJECT, "runs", "go1-dreamwaq-walk", "rsl_dwaq", "*")),
                      key=os.path.getmtime)
        if not runs:
            print("[ERROR] no rsl_dwaq runs found"); sys.exit(1)
        models = glob.glob(os.path.join(runs[-1], "model_*.pt"))
        args.checkpoint = max(models, key=os.path.getmtime)

    print(f"[rsl-ONNX] checkpoint: {args.checkpoint}")
    cenet_out = 19
    ac = ActorCritic_DWAQ(
        args.num_obs + cenet_out,   # actor in = 64
        args.num_priv,              # critic in = 235
        args.num_act,               # 12
        args.num_hist * args.num_obs,  # cenet in = 225
        cenet_out,                  # 19
    )
    ckpt = torch.load(args.checkpoint, map_location="cpu")
    ac.load_state_dict(ckpt["model_state_dict"])
    ac.eval()

    model = DwaqInfer(ac).eval()
    os.makedirs(os.path.dirname(args.output), exist_ok=True)
    torch.onnx.export(
        model,
        (torch.zeros(1, args.num_obs), torch.zeros(1, args.num_hist, args.num_obs)),
        args.output,
        input_names=["observations", "obs_history"],
        output_names=["actions"],
        opset_version=11,
        dynamic_axes={"observations": {0: "batch"}, "obs_history": {0: "batch"}, "actions": {0: "batch"}},
    )
    print(f"[rsl-ONNX] exported -> {args.output}  (no normalization; mean CENet code)")


if __name__ == "__main__":
    main()