"""Minimal tests for the faithful J-lens (jlens_v3) correctness.

Run inside the meru container where torch exists:
    python3 tests/test_jlens_v3.py
Skips gracefully when torch is unavailable.

Key assertion: at the LAST layer the residual-to-residual Jacobian is the
identity, so the faithful J-lens vectors must equal the unembedding rows
(cosine similarity ~1.0). This is the ruler-check that validates the
W_U-probed VJP machinery.
"""
import os
import sys

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

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

from model import GPT, GPTConfig  # noqa: E402

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


def _tiny_model(n_layer=2, d=16, block=32):
    torch.manual_seed(0)
    cfg = GPTConfig(n_layer=n_layer, n_head=2, n_embd=d, block_size=block,
                    bias=False, vocab_size=65, dropout=0.0)
    return GPT(cfg).eval()


def _tiny_data(n=2000, block=32):
    return np.random.RandomState(0).randint(0, 65, n).astype(np.uint16)


def test_shapes_and_finite():
    model = _tiny_model()
    data = _tiny_data()
    batches = jlens_v3.make_batches(data, 32, 4, 3, DEV)
    fv = jlens_v3.compute_faithful_jlens(model, 1, batches, DEV, chunk=8)
    assert fv.shape == (65, 16), f"faithful vecs shape {fv.shape}"
    assert torch.isfinite(fv).all(), "faithful vecs not finite"
    pn = jlens_v3.compute_proxy_norms(model, 1, batches, DEV, 65, chunk=16)
    assert all(np.isfinite(v) for v in pn.values()), "proxy norms not finite"


def test_last_layer_identity():
    """At the final layer J must be identity: faithful vectors == W_U rows."""
    n_layer, d = 2, 16
    model = _tiny_model(n_layer=n_layer, d=d)
    data = _tiny_data()
    batches = jlens_v3.make_batches(data, 32, 4, 3, DEV)
    fv = jlens_v3.compute_faithful_jlens(model, n_layer - 1, batches, DEV,
                                         chunk=8)
    wu = model.lm_head.weight.detach().float()
    sims = torch.nn.functional.cosine_similarity(fv.float(), wu.cpu().float(),
                                                 dim=1)
    assert sims.mean().item() > 0.99, f"identity check failed: {sims.mean():.4f}"


if __name__ == '__main__':
    test_shapes_and_finite()
    test_last_layer_identity()
    print(f"PASS ({DEV})")