"""Regression tests for codex-review fixes (design/reviews/codex-review.md)."""
import torch

from src.config import Config
from src.data import build_examples, decode_tokens, make_batch
from src.model_api import build_model, greedy_decode


def test_mixed_batch_vs_singleton_invariance_rnn():
    cfg = Config(model="rnn")
    m = build_model(cfg)
    ex = build_examples([2, 42, 99], cfg)
    batch = make_batch(ex, cfg)
    out_batch = m(batch["x"], batch["y_in"])
    for i in range(3):
        single = make_batch([ex[i]], cfg)
        out_single = m(single["x"], single["y_in"])
        assert torch.allclose(out_batch["logits"][i], out_single["logits"][0], atol=1e-6), f"row {i}"
        assert torch.allclose(out_batch["halt_steps"][i], out_single["halt_steps"][0], atol=1e-6)


def test_mixed_batch_vs_singleton_invariance_transformer():
    cfg = Config(model="transformer")
    m = build_model(cfg)
    ex = build_examples([2, 42, 99], cfg)
    batch = make_batch(ex, cfg)
    out_batch = m(batch["x"], batch["y_in"])
    for i in range(3):
        single = make_batch([ex[i]], cfg)
        out_single = m(single["x"], single["y_in"])
        assert torch.allclose(out_batch["logits"][i], out_single["logits"][0], atol=1e-6), f"row {i}"


def test_greedy_decode_layout_invariant():
    cfg = Config(model="transformer")
    m = build_model(cfg)
    ex = build_examples([42], cfg)
    single = make_batch(ex, cfg)
    out_a = greedy_decode(m, single["x"], cfg)
    raw = torch.tensor([[4, 2]], dtype=torch.long)   # unpadded — greedy must left-pad internally
    out_b = greedy_decode(m, raw, cfg)
    assert torch.equal(out_a, out_b)


def test_min_steps_floor_reachable():
    """min_steps off-by-one fix: earliest halt is step min_steps (forced-1 halt prob -> exactly floor)."""
    cfg = Config(model="rnn", min_steps=2)
    m = build_model(cfg)
    with torch.no_grad():
        m.halt_head.bias.fill_(50.0)   # p -> 1 at the first free step
    ex = build_examples([2, 7, 42], cfg)
    b = make_batch(ex, cfg)
    out = m(b["x"], b["y_in"])
    assert torch.allclose(out["halt_steps"], torch.full_like(out["halt_steps"], float(cfg.min_steps)), atol=1e-3)


def test_pad_logits_cannot_change_loss():
    """Loss and grads must be insensitive to logits at padded target positions."""
    cfg = Config(model="rnn")
    m = build_model(cfg)
    b = make_batch(build_examples([2, 7], cfg), cfg)
    out = m(b["x"], b["y_in"])
    y_safe = b["y"].clamp(max=cfg.vocab - 1)
    loss = torch.nn.functional.cross_entropy(out["logits"].reshape(-1, cfg.vocab),
                                             y_safe.reshape(-1), reduction="none")
    mask = b["y_mask"].float().reshape(-1)
    l1 = (loss * mask).sum() / mask.sum()
    l1.backward()
    grads = {k: v.grad.clone() for k, v in m.named_parameters() if v.grad is not None}
    m.zero_grad()
    out2 = m(b["x"], b["y_in"])
    logits2 = out2["logits"].clone()
    pad_positions = ~b["y_mask"]                       # (B,T_out)
    logits2[pad_positions] += 100.0                    # huge perturbation on pad positions only
    loss2 = torch.nn.functional.cross_entropy(logits2.reshape(-1, cfg.vocab),
                                              y_safe.reshape(-1), reduction="none")
    l2 = (loss2 * mask).sum() / mask.sum()
    assert torch.allclose(l1.detach(), l2.detach())
    l2.backward()
    for k, v in m.named_parameters():
        if v.grad is not None:
            assert torch.allclose(grads[k], v.grad, atol=1e-6), k


def test_integers_vocab_eos_not_aliased():
    """Value 101 (= next_prime(100)) must not collide with EOS in integers mode."""
    cfg = Config(vocab_mode="integers", range_end=100)
    ex = build_examples([100], cfg)
    _, target = ex[0]
    assert target[0] == 101 and target[1] == cfg.eos_id
    assert cfg.eos_id != 101
    assert cfg.vocab == 103
    assert decode_tokens(target, cfg) == 101


def test_act_weights_sum_to_one():
    """ACT mass conservation: steps must always lie in [min_steps, K] for random states."""
    cfg = Config(model="rnn", max_steps=20, min_steps=2)
    m = build_model(cfg)
    b = make_batch(build_examples([2, 7, 42, 99, 3, 15, 60, 97], cfg), cfg)
    out = m(b["x"], b["y_in"])
    s = out["halt_steps"]
    assert bool((s >= cfg.min_steps).all())
    assert bool((s <= cfg.max_steps).all())