"""Export a trained SB3 checkpoint to the JSON format PolicyNetwork.gd loads.

The exported file contains the deterministic policy MLP (obs -> action means);
the game clamps outputs to [-1, 1] and treats the last value as turbo (> 0).
A parity self-check compares the JSON forward pass against SB3's own
deterministic prediction before writing.

Example:
    .venv/bin/python export_policy.py checkpoints/smoke/final.zip ../Game/bots/rookie.json
"""

import argparse
import json
import pathlib

import numpy as np
import torch
from stable_baselines3 import PPO


def linear_to_layer(linear: torch.nn.Linear, activation: str) -> dict:
    return {
        "weights": linear.weight.detach().cpu().numpy().tolist(),
        "biases": linear.bias.detach().cpu().numpy().tolist(),
        "activation": activation,
    }


def extract_layers(policy) -> list[dict]:
    # Features extractor must be a passthrough (flatten) for this export to
    # be faithful; it has no parameters for our flat "obs" Box space.
    n_extractor_params = sum(p.numel() for p in policy.features_extractor.parameters())
    assert n_extractor_params == 0, "features extractor has weights; export logic needs updating"

    layers = []
    modules = list(policy.mlp_extractor.policy_net)
    for i, module in enumerate(modules):
        if isinstance(module, torch.nn.Linear):
            next_is_tanh = i + 1 < len(modules) and isinstance(modules[i + 1], torch.nn.Tanh)
            assert next_is_tanh or i + 1 >= len(modules), (
                f"unsupported activation after layer {i}: {modules[i + 1] if i + 1 < len(modules) else None}"
            )
            layers.append(linear_to_layer(module, "tanh" if next_is_tanh else "linear"))
        elif not isinstance(module, torch.nn.Tanh):
            raise AssertionError(f"unsupported module in policy net: {module}")
    layers.append(linear_to_layer(policy.action_net, "linear"))
    return layers


def json_forward(layers: list[dict], obs: np.ndarray) -> np.ndarray:
    x = obs
    for layer in layers:
        x = np.asarray(layer["weights"]) @ x + np.asarray(layer["biases"])
        if layer["activation"] == "tanh":
            x = np.tanh(x)
    return x


def main():
    parser = argparse.ArgumentParser(description=__doc__)
    parser.add_argument("checkpoint", help="SB3 checkpoint .zip (e.g. checkpoints/smoke/final.zip)")
    parser.add_argument("output", help="Output JSON path (e.g. ../Game/bots/rookie.json)")
    args = parser.parse_args()

    model = PPO.load(args.checkpoint, device="cpu")
    policy = model.policy
    layers = extract_layers(policy)
    input_size = model.observation_space["obs"].shape[0]

    # Parity check: JSON forward pass must match SB3's deterministic action
    rng = np.random.default_rng(0)
    for _ in range(16):
        obs = rng.uniform(-1, 1, input_size).astype(np.float32)
        expected, _ = model.predict({"obs": obs}, deterministic=True)
        actual = np.clip(json_forward(layers, obs), -1.0, 1.0)
        assert np.allclose(actual, expected, atol=1e-5), f"parity check failed: {actual} vs {expected}"

    output = pathlib.Path(args.output)
    output.parent.mkdir(parents=True, exist_ok=True)
    with open(output, "w") as f:
        json.dump({"input_size": int(input_size), "layers": layers}, f)
    print(f"Exported {args.checkpoint} -> {output} (input size {input_size}, {len(layers)} layers, parity OK)")


if __name__ == "__main__":
    main()