"""Snapshot tests for DiscoMetaNet forward pass — no JAX or weights required.

Uses deterministic seeding with random weights to verify numerical stability
and catch regressions. Expected values were generated by running this test
suite once and recording the outputs.
"""

import torch
import pytest

from disco_torch.meta_net import DiscoMetaNet
from disco_torch.types import UpdateRuleInputs, MetaNetInputOption, TransformConfig
from disco_torch.update_rule import get_input_option


def _make_test_inputs(T=3, B=2, A=4, PS=600):
    """Create deterministic test inputs matching the full Disco103 input spec."""
    torch.manual_seed(123)

    logits = torch.randn(T + 1, B, A)
    beh_logits = torch.randn(T + 1, B, A)
    y = torch.randn(T + 1, B, PS)
    z = torch.randn(T + 1, B, A, PS)
    actions = torch.randint(0, A, (T + 1, B))
    rewards = torch.randn(T, B)
    is_terminal = torch.zeros(T, B)

    tgt_logits = torch.randn(T + 1, B, A)
    tgt_y = torch.randn(T + 1, B, PS)
    tgt_z = torch.randn(T + 1, B, A, PS)
    v_scalar = torch.randn(T + 1, B)
    adv = torch.randn(T, B)
    norm_adv = torch.randn(T, B)
    q = torch.randn(T + 1, B, A)
    qv_adv = torch.randn(T + 1, B, A)
    norm_qv_adv = torch.randn(T + 1, B, A)

    inputs = UpdateRuleInputs(
        observations=torch.zeros(T + 1, B, 8),
        actions=actions,
        rewards=rewards,
        is_terminal=is_terminal,
        agent_out={"logits": logits, "y": y, "z": z},
        behaviour_agent_out={"logits": beh_logits, "y": y.clone(), "z": z.clone()},
    )
    inputs.extra_from_rule = {
        "v_scalar": v_scalar,
        "adv": adv,
        "normalized_adv": norm_adv,
        "q": q,
        "qv_adv": qv_adv,
        "normalized_qv_adv": norm_qv_adv,
        "target_out": {
            "logits": tgt_logits,
            "y": tgt_y,
            "z": tgt_z,
        },
    }
    return inputs


class TestDiscoMetaNetShapes:
    def test_output_shapes(self):
        T, B, A, PS = 3, 2, 4, 600
        torch.manual_seed(42)
        net = DiscoMetaNet(input_option=get_input_option(), prediction_size=PS)
        inputs = _make_test_inputs(T, B, A, PS)
        state = net.initial_meta_rnn_state()

        with torch.no_grad():
            meta_out, new_state = net(inputs, state)

        assert meta_out["pi"].shape == (T, B, A)
        assert meta_out["y"].shape == (T, B, PS)
        assert meta_out["z"].shape == (T, B, PS)
        assert meta_out["meta_input_emb"].shape == (T, B, 1)

        new_h, new_c = new_state
        assert new_h.shape == (128,)
        assert new_c.shape == (128,)

    def test_outputs_are_finite(self):
        T, B, A, PS = 3, 2, 4, 600
        torch.manual_seed(42)
        net = DiscoMetaNet(input_option=get_input_option(), prediction_size=PS)
        inputs = _make_test_inputs(T, B, A, PS)
        state = net.initial_meta_rnn_state()

        with torch.no_grad():
            meta_out, new_state = net(inputs, state)

        for k in ("pi", "y", "z", "meta_input_emb"):
            assert torch.isfinite(meta_out[k]).all(), f"{k} has non-finite values"
        assert torch.isfinite(new_state[0]).all()
        assert torch.isfinite(new_state[1]).all()


class TestDiscoMetaNetSnapshot:
    """Snapshot test: verify exact numerical output with seeded random weights.

    If you intentionally change the network architecture or computation,
    update these expected values by running:
        pytest tests/test_meta_net.py -k snapshot -s
    and copying the printed values.
    """

    def test_snapshot(self):
        T, B, A, PS = 3, 2, 4, 600
        torch.manual_seed(42)
        net = DiscoMetaNet(input_option=get_input_option(), prediction_size=PS)
        inputs = _make_test_inputs(T, B, A, PS)
        state = net.initial_meta_rnn_state()

        with torch.no_grad():
            meta_out, new_state = net(inputs, state)

        # Print actuals for updating snapshot values
        actuals = {
            "pi_mean": meta_out["pi"].mean().item(),
            "y_mean": meta_out["y"].mean().item(),
            "z_mean": meta_out["z"].mean().item(),
            "meta_input_emb_mean": meta_out["meta_input_emb"].mean().item(),
            "new_h_mean": new_state[0].mean().item(),
            "new_c_mean": new_state[1].mean().item(),
            "pi_std": meta_out["pi"].std().item(),
            "y_std": meta_out["y"].std().item(),
        }
        print("\nActual snapshot values:")
        for k, v in actuals.items():
            print(f"  {k}: {v:.10f}")

        # ---- Snapshot expected values (update these if architecture changes) ----
        expected = {
            "pi_mean": 0.2151360512,
            "y_mean": -0.0012468175,
            "z_mean": -0.0000647042,
            "meta_input_emb_mean": 0.0425805487,
            "new_h_mean": -0.0020979117,
            "new_c_mean": -0.0042600036,
            "pi_std": 0.0032218690,
            "y_std": 0.0352731943,
        }
        for k, exp in expected.items():
            assert abs(actuals[k] - exp) < 1e-5, f"{k}: {actuals[k]} != {exp}"
