"""Independent algebra, mixed-layer canonical, and joint-mask intervention checks."""
import json,struct,tempfile,unittest
from pathlib import Path
from unittest.mock import patch
import numpy as np
from dominance.core import InfluenceLM,concentration,masks_from_influences,js_rows
from dominance.run import distances,link
from functional.model import StreamingLM,delta_rule,attention
from crossqkv.core import QKVModel


def fixture(path,kinds):
    rng=np.random.default_rng(922);h=16;nh=4;dim=4
    c={'rms_norm_eps':1e-6,'hidden_size':h,'head_dim':dim,'num_attention_heads':nh,'num_key_value_heads':2,'rope_parameters':{'rope_theta':10000000.,'partial_rotary_factor':.5},'linear_num_key_heads':2,'linear_num_value_heads':4,'linear_key_head_dim':4,'linear_value_head_dim':4,'layer_types':kinds}
    a={}
    def add(name,shape):a['model.language_model.'+name]=(rng.standard_normal(shape)*.15).astype(np.float32)
    add('embed_tokens.weight',(19,h));add('norm.weight',(h,))
    for l,kind in enumerate(kinds):
        p=f'layers.{l}.'
        for k in ('input_layernorm','post_attention_layernorm'):add(p+k+'.weight',(h,))
        for k,shape in [('gate_proj',(24,h)),('up_proj',(24,h)),('down_proj',(h,24))]:add(p+'mlp.'+k+'.weight',shape)
        if kind=='full_attention':
            for k,shape in [('q_proj',(32,h)),('k_proj',(8,h)),('v_proj',(8,h)),('o_proj',(h,16)),('q_norm',(4,)),('k_norm',(4,))]:add(p+'self_attn.'+k+'.weight',shape)
        else:
            for k,shape in [('in_proj_qkv',(32,h)),('in_proj_z',(16,h)),('in_proj_a',(4,h)),('in_proj_b',(4,h)),('out_proj',(h,16)),('conv1d',(32,1,4)),('norm',(4,))]:add(p+'linear_attn.'+k+'.weight',shape)
            add(p+'linear_attn.A_log',(4,));add(p+'linear_attn.dt_bias',(4,))
    header={};body=b''
    for k,v in a.items():
        raw=v.tobytes();header[k]={'shape':list(v.shape),'dtype':'F32','data_offsets':[len(body),len(body)+len(raw)]};body+=raw
    j=json.dumps(header).encode();Path(path,'model.safetensors').write_bytes(struct.pack('<Q',len(j))+j+body)
    return c


class DominanceTests(unittest.TestCase):
    def test_continuous_limits_and_undefined_zero(self):
        x=concentration([[1,1,1,1],[4,0,0,0],[0,0,0,0]])
        np.testing.assert_allclose(x['effective_heads'][:2],[4,1]);np.testing.assert_allclose(x['concentration'][:2],[0,1]);self.assertFalse(x['valid'][2]);self.assertTrue(np.isnan(x['effective_heads'][2]))
        self.assertAlmostEqual(float(concentration([3.,1.])['effective_heads']),1.6)
        self.assertAlmostEqual(float(concentration([30.,10.])['effective_heads']),1.6)
        with self.assertRaises(ValueError):concentration([-1,2])
    def test_mask_ranks_vary_by_token_and_recipient(self):
        d=np.arange(4*3*2).reshape(4,3,2).astype(float);d[:,1]=d[::-1,1]
        w=masks_from_influences(d);s=masks_from_influences(d,False)
        np.testing.assert_array_equal(w.sum(-1),2);np.testing.assert_array_equal(w+s,1)
        self.assertFalse(np.array_equal(w[:,0],w[:,1]))
    def test_js_identity_and_disjoint(self):
        np.testing.assert_allclose(js_rows([[1,0]],[[1,0]]),0)
        np.testing.assert_allclose(js_rows([[1,0]],[[0,1]]),np.log(2))
    def test_mixed_layer_forward_and_effective_qkv(self):
        for kinds in [('linear_attention','full_attention'),('full_attention','linear_attention')]:
            with tempfile.TemporaryDirectory() as d:
                cfg=fixture(d,list(kinds));lm=InfluenceLM(d,cfg,enforce_memory=False);ids=np.array([[1,2,3,4,5,6,7,8,9]])
                cache=Path(d)/'cache';cache.mkdir();final=lm.baseline(ids,cache);ref,_=StreamingLM.forward(lm,ids)
                # baseline returns pre-final-norm residual state, canonical returns normalized state.
                np.testing.assert_allclose(lm.normalization(final,'norm'),ref,atol=2e-6,rtol=2e-6)
                for layer in (0,1):
                    raw=np.load(cache/f'layer-{layer}.npz')['raw'];obs=lm.recipient(raw,layer);h=lm.normalization(raw,f'layers.{layer}.input_layernorm')
                    if kinds[layer]=='full_attention':
                        expected=QKVModel(d,cfg,enforce_memory=False).parts(h,layer,heads=False)
                        np.testing.assert_allclose(obs['P'],expected['p'],atol=2e-6)
                        np.testing.assert_allclose(obs['Q'],expected['q'],atol=2e-6)
                    else:
                        with patch('functional.model.delta_rule',wraps=delta_rule) as fn:lm.linear_attention(h,layer)
                        q,k,v=fn.call_args.args[:3]
                        q=q/np.sqrt(np.sum(q*q,-1,keepdims=True)+1e-6)/np.sqrt(cfg['linear_key_head_dim']);k=k/np.sqrt(np.sum(k*k,-1,keepdims=True)+1e-6)
                        for key,val in [('Q',q),('K',k),('V',v)]:np.testing.assert_allclose(obs[key],val,atol=2e-6)
    def test_joint_recomputation_null_and_saved_raw(self):
        with tempfile.TemporaryDirectory() as d:
            cfg=fixture(d,['full_attention','linear_attention']);lm=InfluenceLM(d,cfg,enforce_memory=False);ids=np.array([[1,2,3,4,5,6,7,8,9]])
            cache=Path(d)/'cache';cache.mkdir();lm.baseline(ids,cache)
            data=dict(np.load(cache/'layer-0.npz'));masks=np.ones((1,9,4),dtype=np.float32)
            actual=lm.intervene(data,0,masks);mid=data['raw']+data['out']-data['heads'].sum(2)
            next_raw=mid+lm.mlp(lm.normalization(mid,'layers.0.post_attention_layernorm'),0)
            expected=lm.recipient(next_raw,1)
            for k in expected:np.testing.assert_allclose(actual[k],expected[k],atol=2e-6)
            info=link(lm,data,0,Path(d)/'effects.npz',99);self.assertEqual(info['source_heads'],4);self.assertLess(info['null_max'],1e-6)
            raw=dict(np.load(Path(d)/'effects.npz'));self.assertEqual(raw['single_Q'].shape,(4,9,4));self.assertEqual(raw['joint_weak_half_Q'].shape,(9,4))
            np.testing.assert_array_equal(raw['mask_weak_half']+raw['mask_strong_half'],1)