"""
Export SKRL PyTorch policy to ONNX format.

Usage:
    python demo/export_skrl_policy_to_onnx.py --checkpoint <path_to_checkpoint.pt> --output <output_dir>

Example:
    python demo/export_skrl_policy_to_onnx.py \
        --checkpoint runs/go1-rough-terrain-walk/skrl/[time]_PPO/checkpoints/best_agent.pt \
        --output exports_go1_rough
"""

import argparse
import os

import numpy as np
import torch
import torch.nn as nn


# ============== PyTorch Checkpoint Export ==============

class SKRLPolicyTorch(nn.Module):
    """PyTorch policy matching SKRL GaussianMixin architecture.

    The architecture is inferred from the checkpoint keys:
    - net.N: Hidden layers (Linear -> ELU)
    - mean_layer: Output layer (no activation)
    """

    def __init__(self, obs_dim, action_dim, hidden_dims=[512, 256, 128]):
        super().__init__()
        self.obs_dim = obs_dim
        self.action_dim = action_dim
        self.hidden_dims = hidden_dims

        layers = []
        in_dim = obs_dim
        for hidden_dim in hidden_dims:
            layers.extend([
                nn.Linear(in_dim, hidden_dim),
                nn.ELU(),
            ])
            in_dim = hidden_dim

        self.net = nn.Sequential(*layers)
        self.mean_layer = nn.Linear(in_dim, action_dim)
        self.log_std = nn.Parameter(torch.zeros(action_dim))

    def forward(self, x):
        return self.mean_layer(self.net(x))


class ONNXPolicyExporterTorch(nn.Module):
    """Exporter wrapper that applies normalization and policy.

    Matches SKRL's RunningStandardScaler behavior:
    - Normalize: (x - mean) / sqrt(var + eps)
    - Clip to [-5.0, 5.0] (clip_threshold)
    """

    def __init__(self, policy, normalizer_mean, normalizer_std):
        super().__init__()
        self.policy = policy
        # Register normalizer buffers as float32 (matching RunningStandardScaler)
        self.register_buffer('mean', normalizer_mean.float())
        self.register_buffer('std', normalizer_std.float())
        self.clip_threshold = 5.0

    def forward(self, x):
        # Normalize (RunningStandardScaler._compute)
        x = (x - self.mean) / (self.std + 1e-8)
        # Clip to RunningStandardScaler clip_threshold
        x = torch.clamp(x, min=-self.clip_threshold, max=self.clip_threshold)
        return self.policy(x)


def infer_architecture_from_state_dict(state_dict):
    """Infer the policy architecture from state dict keys.

    Args:
        state_dict: Policy state dict

    Returns:
        tuple: (obs_dim, action_dim, hidden_dims)
    """
    # Find observation dimension from first layer
    obs_dim = state_dict['net.0.weight'].shape[1]

    # Find action dimension from mean_layer
    action_dim = state_dict['mean_layer.weight'].shape[0]

    # Infer hidden dimensions from net layers
    hidden_dims = []
    net_keys = sorted([k for k in state_dict.keys() if k.startswith('net.') and k.endswith('.weight')])
    for key in net_keys:
        if 'mean_layer' not in key:  # Skip output layer
            layer_idx = int(key.split('.')[1])
            if layer_idx % 2 == 0:  # Only Linear layers, not activation layers
                hidden_dims.append(state_dict[key].shape[0])

    return obs_dim, action_dim, hidden_dims


def export_torch_policy_to_onnx(checkpoint_path, output_path, verbose=False):
    """Export PyTorch SKRL policy to ONNX.

    Args:
        checkpoint_path: Path to .pt checkpoint file
        output_path: Directory to save ONNX file
        verbose: Whether to print verbose output
    """
    print(f"Loading PyTorch checkpoint from: {checkpoint_path}")
    state = torch.load(checkpoint_path, map_location='cpu', weights_only=False)

    policy_state = state['policy']
    normalizer_state = state['state_preprocessor']

    # Infer architecture
    obs_dim, action_dim, hidden_dims = infer_architecture_from_state_dict(policy_state)
    print(f"Inferred architecture: obs={obs_dim}, action={action_dim}, hidden={hidden_dims}")

    # Create model
    print("Creating PyTorch model...")
    policy = SKRLPolicyTorch(obs_dim, action_dim, hidden_dims)

    # Remove value layers from state dict (keep only policy)
    policy_state_policy_only = {k: v for k, v in policy_state.items()
                                if not k.startswith('value') and 'value_layer' not in k}
    policy.load_state_dict(policy_state_policy_only, strict=False)
    policy.eval()

    # Load normalizer
    mean = normalizer_state['running_mean']
    std = torch.sqrt(normalizer_state['running_variance'])

    # Create exporter
    exporter = ONNXPolicyExporterTorch(policy, mean, std)
    exporter.eval()

    # Create output directory
    os.makedirs(output_path, exist_ok=True)
    output_file = os.path.join(output_path, "policy.onnx")

    # Export to ONNX
    print(f"Exporting policy to: {output_file}")
    obs = torch.zeros(1, obs_dim, dtype=torch.float32)

    torch.onnx.export(
        exporter,
        obs,
        output_file,
        export_params=True,
        opset_version=11,
        verbose=verbose,
        input_names=["observations"],
        output_names=["actions"],
        dynamic_axes={},
    )

    print(f"Successfully exported policy to ONNX: {output_file}")

    # Save PyTorch model and normalizer
    torch.save(policy.state_dict(), os.path.join(output_path, "policy.pt"))
    np.savez(os.path.join(output_path, "normalizer.npz"),
             mean=mean.numpy(), std=std.numpy())

    return output_file


if __name__ == "__main__":
    parser = argparse.ArgumentParser(description="Export SKRL PyTorch policy to ONNX")
    parser.add_argument("--checkpoint", type=str, required=True, help="Path to checkpoint file (.pt)")
    parser.add_argument("--output", type=str, default="exports", help="Output directory")
    parser.add_argument("--verbose", action="store_true", help="Verbose output")

    args = parser.parse_args()

    output_file = export_torch_policy_to_onnx(
        checkpoint_path=args.checkpoint,
        output_path=args.output,
        verbose=args.verbose
    )

    print(f"\nExported to: {output_file}")