"""Minimal tests for the loss-reweighting experiment (no pytest needed).

Run inside the meru container where torch exists:
    python3 tests/test_loss_reweight.py
Skips gracefully when torch is unavailable (e.g. on voidlaptop bare metal).
"""
import os
import sys

sys.path.insert(0, os.path.join(os.path.dirname(__file__), '..', 'src'))

try:
    import numpy as np
    import torch
    import torch.nn.functional as F
    import loss_reweight as lr
except ImportError:
    print("SKIP: torch not available (run inside the meru container)")
    sys.exit(0)

DEV = 'cuda' if torch.cuda.is_available() else 'cpu'


def test_batch_determinism():
    """Paired-model guarantee: identical minibatch order per seed, different across seeds."""
    data = np.random.RandomState(0).randint(0, 65, 50000).astype(np.uint16)
    g1a = torch.Generator().manual_seed(20260731)
    g1b = torch.Generator().manual_seed(20260731)
    x1, y1 = lr._batch(data, 128, 16, g1a, DEV)
    x1b, y1b = lr._batch(data, 128, 16, g1b, DEV)
    assert torch.equal(x1, x1b) and torch.equal(y1, y1b), "same seed must give same batch"
    g2 = torch.Generator().manual_seed(20260732)
    x2, _ = lr._batch(data, 128, 16, g2, DEV)
    assert not torch.equal(x1, x2), "different seeds must give different batches"


def test_weighting_semantics():
    """control == plain CE; ctrl_random picks a fixed set of non-q positions."""
    q_id = 3
    y = torch.randint(0, 65, (2, 4), device=DEV)
    logits = torch.randn(2, 4, 65, device=DEV)
    l_ctl = lr._weighted_loss(logits, y, 'control', q_id, 0, 65, DEV)
    l_ce = F.cross_entropy(logits.view(-1, 65), y.view(-1))
    assert abs(l_ctl.item() - l_ce.item()) < 1e-6, "control must equal plain CE"
    l_r1 = lr._weighted_loss(logits, y, 'ctrl_random', q_id, 7, 65, DEV)
    l_r2 = lr._weighted_loss(logits, y, 'ctrl_random', q_id, 7, 65, DEV)
    assert abs(l_r1.item() - l_r2.item()) < 1e-9, "ctrl_random must be deterministic"


def test_forward_full_logits():
    """Regression: model(X, Y)[0] must be full (B, T, V) logits.
    Karpathy nanoGPT returns LAST-position logits only when targets=None."""
    from model import GPT, GPTConfig
    torch.manual_seed(0)
    cfg = GPTConfig(n_layer=2, n_head=2, n_embd=16, block_size=32, bias=False,
                    vocab_size=65, dropout=0.0)
    m = GPT(cfg).to(DEV)
    x = torch.randint(0, 65, (2, 32), device=DEV)
    y = torch.randint(0, 65, (2, 32), device=DEV)
    logits = m(x, y)[0]
    assert logits.shape == (2, 32, 65), f"full logits expected, got {logits.shape}"


if __name__ == '__main__':
    test_batch_determinism()
    test_weighting_semantics()
    test_forward_full_logits()
    print(f"PASS ({DEV})")