#!/usr/bin/env python3
"""Export DreamWaQ CENetActorModel to ONNX for MuJoCo sim2sim.

输入: obs(45) + obs_history(225) → 输出: action(12)（确定性，mean CENet code）

用法:
    uv run scripts/export_dreamwaq_onnx_new.py
    uv run scripts/export_dreamwaq_onnx_new.py --checkpoint runs/.../model_2850.pt
"""
import argparse, os, sys, glob
import torch

sys.path.insert(0, os.path.dirname(os.path.dirname(os.path.abspath(__file__))))
from motrix_rl.rslrl.torch.models.cenet_actor import CENetActorModel
from tensordict import TensorDict

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


class DwaqInfer(torch.nn.Module):
    """确定性推理：VAE mean code + Actor MLP。"""
    def __init__(self, model: CENetActorModel):
        super().__init__()
        self.vae = model.vae
        self.actor_mlp = model.mlp

    def forward(self, obs, obs_history):
        """obs: (N,45), obs_history: (N,225) → action: (N,12)"""
        code = self.vae.deterministic_code(obs_history)  # (N,19)
        latent = torch.cat([code, obs], dim=-1)  # (N,64)
        return self.actor_mlp(latent)


def main():
    p = argparse.ArgumentParser()
    p.add_argument("--checkpoint", default=None)
    p.add_argument("--output", default=None)
    args = p.parse_args()

    if args.checkpoint:
        ckpt = args.checkpoint
    else:
        models = sorted(glob.glob(os.path.join(PROJECT, "runs/go1-dreamwaq-walk/rslrl/*/model_*.pt")))
        if not models:
            print("No checkpoints found"); sys.exit(1)
        ckpt = models[-1]
    print(f"[Export] checkpoint: {ckpt}")

    # 加载模型
    dummy = TensorDict({
        "policy": torch.zeros(1, 45),
        "obs_history": torch.zeros(1, 225),
        "privileged_obs": torch.zeros(1, 247),
    }, batch_size=[1])
    model = CENetActorModel(
        dummy, {"actor": ["policy", "obs_history"]}, "actor", 12,
        hidden_dims=[512, 256, 128], activation="elu", stochastic=True,
        init_noise_std=1.0, cenet_in_dim=225, cenet_out_dim=19,
    )
    data = torch.load(ckpt, map_location="cpu")
    if "actor_state_dict" in data:
        model.load_state_dict(data["actor_state_dict"])
    else:
        model.load_state_dict(data)
    model.eval()

    # 包装为推理模型
    infer = DwaqInfer(model)
    infer.eval()

    # 导出 ONNX
    out_path = args.output or os.path.join(os.path.dirname(ckpt), "policy.onnx")
    dummy_obs = torch.zeros(1, 45)
    dummy_hist = torch.zeros(1, 225)
    torch.onnx.export(
        infer, (dummy_obs, dummy_hist), out_path,
        input_names=["obs", "obs_history"],
        output_names=["actions"],
        opset_version=11,
        dynamic_axes={"obs": {0: "batch"}, "obs_history": {0: "batch"}, "actions": {0: "batch"}},
    )
    print(f"[Export] saved: {out_path}")


if __name__ == "__main__":
    main()