"""Export a trained SB3 PPO policy to ONNX for the .NET backend.

The exported model has a single float input of shape [batch, 7] (the observation
vector) and a single float output of shape [batch, 4] (action logits). Argmax
over the last axis gives the action.
"""
from __future__ import annotations

import argparse
from pathlib import Path

import torch
from stable_baselines3 import PPO


class OnnxablePolicy(torch.nn.Module):
    """Wraps the SB3 policy's actor path for ONNX export.

    Standard SB3 MlpPolicy for a discrete action space:
        features = policy.extract_features(obs)
        latent_pi = policy.mlp_extractor.forward_actor(features)
        logits = policy.action_net(latent_pi)
    """

    def __init__(self, policy: torch.nn.Module) -> None:
        super().__init__()
        self.policy = policy

    def forward(self, obs: torch.Tensor) -> torch.Tensor:
        features = self.policy.extract_features(obs)
        latent_pi = self.policy.mlp_extractor.forward_actor(features)
        return self.policy.action_net(latent_pi)


def main() -> None:
    parser = argparse.ArgumentParser()
    parser.add_argument("--checkpoint", type=Path, required=True,
                        help="path to SB3 .zip checkpoint")
    parser.add_argument("--out", type=Path, required=True,
                        help="destination .onnx path")
    parser.add_argument("--opset", type=int, default=17)
    args = parser.parse_args()

    args.out.parent.mkdir(parents=True, exist_ok=True)

    model = PPO.load(str(args.checkpoint), device="cpu")
    model.policy.eval()

    onnxable = OnnxablePolicy(model.policy)
    onnxable.eval()

    dummy = torch.randn(1, 7, dtype=torch.float32)
    torch.onnx.export(
        onnxable,
        dummy,
        str(args.out),
        input_names=["obs"],
        output_names=["logits"],
        dynamic_axes={"obs": {0: "batch"}, "logits": {0: "batch"}},
        opset_version=args.opset,
        dynamo=False,
    )
    print(f"exported ONNX policy to {args.out}")


if __name__ == "__main__":
    main()