import copy,random,unittest
import torch
from torch.nn import functional as F
from lightnest.gradient_interference_experiment import RoutedXorPolicy,freeze,teacher_metrics
from lightnest.xor_patch_experiment import inputs,XorPolicy
from lightnest.multifamily_patch import make_task,teacher_states

class GradientRoutingTests(unittest.TestCase):
    def setUp(self):
        torch.set_num_threads(2);torch.manual_seed(17)
        self.states=teacher_states(make_task(random.Random(31)))
    def test_byte_gradient_isolation(self):
        for detached in [False,True]:
            m=RoutedXorPolicy(detached);c,b,p,y=inputs(self.states);mask=p<256
            _,bits=m(c,b);F.binary_cross_entropy_with_logits(bits[torch.arange(len(p))[mask],p[mask]],y[mask]).backward()
            backbone=sum(float(v.grad.abs().sum()) for k,v in m.named_parameters() if not k.startswith('replacement.') and v.grad is not None)
            self.assertGreater(float(m.replacement.weight.grad.abs().sum()),0)
            if detached:self.assertEqual(backbone,0)
            else:self.assertGreater(backbone,0)
    def test_pointer_gradient_unaffected(self):
        a=RoutedXorPolicy();b=copy.deepcopy(a);b.detach_byte=True;c,x,p,y=inputs(self.states)
        for m in [a,b]:F.cross_entropy(m(c,x)[0],p).backward()
        self.assertGreater(float(a.fusion.weight.grad.abs().sum()),0)
        for (k,v),(l,w) in zip(a.named_parameters(),b.named_parameters()):
            if v.grad is None:self.assertIsNone(w.grad)
            else:self.assertTrue(torch.equal(v.grad,w.grad),k)
    def test_inference_equivalent_to_existing_xor(self):
        a=RoutedXorPolicy();b=XorPolicy(copy.deepcopy(a));c,x,_,_=inputs(self.states)
        # Existing wrapper expects absolute-network logits, so compare with plain PatchNetwork.
        from lightnest.learned_patch import PatchNetwork
        net=PatchNetwork();net.load_state_dict(a.state_dict());b=XorPolicy(net)
        for lhs,rhs in zip(a(c,x),b(c,x)):self.assertTrue(torch.equal(lhs,rhs))
        a.detach_byte=True
        for lhs,rhs in zip(a(c,x),b(c,x)):self.assertTrue(torch.equal(lhs,rhs))
    def test_frozen_disjoint_and_reproducible(self):
        train,suites=freeze();train2,suites2=freeze();self.assertEqual(train,train2);self.assertEqual(suites,suites2)
        self.assertFalse({(t.command,t.binary) for t in train}&{(t.command,t.binary) for ts in suites.values() for t in ts})
if __name__=='__main__':unittest.main()