"""The models must actually train, not merely have correct gradients.

"No trainer in the library" is a scope decision. It must not quietly become
"nobody ever checked that it converges" — a block can pass gradcheck and still
never learn. These tests are the guard against that reading.
"""

import warnings

import pytest
import torch
import torch.nn as nn

import torch_dimensions as td


def factory(cls, plan=None):
    def build(lat, d_model):
        with warnings.catch_warnings():
            warnings.simplefilter("ignore")
            return cls(d_model, len(lat.axis_names), lat, plan=plan)

    return build


@pytest.mark.parametrize("cls", [td.LSTM, td.GRU])
def test_the_rnn_family_learns_a_task_that_needs_axial_mixing(cls):
    result = td.testing.check_trainable(factory(cls), steps=150)
    assert result["ratio"] >= 3.0, result
    assert result["held_out"] < result["initial"]


def test_a_model_that_never_sweeps_the_needed_axis_cannot_learn_it():
    """The negative that makes the positive mean something. The task is a
    cumulative sum along `w`; a plan that only sweeps time has no path for
    information to travel along `w` at all."""
    blind = factory(td.LSTM, plan=td.ScanPlan.from_list(["time", "time"]))
    result = td.testing.check_trainable(blind, steps=150, raise_on_failure=False)
    assert result["ratio"] < 3.0, f"blind model should not solve this: {result}"


def test_check_trainable_raises_with_an_actionable_message():
    blind = factory(td.LSTM, plan=td.ScanPlan.from_list(["time", "time"]))
    with pytest.raises(AssertionError, match="did not learn"):
        td.testing.check_trainable(blind, steps=100)


def test_check_trainable_reports_the_numbers_it_judged_on():
    result = td.testing.check_trainable(factory(td.LSTM), steps=50, min_ratio=0.0)
    assert set(result) == {"initial", "final", "held_out", "ratio"}
    assert all(isinstance(v, float) for v in result.values())


def test_a_sparse_lattice_model_still_learns():
    """Masking absent cells must not sever the gradient path for present ones."""
    valid = torch.tensor(
        [[True, True, True, False], [True, True, True, True], [True, False, True, True]]
    )
    lat = td.Lattice(shape=(3, 4), names=("h", "w"), valid=valid, time=True)
    torch.manual_seed(0)
    model, head = td.LSTM(16, 3, lat), nn.Linear(16, 1)
    opt = torch.optim.Adam([*model.parameters(), *head.parameters()], lr=1e-2)
    g = torch.Generator().manual_seed(0)

    def draw():
        x = torch.randn(8, 5, 3, 4, 16, generator=g)
        return x, x[..., :1].cumsum(dim=3) * lat.mask().to(x.dtype)

    first = None
    for _ in range(150):
        x, y = draw()
        loss = (head(model(x)) * lat.mask().to(x.dtype) - y).pow(2).mean()
        first = first if first is not None else loss.item()
        opt.zero_grad()
        loss.backward()
        opt.step()
    assert loss.item() < first / 3, f"{first:.4f} -> {loss.item():.4f}"


def test_training_updates_every_parameter():
    """A parameter that never moves is dead weight the gradient check would
    not catch, since it can receive a gradient of exactly zero forever."""
    lat = td.Lattice(shape=(2, 3), time=True)
    torch.manual_seed(0)
    model = td.LSTM(8, 3, lat)
    before = {n: p.detach().clone() for n, p in model.named_parameters()}
    opt = torch.optim.Adam(model.parameters(), lr=1e-2)
    for _ in range(5):
        loss = model(torch.randn(4, 3, 2, 3, 8)).pow(2).mean()
        opt.zero_grad()
        loss.backward()
        opt.step()
    unmoved = [n for n, p in model.named_parameters() if torch.equal(p, before[n])]
    assert not unmoved, unmoved