import torch
from lightnest.global_patch_context import GlobalPatchNetwork,freeze,examples
from lightnest.learned_patch import batch,PatchNetwork
from lightnest.probe_receptive_field import measure

def test_frozen_pairs_and_labels():
    pairs=freeze(12,2,'combined');ex=examples(pairs)
    assert len(ex)==4
    assert all(sum(a!=b for a,b in zip(e.binary,e.target))==1 for e in ex)
    assert freeze(12,2,'combined')==pairs

def test_global_context_breaks_local_invariance():
    torch.set_num_threads(2);torch.manual_seed(82)
    pairs=freeze(12,2,'combined');local=PatchNetwork();global_model=GlobalPatchNetwork()
    global_model.load_state_dict(local.state_dict(),strict=False)
    with torch.no_grad(): global_model.global_scale.fill_(1.)
    assert all(x['pair_difference_change']==0 for x in measure(local,pairs))
    assert any(x['pair_difference_change']>1e-5 for x in measure(global_model,pairs))
    assert not global_model.native_packer_compatible

def test_outputs_and_gradients():
    torch.set_num_threads(2);model=GlobalPatchNetwork();c,b,p,y=batch(examples(freeze(13,1,'train')))
    scores,bits=model(c,b)
    assert scores.shape==(2,257) and bits.shape==(2,256,8)
    torch.nn.functional.cross_entropy(scores,p).backward()
    assert model.global_attention.in_proj_weight.grad.abs().sum()>0
    assert torch.isfinite(scores).all()