"""Frozen-model head interventions and continuous influence concentration."""
import numpy as np
from functional.model import StreamingLM,PREFIX,rope,causal_conv,attention


def concentration(d):
    """Last axis donor head; zero total is undefined, represented by NaNs."""
    d=np.asarray(d,dtype=np.float64)
    if not np.isfinite(d).all() or np.any(d<0):raise ValueError('Finite nonnegative influences required')
    h=d.shape[-1];total=d.sum(-1);valid=total>0
    p=np.divide(d,total[...,None],out=np.zeros_like(d),where=valid[...,None]);square=np.sum(p*p,-1)
    neff=np.divide(1,square,out=np.full_like(total,np.nan),where=valid)
    ordered=np.sort(p,axis=-1)[...,::-1]
    return {'total':total,'max_absolute':d.max(-1),'valid':valid,'effective_heads':neff,'effective_fraction':neff/h,
            'concentration':np.where(valid,(square-1/h)/(1-1/h),np.nan),
            'top1':np.where(valid,ordered[...,0],np.nan),'top25':np.where(valid,ordered[...,:max(1,h//4)].sum(-1),np.nan),
            'top50':np.where(valid,ordered[...,:max(1,h//2)].sum(-1),np.nan),
            'bottom25':np.where(valid,ordered[...,-max(1,h//4):].sum(-1),np.nan),
            'heads90':np.where(valid,1+np.sum(np.cumsum(ordered,axis=-1)<.9,axis=-1),np.nan)}


def masks_from_influences(d,weak=True):
    """d [donor,T,recipient]; output masks [recipient,T,donor]."""
    ranked=np.argsort(d.transpose(2,1,0),axis=-1,kind='stable');n=d.shape[0]//2
    selected=ranked[...,:n] if weak else ranked[...,-n:]
    mask=np.zeros_like(ranked,dtype=np.float32);np.put_along_axis(mask,selected,1,axis=-1)
    return mask


def js_rows(a,b):
    a=np.asarray(a,dtype=np.float64);b=np.asarray(b,dtype=np.float64);m=(a+b)/2
    def term(x):return np.where(x>0,x*np.log(np.maximum(x,1e-300)/np.maximum(m,1e-300)),0)
    return np.maximum(0,.5*np.sum(term(a)+term(b),axis=-1))


class InfluenceLM(StreamingLM):
    capturing=False
    def linear(self,x,name,rows=256):
        y=super().linear(x,name,rows)
        if self.capturing and (name.endswith('self_attn.o_proj') or name.endswith('linear_attn.out_proj')):
            layer=int(name.split('.')[1]);nh=self.config['num_attention_heads'] if 'self_attn' in name else self.config['linear_num_value_heads']
            xx=x.reshape(*x.shape[:2],nh,-1);dim=xx.shape[-1];t=self.tensors[PREFIX+name+'.weight'];od,idim=t.shape
            out=np.empty((*x.shape[:2],nh,od),dtype=np.float32)
            for start in range(0,od,256):
                self.guard();stop=min(od,start+256);w=t.read(slice(start*idim,stop*idim)).reshape(stop-start,idim)
                for head in range(nh):out[:,:,head,start:stop]=xx[:,:,head]@w[:,head*dim:(head+1)*dim].T
                del w;t.evict()
            err=float(np.linalg.norm(out.sum(2).astype(np.float64)-y)/max(np.linalg.norm(y.astype(np.float64)),1e-30))
            if err>2e-5:raise AssertionError(('head sum mismatch',layer,err))
            self.head_outputs=out;self.head_sum_errors.append({'layer':layer,'relative_l2':err})
        return y

    def baseline(self,ids,root):
        h=self.embedding(ids);self.head_sum_errors=[]
        for layer,kind in enumerate(self.config['layer_types']):
            raw=h.copy();n=self.normalization(h,f'layers.{layer}.input_layernorm');self.capturing=True
            out,_=self.full_attention(n,layer) if kind=='full_attention' else (self.linear_attention(n,layer),None)
            self.capturing=False;heads=self.head_outputs;self.head_outputs=None
            mid=raw+out;h=mid+self.mlp(self.normalization(mid,f'layers.{layer}.post_attention_layernorm'),layer)
            np.savez(root/f'layer-{layer}.npz',raw=raw,out=out,heads=heads,next_raw=h)
            print('baseline layer',layer,flush=True)
        return h

    def recipient(self,raw,layer):
        h=self.normalization(raw,f'layers.{layer}.input_layernorm');c=self.config;b,t=h.shape[:2]
        if c['layer_types'][layer]=='full_attention':
            p=f'layers.{layer}.self_attn.';nh=c['num_attention_heads'];nk=c['num_key_value_heads'];d=c['head_dim']
            q=self.linear(h,p+'q_proj').reshape(b,t,nh,2*d)[...,:d];k=self.linear(h,p+'k_proj').reshape(b,t,nk,d);v=self.linear(h,p+'v_proj').reshape(b,t,nk,d)
            q=self.normalization(q,p+'q_norm');k=self.normalization(k,p+'k_norm');r=c['rope_parameters'];q,k=(rope(z,r['rope_theta'],r['partial_rotary_factor']) for z in (q,k))
            k=np.repeat(k,nh//nk,axis=2);v=np.repeat(v,nh//nk,axis=2)
            _,prob=attention(q,k,v)
            return {'Q':q,'K':k,'V':v,'P':prob}
        p=f'layers.{layer}.linear_attn.';kh=c['linear_num_key_heads'];vh=c['linear_num_value_heads'];kd=c['linear_key_head_dim'];vd=c['linear_value_head_dim']
        x=causal_conv(self.linear(h,p+'in_proj_qkv'),self.small(p+'conv1d.weight')[:,0,:]);q,k,v=np.split(x,[kh*kd,2*kh*kd],axis=-1)
        q=q.reshape(b,t,kh,kd);k=k.reshape(b,t,kh,kd);v=v.reshape(b,t,vh,vd)
        if vh!=kh:q,k=np.repeat(q,vh//kh,axis=2),np.repeat(k,vh//kh,axis=2)
        q=q/np.sqrt(np.sum(q*q,axis=-1,keepdims=True)+1e-6)/np.float32(np.sqrt(kd));k=k/np.sqrt(np.sum(k*k,axis=-1,keepdims=True)+1e-6)
        return {'Q':q,'K':k,'V':v}

    def intervene(self,cache,layer,masks):
        # Each scenario has a full coherent token-dependent deletion mask.
        removed=np.einsum('bth,thd->btd',masks,cache['heads'][0],optimize=True)
        mid=cache['raw']+cache['out']-removed
        nxt=mid+self.mlp(self.normalization(mid,f'layers.{layer}.post_attention_layernorm'),layer)
        return self.recipient(nxt,layer+1)