"""Verify export_onnx.py produces an ONNX file that matches the SB3 policy."""
from pathlib import Path

import numpy as np
import pytest


REPO_ROOT = Path(__file__).resolve().parents[2]
CHECKPOINT = REPO_ROOT / "checkpoints" / "ppo_lander_final.zip"
ONNX_OUT = REPO_ROOT / "models" / "ppo_lander.onnx"


@pytest.fixture(scope="module")
def exported_onnx():
    if not CHECKPOINT.exists():
        pytest.skip(f"no checkpoint at {CHECKPOINT} — run train.py first")
    # export
    import subprocess
    import sys
    training_dir = REPO_ROOT / "Training"
    result = subprocess.run(
        [sys.executable, str(training_dir / "export_onnx.py"),
         "--checkpoint", str(CHECKPOINT), "--out", str(ONNX_OUT)],
        capture_output=True, text=True, cwd=str(REPO_ROOT),
    )
    assert result.returncode == 0, f"export failed: {result.stderr}"
    assert ONNX_OUT.exists(), f"ONNX file not produced at {ONNX_OUT}"
    return ONNX_OUT


def test_onnx_file_exists(exported_onnx):
    assert exported_onnx.stat().st_size > 0


def test_onnx_input_output_shapes(exported_onnx):
    import onnx
    model = onnx.load(str(exported_onnx))
    assert len(model.graph.input) == 1
    in_shape = [d.dim_value for d in model.graph.input[0].type.tensor_type.shape.dim]
    # (batch, 7)
    assert in_shape[-1] == 7
    assert len(model.graph.output) >= 1
    out_shape = [d.dim_value for d in model.graph.output[0].type.tensor_type.shape.dim]
    assert out_shape[-1] == 4  # 4 discrete actions


def test_onnx_output_matches_sb3(exported_onnx):
    """Sanity: ONNX inference for 100 random obs must argmax to the same action as SB3."""
    import onnxruntime
    from stable_baselines3 import PPO

    model = PPO.load(str(CHECKPOINT))
    session = onnxruntime.InferenceSession(str(exported_onnx))
    input_name = session.get_inputs()[0].name

    rng = np.random.default_rng(1234)
    obs_batch = rng.normal(size=(100, 7)).astype(np.float32)

    # ONNX: (100, 4) logits
    onnx_logits = session.run(None, {input_name: obs_batch})[0]
    onnx_actions = onnx_logits.argmax(axis=-1)

    # SB3 deterministic prediction
    sb3_actions, _ = model.predict(obs_batch, deterministic=True)

    mismatches = int((onnx_actions != sb3_actions).sum())
    # Allow up to 2 mismatches out of 100 for float rounding on marginal states.
    assert mismatches <= 2, f"{mismatches}/100 argmax mismatches between ONNX and SB3"