from __future__ import annotations

import pytest
import torch

from mouse_control import SimpleNet


@pytest.fixture
def model() -> SimpleNet:
    return SimpleNet()


@pytest.mark.parametrize("batch_size", [1, 4, 16, 64])
def test_forward_with_unsqueezed_input(model: SimpleNet, batch_size: int) -> None:
    out = model(torch.randn(batch_size, 1, 2))
    assert out.shape == (batch_size, 10, 2)


@pytest.mark.parametrize("batch_size", [1, 8, 32])
def test_forward_with_flat_input(model: SimpleNet, batch_size: int) -> None:
    """SimpleNet's first layer flattens, so (N, 2) is also valid."""
    out = model(torch.randn(batch_size, 2))
    assert out.shape == (batch_size, 10, 2)


def test_layers_have_expected_dimensions(model: SimpleNet) -> None:
    assert model.fc1.in_features == 2
    assert model.fc1.out_features == 64
    assert model.fc2.in_features == 64
    assert model.fc2.out_features == 32
    assert model.fc3.in_features == 32
    assert model.fc3.out_features == 20


def test_model_has_trainable_parameters(model: SimpleNet) -> None:
    params = list(model.parameters())
    assert params, "SimpleNet should expose parameters"
    assert all(p.requires_grad for p in params)


def test_forward_is_differentiable(model: SimpleNet) -> None:
    inp = torch.randn(2, 1, 2, requires_grad=True)
    out = model(inp)
    out.sum().backward()
    assert inp.grad is not None
    assert inp.grad.shape == inp.shape


def test_onnx_export_round_trips(tmp_path, model: SimpleNet) -> None:
    """Trained or not, the model architecture must round-trip via ONNX."""
    import numpy as np
    import onnxruntime as ort

    onnx_path = tmp_path / "round_trip.onnx"
    dummy = torch.randn(1, 1, 2)
    model.eval()
    torch.onnx.export(
        model, dummy, str(onnx_path),
        input_names=["input"], output_names=["output"],
        external_data=False,
    )
    assert onnx_path.exists()
    assert not onnx_path.with_suffix(onnx_path.suffix + ".data").exists(), \
        "external_data=False must produce a single-file ONNX"

    session = ort.InferenceSession(str(onnx_path))
    inp = np.array([[[42.0, -17.0]]], dtype=np.float32)
    out = session.run(None, {"input": inp})[0]
    assert out.shape == (1, 10, 2)

    with torch.no_grad():
        torch_out = model(torch.from_numpy(inp)).numpy()
    np.testing.assert_allclose(out, torch_out, atol=1e-4)