"""Read-only FP32 effective weights for an existing pure-LoRA checkpoint."""
import re
import numpy as np
from functional.model import PREFIX,StreamingLM
from pilot.tensors import load_tensors
BASE_REVISION='15852e8c16360a2fea060d615a32b45270f8a8fc'
TARGETS={'in_proj_qkv','in_proj_a','in_proj_b','out_proj'}

class EffectiveWeight:
    def __init__(self,base,a,b,scale):
        assert a.ndim==b.ndim==2 and (b.shape[0],a.shape[1])==base.shape and a.shape[0]==b.shape[1]
        self.base,self.a,self.b,self.scale=base,a,b,scale
        self.shape,self.numel,self.dtype=base.shape,base.numel,'F32'
    def read(self,selector=slice(None)):
        if not isinstance(selector,slice):raise ValueError('Only contiguous row chunks supported')
        start,stop,step=selector.indices(self.numel);dim=self.shape[1]
        if step!=1 or start%dim or stop%dim:raise ValueError('Row aligned chunks required')
        w=self.base.read(slice(start,stop)).reshape(-1,dim)
        return (w+self.scale*(self.b[start//dim:stop//dim]@self.a)).astype(np.float32).reshape(-1)
    def evict(self):self.base.evict()


def install(lm,directory,meta,multiplier=1.):
    assert meta['base_model']=='Qwen/Qwen3.5-2B' and meta['base_revision']==BASE_REVISION and meta['gate'] is None and meta['arm']=='lora'
    tensors=load_tensors(directory);pairs={};pattern=re.compile(r'^base_model\.model\.model\.(layers\.(\d+)\.linear_attn\.(in_proj_qkv|in_proj_a|in_proj_b|out_proj))\.lora_([AB])\.default\.weight$')
    for name,tensor in tensors.items():
        match=pattern.fullmatch(name)
        if not match:raise ValueError('Unexpected trainable tensor: '+name)
        target=PREFIX+match[1]+'.weight';assert lm.config['layer_types'][int(match[2])]=='linear_attention'
        assert target in lm.tensors
        value=tensor.read().reshape(tensor.shape).copy();tensor.evict();assert np.isfinite(value).all()
        pairs.setdefault(target,{})[match[4]]=value
    expected={PREFIX+f'layers.{layer}.linear_attn.{role}.weight' for layer,kind in enumerate(lm.config['layer_types']) if kind=='linear_attention' for role in TARGETS}
    assert set(pairs)==expected and len(tensors)==meta['n_tensors']
    scale=meta['lora_alpha']/meta['lora_rank']*multiplier
    for target,pair in pairs.items():
        assert set(pair)=={'A','B'} and pair['A'].shape[0]==meta['lora_rank']
        lm.tensors[target]=EffectiveWeight(lm.tensors[target],pair['A'],pair['B'],scale)
    return pairs,scale


class ExplicitLoRA(StreamingLM):
    """Independent reference: base projection + scaled low-rank output, no W merge."""
    def linear(self,x,name,rows=256):
        y=super().linear(x,name,rows);pair=self.pairs.get(PREFIX+name+'.weight')
        if pair is not None:y=y+self.scale*((x@pair['A'].T)@pair['B'].T)
        return y.astype(np.float32)