import numpy as np
import pytest
import torch
from pathlib import Path

FIXTURE = Path(__file__).parent / "fixtures" / "parity"
CKPT = FIXTURE / "ckpt"

pytestmark = pytest.mark.skipif(not CKPT.exists(), reason="parity fixture missing")


@pytest.fixture(scope="module")
def model_and_batch():
    import sys
    sys.path.insert(0, str(Path(__file__).resolve().parent.parent / "scripts"))
    from make_parity_fixture import PAD, build_request
    from groot1.model import GrootDecisionModel, collate
    from groot1.serialize import EncodedQuestion, EncodedState

    model = GrootDecisionModel.from_pretrained(CKPT)
    chunks, eqs = build_request(np.random.default_rng(0), 1024)
    state = EncodedState(chunks=chunks, tokens=0)
    batch = collate([state], eqs, PAD, q_state=[0] * len(eqs))
    return model, batch, [q.n_options for q in eqs]


def test_forward_shapes_and_probabilities(model_and_batch):
    model, batch, ks = model_and_batch
    with torch.inference_mode():
        out = model(batch)
    assert out["probabilities"].shape[0] == len(ks)
    assert out["confidence_logit"].shape == (len(ks),)
    for qi, k in enumerate(ks):
        p = out["probabilities"][qi, :k]
        assert p.sum().item() == pytest.approx(1.0, abs=1e-5)
        assert torch.isfinite(out["confidence_logit"][qi])


def test_outputs_match_reference_npy(model_and_batch):
    model, batch, ks = model_and_batch
    ref = np.load(FIXTURE / "ref.npy")
    nopt, nq = sum(ks), len(ks)
    with torch.inference_mode():
        out = model(batch)
    probs = np.concatenate([out["probabilities"][q, :k].numpy() for q, k in enumerate(ks)])
    conf = out["confidence_logit"].numpy()
    logits = np.concatenate([out["logits"][q, :k].numpy() for q, k in enumerate(ks)])
    got = np.concatenate([probs, conf, logits])
    assert np.abs(got - ref).max() < 1e-5