import torch
import torch.nn.functional as F

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


def _batch(cfg, inputs=(2, 7, 42, 99)):
    ex = build_examples(list(inputs), cfg)
    return make_batch(ex, cfg)


def test_rnn_shapes_and_halting_range():
    cfg = Config(model="rnn", max_steps=20, min_steps=2)
    m = build_model(cfg)
    b = _batch(cfg)
    out = m(b["x"], b["y_in"])
    assert out["logits"].shape == (4, b["y"].shape[1], cfg.vocab)
    s = out["halt_steps"]
    assert s.shape == (4,)
    assert bool((s >= cfg.min_steps).all())
    assert bool((s <= cfg.max_steps).all())


def test_rnn_fixed_k_ablation():
    cfg = Config(model="rnn", halting=False, max_steps=20)
    m = build_model(cfg)
    b = _batch(cfg)
    out = m(b["x"], b["y_in"])
    assert torch.allclose(out["halt_steps"], torch.full((4,), 20.0))


def test_transformer_shapes():
    cfg = Config(model="transformer")
    m = build_model(cfg)
    b = _batch(cfg)
    out = m(b["x"], b["y_in"])
    assert out["logits"].shape == (4, b["y"].shape[1], cfg.vocab)
    assert out["halt_steps"] is None


def test_transformer_integers_mode():
    cfg = Config(model="transformer", vocab_mode="integers")
    m = build_model(cfg)
    b = _batch(cfg)
    out = m(b["x"], b["y_in"])
    assert out["logits"].shape == (4, b["y"].shape[1], cfg.vocab)


def test_rnn_fully_tied_no_gru_decoder():
    """Regression: design review fix — everything recurrent must be the ONE tied cell."""
    import torch.nn as nn
    cfg = Config(model="rnn")
    m = build_model(cfg)
    assert not any(isinstance(mod, nn.GRUCell) for mod in m.modules())
    assert not any(isinstance(mod, (nn.GRU, nn.LSTM, nn.RNN)) for mod in m.modules())


def test_param_counts_logged_not_gated():
    r = build_model(Config(model="rnn")).param_count()
    t = build_model(Config(model="transformer")).param_count()
    assert r > 1000 and t > 1000


def test_forward_backward_no_nan_both_models():
    for model_name in ("rnn", "transformer"):
        cfg = Config(model=model_name)
        m = build_model(cfg)
        b = _batch(cfg)
        out = m(b["x"], b["y_in"])
        y_safe = b["y"].clamp(max=cfg.vocab - 1)
        loss = F.cross_entropy(out["logits"].reshape(-1, cfg.vocab), y_safe.reshape(-1), reduction="none")
        mask = b["y_mask"].float().reshape(-1)
        loss = (loss * mask).sum() / mask.sum()
        loss.backward()
        assert torch.isfinite(loss).item(), f"{model_name} loss NaN"


def test_greedy_decode_shape():
    from src.model_api import greedy_decode
    cfg = Config(model="rnn")
    m = build_model(cfg)
    x = torch.tensor([[4, 2], [9, 9]], dtype=torch.long)
    gen = greedy_decode(m, x, cfg)
    assert gen.shape == (2, cfg.max_out_len)