"""Remote-only independent checks for matching, head decomposition and interventions."""
import itertools
import json
from pathlib import Path
import struct
import tempfile
import unittest
import numpy as np
from crossqkv.core import (QKVModel,InterventionLM,assignment,inverse,distribution_metrics,
                           output_metrics,cross_head_output_metrics,value_gram_cka)
from functional.model import StreamingLM


def fixture(directory):
    rng=np.random.default_rng(31);arrays={}
    cfg={'hidden_size':16,'head_dim':4,'num_attention_heads':4,'num_key_value_heads':2,'rms_norm_eps':1e-6,
         'rope_parameters':{'rope_theta':10000000.,'partial_rotary_factor':.5}}
    for layer in (0,1):
        p=f'model.language_model.layers.{layer}.self_attn.'
        for role,shape in [('q_proj',(32,16)),('k_proj',(8,16)),('v_proj',(8,16)),('o_proj',(16,16)),('q_norm',(4,)),('k_norm',(4,))]:
            arrays[p+role+'.weight']=(rng.standard_normal(shape)*.15).astype(np.float32)
    header={};offset=0;body=b''
    for name,arr in arrays.items():
        raw=arr.tobytes();header[name]={'shape':list(arr.shape),'dtype':'F32','data_offsets':[offset,offset+len(raw)]};offset+=len(raw);body+=raw
    h=json.dumps(header).encode();Path(directory,'model.safetensors').write_bytes(struct.pack('<Q',len(h))+h+body)
    return cfg,rng.standard_normal((2,9,16),dtype=np.float32)


class CrossQKVTests(unittest.TestCase):
    def test_assignment_against_exhaustive_and_inverse(self):
        rng=np.random.default_rng(11);cost=rng.uniform(size=(4,4))
        p=assignment(cost)
        best=min((sum(cost[i,j] for i,j in enumerate(perm)),perm) for perm in itertools.permutations(range(4)))
        self.assertEqual(tuple(p),best[1]);self.assertEqual([p[inverse(p)[i]] for i in range(4)],list(range(4)))

    def test_probability_identity_and_nontrivial_queries(self):
        rng=np.random.default_rng(12);p=np.zeros((2,3,9,9),dtype=np.float32)
        for t in range(9):
            x=rng.uniform(size=(2,3,t+1));p[:,:,t,:t+1]=x/x.sum(-1,keepdims=True)
        metrics=distribution_metrics(p,p)
        np.testing.assert_allclose(metrics['js'],0,atol=1e-14);np.testing.assert_array_equal(metrics['top5'],np.ones(2))
        changed=p.copy();changed[:,:,:7]=0
        # Primary query rows do not include the inevitable first row or t<7.
        np.testing.assert_allclose(distribution_metrics(p,changed)['js'],0,atol=1e-14)

    def test_profile_matches_canonical_and_head_sum(self):
        with tempfile.TemporaryDirectory() as d:
            cfg,x=fixture(d);lm=QKVModel(d,cfg,enforce_memory=False)
            for layer in (0,1):
                parts=lm.parts(x,layer);expected,p=StreamingLM.full_attention(lm,x,layer)
                np.testing.assert_allclose(parts['output'],expected,atol=2e-6,rtol=2e-6)
                np.testing.assert_allclose(parts['p'],p,atol=2e-7)
                np.testing.assert_allclose(parts['heads'].sum(2),expected,atol=2e-6,rtol=2e-6)
                np.testing.assert_allclose(lm.transplant(parts,parts,layer,'qk_map',perm=list(range(4))),expected,atol=2e-6)
                np.testing.assert_allclose(lm.transplant(parts,parts,layer,'qk_head',1,1),expected,atol=2e-6)
                np.testing.assert_allclose(lm.transplant(parts,parts,layer,'v_group',1,1),expected,atol=2e-6)
                np.testing.assert_allclose(lm.transplant(parts,parts,layer,'qkv_head',1,1),expected,atol=2e-6)

    def test_recipient_value_gate_output_and_gqa_group_stay_fixed(self):
        with tempfile.TemporaryDirectory() as d:
            cfg,x=fixture(d);lm=QKVModel(d,cfg,enforce_memory=False);a,b=lm.parts(x,0),lm.parts(x,1)
            v=a['v'].copy();v[:,:,1]=b['v'][:,:,0]
            context=np.einsum('bhts,bshd->bthd',a['p'],np.repeat(v,2,axis=2))
            expected=lm.project(context,a['gate'],0)
            np.testing.assert_allclose(lm.transplant(a,b,0,'v_group',1,0),expected,atol=2e-6)
            p=a['p'].copy();p[:,2]=b['p'][:,1]
            context=np.einsum('bhts,bshd->bthd',p,np.repeat(a['v'],2,axis=2))
            np.testing.assert_allclose(lm.transplant(a,b,0,'qk_head',2,1),lm.project(context,a['gate'],0),atol=2e-6)

    def test_head_matrix_metrics_and_gram_are_diagnostic(self):
        rng=np.random.default_rng(13);x=rng.standard_normal((2,9,3,6),dtype=np.float32)
        y=rng.standard_normal((2,9,2,6),dtype=np.float32);m=cross_head_output_metrics(x,y)
        for a in range(3):
            for b in range(2):
                ref=output_metrics(x[:,:,a],y[:,:,b])
                for key in ref:np.testing.assert_allclose(m[key][:,a,b],ref[key],atol=1e-12)
        v=rng.standard_normal((2,9,6),dtype=np.float32);r=np.eye(6,dtype=np.float32)[::-1]
        np.testing.assert_allclose(value_gram_cka(v,v@r),np.ones(2),atol=1e-12)
        self.assertGreater(float(np.linalg.norm(v-v@r)),1.)

    def test_cached_mode_rejects_future_state_reuse(self):
        with tempfile.TemporaryDirectory() as d:
            cfg,x=fixture(d);lm=InterventionLM(d,cfg,enforce_memory=False)
            c={'i':0,'j':1,'kind':'qk_head','a':2,'b':1,'perm':None}
            with self.assertRaises(AssertionError):lm.configure(c,reverse=False,cached=True)
            lm.configure(c,reverse=True,cached=True)
            lm.full_attention(x,0);out,_=lm.full_attention(x,1)
            own=lm.parts(x,1);expected=lm.transplant(own,lm.donor_cache,1,'qk_head',1,2)
            np.testing.assert_allclose(out,expected,atol=2e-6)


    def test_controls_cover_multiple_projection_layers(self):
        from crossqkv.run import controls_for_profiles
        rng=np.random.default_rng(20)
        h=rng.standard_normal((2,9,16),dtype=np.float32)
        probs=np.zeros((2,4,9,9),dtype=np.float32)
        for t in range(9):probs[:,:,t,:t+1]=1/(t+1)
        profiles={layer:{'p':probs.copy(),'v':rng.standard_normal((2,9,2,4),dtype=np.float32)} for layer in (0,1)}
        rows=[];controls_for_profiles(profiles,0,[{'id':'one'},{'id':'two'}],rows,h)
        self.assertEqual({r['projection_layer'] for r in rows if 'projection_layer' in r},{0,1})
        self.assertTrue(any(r['control']=='value_gram_random_projection_pair' for r in rows))
        for r in rows:
            for value in r.values():
                if isinstance(value,float):self.assertTrue(np.isfinite(value))


    def test_topk_tail_rank_can_change_with_tiny_probability_difference(self):
        from crossqkv.conditional import tv_metrics
        a=np.zeros((1,1,9,9),dtype=np.float64);b=a.copy()
        a[0,0,0,0]=b[0,0,0,0]=1.
        for t in range(1,9):
            delta=np.linspace(-1e-6,1e-6,t)
            a[0,0,t,:t]=.04/t+delta;b[0,0,t,:t]=.04/t-delta
            a[0,0,t,t]=b[0,0,t,t]=.96
        metric=distribution_metrics(a,b);tv=tv_metrics(a,b)
        self.assertLess(float(metric['js'][0]),1e-6)
        self.assertLess(float(tv['total_variation'][0]),1e-4)
        self.assertLess(float(metric['top5'][0]),.8)


if __name__=='__main__':unittest.main()