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

The exported file contains the deterministic policy MLP. For a MultiDiscrete
(curriculum generation 4+) model, the raw output is 32 per-head logits
decoded via ShipActionCodec.from_logits (argmax per head, mapped through
ACTION_HEADS' bin values below) and an "action_space" block is written to
the JSON so the game knows to decode it that way. For an older continuous
model, output is 7 action means, clamped to [-1, 1] and the last value
treated as turbo (> 0) — no "action_space" block, matching every export
before generation 4 (e.g. Game/bots/promoted/easy.json). A parity self-check
compares the JSON forward pass against SB3's own deterministic prediction
before writing, in either case.

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

import argparse
import json
import pathlib

import gymnasium as gym
import numpy as np
import torch
from stable_baselines3 import PPO

# MUST exactly match Game/scripts/ship_action_codec.gd's HEADS (name, order,
# and bin values) — this is what gets written into every generation-4
# export's "action_space" block, and PolicyNetwork.gd/AIShipController never
# re-derive it, they just decode against whatever's in the file. Sizes are
# cross-checked against the live model's action_space.nvec below (a real
# assertion), but bin *values* have no automated cross-language check —
# treat any edit to either file as requiring the other.
ACTION_HEADS = [
    {"name": "rot_x", "bins": [-1.0, -0.5, 0.0, 0.5, 1.0]},
    {"name": "rot_y", "bins": [-1.0, -0.5, 0.0, 0.5, 1.0]},
    {"name": "rot_z", "bins": [-1.0, -0.5, 0.0, 0.5, 1.0]},
    {"name": "thrust_x", "bins": [-1.0, -0.5, 0.0, 0.5, 1.0]},
    {"name": "thrust_y", "bins": [-0.5, 0.0, 0.45, 0.75, 1.0]},
    {"name": "thrust_z", "bins": [-1.0, -0.5, 0.0, 0.5, 1.0]},
    {"name": "turbo", "bins": [0.0, 1.0]},
]


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]
    is_multi_discrete = isinstance(model.action_space, gym.spaces.MultiDiscrete)

    output_data = {"input_size": int(input_size), "layers": layers}
    rng = np.random.default_rng(0)

    if is_multi_discrete:
        head_sizes = [len(head["bins"]) for head in ACTION_HEADS]
        nvec = [int(n) for n in model.action_space.nvec]
        assert nvec == head_sizes, (
            f"model action_space.nvec {nvec} doesn't match ACTION_HEADS sizes {head_sizes} — "
            "update ACTION_HEADS to match ship_action_codec.gd's HEADS"
        )
        output_data["action_space"] = {"type": "multi_discrete", "heads": ACTION_HEADS}

        # Index-level parity check: deterministic=True now returns one
        # argmax index per head (not a float to clip), so compare argmax of
        # the JSON forward pass's raw logits, sliced per head, against SB3's
        # own chosen indices — a head-order mistake here would otherwise
        # train/export cleanly and only surface as silently wrong in-game
        # behaviour (e.g. pitch commands driving strafe thrusters).
        offsets = []
        running = 0
        for size in head_sizes:
            offsets.append((running, running + size))
            running += size
        for _ in range(16):
            obs = rng.uniform(-1, 1, input_size).astype(np.float32)
            expected, _ = model.predict({"obs": obs}, deterministic=True)
            logits = json_forward(layers, obs)
            actual = np.array([int(np.argmax(logits[start:end])) for start, end in offsets])
            assert np.array_equal(actual, expected), f"parity check failed: {actual} vs {expected}"
    else:
        # Legacy continuous parity check, unchanged: JSON forward pass must
        # match SB3's deterministic action mean.
        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(output_data, f)
    action_space_label = "multi_discrete" if is_multi_discrete else "continuous"
    print(
        f"Exported {args.checkpoint} -> {output} (input size {input_size}, {len(layers)} layers, "
        f"action_space={action_space_label}, parity OK)"
    )


if __name__ == "__main__":
    main()