import os, sys
sys.path.insert(0, os.path.join(os.path.dirname(__file__), "..", "src"))
import torch
import torch.nn as nn
from zbp_scaling.model import ScalingLM
from zbp_scaling.zbp import ZBPConfig, zbp_blocks


def test_exact_mode_matches_bp():
    torch.manual_seed(0)
    m = ScalingLM(101, 64, 2, 4, 32, cfg=ZBPConfig(mode="bp"))
    x = torch.randint(0, 101, (2, 32)); y = torch.randint(0, 101, (2, 32))
    g_bp = torch.autograd.grad(nn.functional.cross_entropy(m(x).flatten(0, 1), y.flatten()), list(m.parameters()))
    for b in zbp_blocks(m):
        b.cfg = ZBPConfig(mode="exact")
    g = torch.autograd.grad(nn.functional.cross_entropy(m(x).flatten(0, 1), y.flatten()), list(m.parameters()))
    err = max((a - b).norm().item() / (b.norm().item() + 1e-12) for a, b in zip(g, g_bp))
    assert err < 1e-5, err


def test_zbp_cd_trains_shape():
    torch.manual_seed(0)
    m = ScalingLM(101, 64, 2, 4, 32, cfg=ZBPConfig(mode="cd", n_probes=4, eps=0.1, probe_chunk=4))
    x = torch.randint(0, 101, (2, 32)); y = torch.randint(0, 101, (2, 32))
    nn.functional.cross_entropy(m(x).flatten(0, 1), y.flatten()).backward()
    assert all(p.grad is not None and torch.isfinite(p.grad).all() for p in m.parameters())