"""Final normalized feature extraction, preserving the qualified CUDA engine."""
import hashlib,json
from pathlib import Path
import numpy as np

def sha(path):return hashlib.sha256(Path(path).read_bytes()).hexdigest()

def token_ids_for(engine,state,block,n):
    text=engine._render(state['parts'],block);raw=engine.t.encode(text,add_special_tokens=False);ids=[]
    for c in 'ABCDEFGHIJ'[:n]:
        ext=engine.t.encode(text+c,add_special_tokens=False)
        if ext[:-1]!=raw or len(ext)!=len(raw)+1:raise ValueError('unstable answer token continuation')
        ids.append(ext[-1])
    return ids

def capture(engine,state,block,n):
    """Prehook sees precisely the post-final-norm input used by the vocabulary head."""
    captured=[]
    def hook(module,args):
        h=args[0]
        if h.ndim!=2 or h.shape[0]!=1:raise ValueError('unexpected final head input shape')
        captured.append(h.detach().float().cpu().numpy()[0].copy())
    handle=engine.model.lm_head.register_forward_pre_hook(hook)
    try:result=engine.ask(state,block,n,execution='cached')
    finally:handle.remove()
    if len(captured)!=1:raise ValueError('expected exactly one vocabulary projection')
    ids=token_ids_for(engine,state,block,n)
    if not np.isfinite(captured[0]).all():raise ValueError('nonfinite hidden features')
    return captured[0],ids,result

def copy_letter_head(engine):
    ids=[]
    for c in 'ABCDEFGHIJ':
        pieces=engine.t.encode(c,add_special_tokens=False)
        if len(pieces)!=1:raise ValueError('letters must be one token')
        ids.append(pieces[0])
    head=engine.model.lm_head;w=head.weight.detach()[ids].float().cpu().numpy().copy();b=getattr(head,'bias',None)
    return {'weight':w,'bias':np.zeros(10,np.float32) if b is None else b.detach()[ids].float().cpu().numpy().copy(),'token_ids':np.array(ids,np.int64),'vocabulary_size':np.array(head.weight.shape[0],np.int64),'bias_present':np.array(b is not None)}

def atomic_json(path,value):
    path=Path(path);tmp=path.with_suffix(path.suffix+'.tmp');tmp.write_text(json.dumps(value,sort_keys=True,indent=2));tmp.replace(path)

def write_document(path,features,token_ids,results,metadata,identity):
    """NPZ first, then checksummed completion JSON. No cumulative shard rewrite."""
    path=Path(path);n=len(features)
    if not n or not(len(token_ids)==len(results)==len(metadata)==n):raise ValueError('document arrays must align')
    h=np.asarray(features,np.float32);logits=np.zeros((n,10),np.float32);tokens=np.full((n,10),-1,np.int64);widths=[]
    if h.ndim!=2 or not np.isfinite(h).all():raise ValueError('invalid document feature matrix')
    rows=[]
    for i,(ids,r,m) in enumerate(zip(token_ids,results,metadata)):
        values=np.asarray(r['letter_logits'],np.float32);width=len(ids)
        if not 2<=width<=10 or values.shape!=(width,) or not np.isfinite(values).all():raise ValueError('invalid teacher row')
        logits[i,:width]=values;tokens[i,:width]=ids;widths.append(width)
        rows.append({**m,'n':width,'mass':float(r['mass']),'top_is_letter':bool(r['top_is_letter']),'prompt_tokens':int(r['prompt_tokens']),'branch_tokens':int(r['branch_tokens']),'reused_prefix_tokens':int(r['reused_prefix_tokens'])})
    tmp=path.with_suffix('.npz.tmp')
    with tmp.open('wb') as f:np.savez(f,hidden=h,teacher_logits=logits,token_ids=tokens,widths=np.asarray(widths,np.int64))
    tmp.replace(path)
    report={'complete':True,'identity':identity,'rows':rows,'npz_sha256':sha(path),'feature_shape':list(h.shape),'feature_kind':'final-normalized-lm-head-input','storage':'uncompressed-npz-float32'}
    atomic_json(path.with_suffix('.json'),report)
    return report

def validate_document(path,identity,expected_ids):
    path=Path(path);meta=path.with_suffix('.json')
    if not path.exists() or not meta.exists():return False
    r=json.loads(meta.read_text())
    if r.get('complete')is not True or r.get('identity')!=identity or r.get('npz_sha256')!=sha(path) or [x['id'] for x in r['rows']]!=list(expected_ids):raise ValueError('resumed document identity mismatch')
    return True

class TrainedHeads:
    """Validated saved answer/correctness heads. No backbone or label access."""
    def __init__(self,directory,*,expected_report_sha256=None,expected_weights_sha256=None):
        directory=Path(directory);report_path=directory/'report.json';weights_path=directory/'weights.npz'
        self.report_sha256=sha(report_path);self.weights_sha256=sha(weights_path)
        if expected_report_sha256 is not None and self.report_sha256!=expected_report_sha256:raise ValueError('selected head report changed')
        if expected_weights_sha256 is not None and self.weights_sha256!=expected_weights_sha256:raise ValueError('selected head weights changed')
        self.report=json.loads(report_path.read_text())
        if self.report['weights_sha256']!=self.weights_sha256:raise ValueError('trained head weight checksum mismatch')
        self.models={}
        with np.load(weights_path,allow_pickle=False) as archive:
            for kind,entries in self.report['models'].items():
                self.models[kind]={}
                for key,fields in entries.items():
                    model={name:archive[value['array']].astype(np.float64) if isinstance(value,dict) and set(value)=={'array'} else value for name,value in fields.items()}
                    if any(not np.isfinite(value).all() for value in model.values() if isinstance(value,np.ndarray)):raise ValueError('nonfinite trained head')
                    if kind=='heads':
                        if model['weight'].ndim!=2 or model['weight'].shape[0]!=10 or model['bias'].shape!=(10,):raise ValueError('invalid answer head shape')
                    elif kind=='correctness':
                        if model['mean'].ndim!=1 or model['scale'].shape!=model['mean'].shape or np.any(model['scale']<=0) or model['weights'].shape!=(len(model['mean'])+1,):raise ValueError('invalid correctness head shape')
                    else:raise ValueError('unknown trained model kind')
                    self.models[kind][key]=model
        if not self.models.get('heads') or not self.models.get('correctness'):raise ValueError('both trained answer and correctness heads required')

    def answer(self,hidden,key,n):
        if key not in self.models['heads']:raise ValueError('no trained answer head for '+str(key))
        if type(n)is not int or not 2<=n<=10:raise ValueError('invalid head output width')
        h=np.asarray(hidden,np.float64);model=self.models['heads'][key]
        if h.shape!=(model['weight'].shape[1],) or not np.isfinite(h).all():raise ValueError('invalid hidden state for answer head')
        return (model['weight']@h+model['bias'])[:n]

    def correctness(self,task,branches):
        from scipy.special import expit
        if task not in self.models['correctness'] or not branches:raise ValueError('no trained whole-answer correctness head')
        hidden=np.asarray([r['hidden'] for r in branches],np.float64);probs=[]
        if hidden.ndim!=2 or not np.isfinite(hidden).all():raise ValueError('invalid correctness hidden representations')
        for row in branches:
            p=np.asarray(row['probabilities'],float)
            if p.ndim!=1 or len(p)<2 or not np.isfinite(p).all() or np.any(p<0) or abs(p.sum()-1)>1e-6:raise ValueError('invalid correctness answer distribution')
            probs.append(p)
        feature=np.r_[hidden.mean(0),min(float(p.max()) for p in probs),min(float(np.sort(p)[-1]-np.sort(p)[-2]) for p in probs),max(float(-np.sum(p*np.log(np.maximum(p,1e-300)))/np.log(len(p))) for p in probs),float(len(branches)),float(sum(len(p) for p in probs))]
        model=self.models['correctness'][task]
        if feature.shape!=model['mean'].shape:raise ValueError('whole-answer confidence feature width mismatch')
        value=float(expit(np.r_[1.,(feature-model['mean'])/model['scale']]@model['weights']))
        return {'reliability':value,'features':feature.tolist(),'task':task,'target':'whole_return_exact_correctness','calibration_established':False,'heads_sha256':self.weights_sha256}


def semantic_head_key(task,block):
    from solomon.engine_contract import SUFFICIENCY,SINGLE_R,SINGLE_S,ORDERED_R,ORDERED_S,LABEL,BOOLEAN_TASK
    if task not in ('boolean','single','ordered','multilabel','entity'):raise ValueError('explicit answering task required')
    if block.startswith(SUFFICIENCY) and task in ('single','ordered'):return task+'/sufficiency3'
    for prefix,t,form in ((SINGLE_R,'single','R'),(SINGLE_S,'single','S'),(ORDERED_R,'ordered','R'),(ORDERED_S,'ordered','S')):
        if block.startswith(prefix) and task==t:return task+'/choice'+form
    if block.startswith(LABEL) and task=='multilabel':return 'multilabel/state4'
    if block.startswith(BOOLEAN_TASK):
        if task in ('boolean','entity'):return task+'/state4'
        if task=='ordered':return 'ordered/threshold4'
    raise ValueError('task block has no trained semantic head mapping')


from solomon.engine_reasoning import ReasoningEngine
class HeadsEngine(ReasoningEngine):
    """Question-only backbone with trained small heads; native LM head retained for reasoning.

    Wrap a complete readout in task_context(task), or supply ask(head_key=...).
    Answer projection is deliberately evaluated in CPU float64 on the exact copied
    float32 hidden vector to match saved-feature evaluation; no speedup is claimed.
    """
    def __init__(self,heads_dir,*,expected_base_fingerprint,expected_report_sha256=None,expected_weights_sha256=None,**kwargs):
        import threading
        self.heads=TrainedHeads(heads_dir,expected_report_sha256=expected_report_sha256,expected_weights_sha256=expected_weights_sha256)
        super().__init__(**kwargs)
        if self.identity['fingerprint']!=expected_base_fingerprint:raise ValueError('trained head backbone/adapter runtime mismatch')
        self.base_identity=dict(self.identity);self._head_task=None;self._head_lock=threading.RLock()
        self.identity={**self.identity,'base_fingerprint':expected_base_fingerprint,'trained_heads_sha256':self.heads.weights_sha256,'trained_heads_report_sha256':self.heads.report_sha256,'heads_runtime_sha256':sha(__file__),'answer_projection':'trained-semantic-head-float64-cpu'}
        self.identity.pop('fingerprint',None);self.identity['fingerprint']=hashlib.sha256(json.dumps(self.identity,sort_keys=True).encode()).hexdigest()

    def task_context(self,task):
        from contextlib import contextmanager
        @contextmanager
        def context():
            with self._head_lock:
                previous=self._head_task;self._head_task=task
                try:yield self
                finally:self._head_task=previous
        return context()

    def prefill(self,document):
        with self._head_lock:return super().prefill(document)

    def ask(self,state,block,n_letters,execution='cached',head_key=None):
        import copy
        with self._head_lock:
            key=head_key or semantic_head_key(self._head_task,block)
            if key not in self.heads.models['heads']:raise ValueError('missing trained head '+key)
            if self.placement!='question':raise ValueError('trained heads require question-only adaptation')
            if execution not in ('cached','full'):raise ValueError('invalid execution')
            torch=self.torch;p=state['prefix_tokens'];previous=self.ctx['start']
            try:
                with torch.inference_mode():
                    text=self._render(state['parts'],block);ids,embeds,pos,_=self._encode(state['parts'],text,state['features'])
                    if not torch.equal(ids[:,:p],state['prefix_ids']):raise ValueError('head question prefix mismatch')
                    self.ctx['start']=(0 if execution=='cached' or self.placement=='full' else p) if self.adapter else None
                    if execution=='cached':h=self.lm(inputs_embeds=embeds[:,p:],position_ids=pos[...,p:],past_key_values=copy.deepcopy(state['cache']),use_cache=True).last_hidden_state
                    else:h=self.lm(inputs_embeds=embeds,position_ids=pos,use_cache=False).last_hidden_state
                    hidden=h[0,-1].float().cpu().numpy().copy()
                logits=self.heads.answer(hidden,key,n_letters);probs=np.exp(logits-logits.max());probs/=probs.sum()
                return {'letter_logits':logits,'probabilities':probs,'prediction':int(probs.argmax()),'hidden':hidden.tolist(),'head_key':key,'mass':None,'top_is_letter':None,'mass_available':False,'execution':execution,'fallback':'','prompt_tokens':int(ids.shape[1]),'branch_tokens':int(ids.shape[1])-p,'input_tokens':int(ids.shape[1])-p,'reused_prefix_tokens':p if execution=='cached' else 0,'generation_tokens':0,'fingerprint':self.identity['fingerprint']}
            finally:self.ctx['start']=previous

    def estimate(self,state,block,n,stage='fast',max_new_tokens=512):
        if stage!='fast' or not any('image' in part for part in state['parts']):
            return super().estimate(state,block,n,stage=stage,max_new_tokens=max_new_tokens)
        if type(n)is not int or not 2<=n<=10:raise ValueError('invalid letter count')
        # Reuse prefetched vision features; no transformer branch is evaluated.
        with self._head_lock,self.torch.inference_mode():
            ids,_,_,_=self._encode(state['parts'],self._render(state['parts'],block),state['features'])
            prefix=state['prefix_tokens']
            if not self.torch.equal(ids[:,:prefix],state['prefix_ids']):raise ValueError('image budget prefix mismatch')
            return {'input_tokens':int(ids.shape[1])-prefix,'generated_tokens':0,'branches':1,
                    'full_prompt_tokens':int(ids.shape[1]),'reused_prefix_tokens':prefix}

    def predict_correctness(self,task,branches):return self.heads.correctness(task,branches)

    def reason(self,*args,**kwargs):
        with self._head_lock:return super().reason(*args,**kwargs)