"""Numerical, shape, gradient, masking, and causality checks for the course."""
import numpy as np
import torch
import numpy_core as n
import torch_core as t
from train_copy import dataset


def check() -> None:
    torch.set_num_threads(1)
    rng = np.random.default_rng(19)
    q, k, v = rng.normal(size=(2, 3, 4, 8)), rng.normal(size=(2, 3, 6, 8)), rng.normal(size=(2, 3, 6, 5))
    visible = np.ones((2, 1, 4, 6), dtype=bool)
    visible[..., -1] = False
    result, weights = n.attention(q, k, v, visible)
    tensors = [torch.tensor(x, requires_grad=True) for x in (q, k, v)]
    torch_result, torch_weights = t.attention(tensors[0], tensors[1], tensors[2], torch.tensor(visible))
    np.testing.assert_allclose(result, torch_result.detach().numpy(), atol=1e-12)
    np.testing.assert_allclose(weights.sum(-1), 1.0)
    assert np.all(weights[..., -1] == 0)
    changed_v = v.copy(); changed_v[..., -1, :] += 10000
    np.testing.assert_allclose(n.attention(q, k, changed_v, visible)[0], result)
    torch_result.square().sum().backward()
    assert all(x.grad is not None and torch.isfinite(x.grad).all() for x in tensors)
    for implementation, data, mask in [(n, (q,k,v), np.zeros_like(visible)),
                                     (t, tensors, torch.zeros_like(torch.tensor(visible)))]:
        try:
            implementation.attention(*data, allowed=mask)
        except ValueError:
            pass
        else:
            raise AssertionError("All-masked rows must be rejected")

    np.testing.assert_allclose(n.positions(9,8), t.positions(9,8,dtype=torch.float64).numpy(), atol=1e-12)
    np.testing.assert_allclose(n.embedding_demo()[2], t.embedding_demo()[2].detach().numpy())
    x = rng.normal(size=(2,4,8))
    np.testing.assert_allclose(n.layer_norm(x), t.layer_norm(torch.tensor(x)).numpy(), atol=1e-12)
    mha_n, mha_t = n.MultiHead(8,2,seed=3), t.MultiHead(8,2,seed=3).double()
    # Copy full precision arrays; converting previously rounded float32 weights to double cannot undo rounding.
    with torch.no_grad():
        for name in ['wq','wk','wv','wo']:
            getattr(mha_t,name).weight.copy_(torch.tensor(getattr(mha_n,name).T))
    np.testing.assert_allclose(mha_n(x,x)[0], mha_t(torch.tensor(x),torch.tensor(x))[0].detach().numpy(), atol=1e-12)
    permutation = np.array([2,0,3,1])
    np.testing.assert_allclose(mha_n(x[:,permutation],x[:,permutation])[0],mha_n(x,x)[0][:,permutation],atol=1e-12)

    source = torch.tensor([[4,7,0],[9,4,6]])
    prefix = torch.tensor([[1,4,7,0],[1,9,4,6]])
    model = t.TinyTransformer().eval()
    with torch.no_grad():
        baseline = model(source,prefix)
        changed = prefix.clone(); changed[:,3] = 11
        torch.testing.assert_close(model(source,changed)[:,:3],baseline[:,:3])
        padded_source = torch.cat([source,torch.zeros(2,2,dtype=torch.long)],dim=1)
        torch.testing.assert_close(model(padded_source,prefix),baseline,atol=2e-6,rtol=2e-6)
        memory, mask = model.encode(source)
        for i in range(prefix.shape[1]):
            incremental = model.decode(prefix[:,:i+1],memory,mask)[:,-1]
            torch.testing.assert_close(incremental,baseline[:,i],atol=2e-6,rtol=2e-6)
    train, held_out = dataset()
    assert not set(map(tuple,train)) & set(map(tuple,held_out))
    assert n.demo_forward().shape == (2,3,12)
    print("PASS: NumPy/PyTorch agreement, normalized weights, masked-value invariance, gradients,")
    print("positions, embeddings, normalization, heads, permutation equivariance, decoder causality,")
    print("padding invariance, incremental-prefix equivalence, and disjoint evaluation sequences.")


if __name__ == "__main__":
    check()