"""Independent existing-LoRA arithmetic, head decomposition and runtime checks."""
import json,struct,tempfile,unittest
from pathlib import Path
import numpy as np
from tests.test_dominance import fixture
from dominance.core import InfluenceLM
from finetune_compare.lora import install,EffectiveWeight,ExplicitLoRA,BASE_REVISION
from functional.model import StreamingLM

def adapter(directory,cfg,lm):
    rng=np.random.default_rng(772);arrays={}
    for layer,kind in enumerate(cfg['layer_types']):
        if kind!='linear_attention':continue
        for role in ('in_proj_qkv','in_proj_a','in_proj_b','out_proj'):
            shape=lm.tensors[f'model.language_model.layers.{layer}.linear_attn.{role}.weight'].shape
            prefix=f'base_model.model.model.layers.{layer}.linear_attn.{role}.'
            arrays[prefix+'lora_A.default.weight']=rng.standard_normal((2,shape[1]),dtype=np.float32)*.1
            arrays[prefix+'lora_B.default.weight']=rng.standard_normal((shape[0],2),dtype=np.float32)*.1
    header={};body=b''
    for name,x in arrays.items():
        raw=x.tobytes();header[name]={'shape':list(x.shape),'dtype':'F32','data_offsets':[len(body),len(body)+len(raw)]};body+=raw
    h=json.dumps(header).encode();Path(directory,'adapter.safetensors').write_bytes(struct.pack('<Q',len(h))+h+body)
    return {'base_model':'Qwen/Qwen3.5-2B','base_revision':BASE_REVISION,'arm':'lora','gate':None,'lora_rank':2,'lora_alpha':4,'n_tensors':len(arrays)}

class FineTuneTests(unittest.TestCase):
    def test_effective_rows_equal_explicit_low_rank_output(self):
        with tempfile.TemporaryDirectory() as base,tempfile.TemporaryDirectory() as ad:
            cfg=fixture(base,['linear_attention','full_attention']);lm=InfluenceLM(base,cfg,enforce_memory=False);meta=adapter(ad,cfg,lm);before={k:t.read().copy() for k,t in lm.tensors.items() if 'linear_attn.' in k and t.numel<10000};pairs,scale=install(lm,Path(ad),meta)
            rng=np.random.default_rng(91);x=rng.standard_normal((2,9,16),dtype=np.float32)
            for name,p in pairs.items():
                t=lm.tensors[name];w=before[name].reshape(t.shape);actual=lm.linear(x,name.removeprefix('model.language_model.').removesuffix('.weight'),rows=3);expected=x@w.T+scale*((x@p['A'].T)@p['B'].T)
                np.testing.assert_allclose(actual,expected,atol=2e-6,rtol=2e-6)
                with self.assertRaises(ValueError):t.read(slice(1,10))
    def test_forward_and_head_decomposition_match_independent_overlay(self):
        with tempfile.TemporaryDirectory() as base,tempfile.TemporaryDirectory() as ad:
            cfg=fixture(base,['linear_attention','full_attention']);lm=InfluenceLM(base,cfg,enforce_memory=False);meta=adapter(ad,cfg,lm);pairs,scale=install(lm,Path(ad),meta);ids=np.array([[1,2,3,4,5,6,7,8,9]])
            out=Path(base)/'capture';out.mkdir();raw=lm.baseline(ids,out)
            ref=ExplicitLoRA(base,cfg,enforce_memory=False);ref.pairs,ref.scale=pairs,scale;expected,_=ref.forward(ids)
            np.testing.assert_allclose(lm.normalization(raw,'norm'),expected,atol=2e-6,rtol=2e-6)
            self.assertTrue(all(v['relative_l2']<2e-5 for v in lm.head_sum_errors))
    def test_zero_scale_preserves_original_and_bad_metadata_rejected(self):
        with tempfile.TemporaryDirectory() as base,tempfile.TemporaryDirectory() as ad:
            cfg=fixture(base,['full_attention','linear_attention']);lm=InfluenceLM(base,cfg,enforce_memory=False);meta=adapter(ad,cfg,lm);ids=np.array([[1,2,3,4,5,6,7,8,9]]);expected,_=lm.forward(ids);install(lm,Path(ad),meta,multiplier=0);actual,_=lm.forward(ids);np.testing.assert_array_equal(actual,expected)
            meta['gate']='custom'
            with self.assertRaises(AssertionError):install(lm,Path(ad),meta)