"""Question-only CUDA runtime for immutable semantic-head checkpoints.

v1.1 (Solomon v1.1):
  * precision: 'bf16' (default: bf16 weights, fp32 recurrence, the trainer's mode), 'fp32' (the qualified v1
    numerics) or 'int8' (torchao weight-only int8 on the decoder linears). Recorded in the identity.
  * head layout: any subset of the ten semantic heads loads, provided the yes/no head boolean/state4 is present.
    The v1.1 heads file carries ONE merged yes/no head (boolean/state4) and no entity/multilabel heads; the
    serving layer routes multi-label branches to it through the binding's head routing table. A non-v1 layout
    is recorded in the identity (answer_heads), so a binding pins it.
  * the document prefill keeps its final hidden states (and any tapped layer) for the trained evidence head.
"""
import copy
from contextlib import contextmanager
import hashlib
import json
from pathlib import Path
import threading
import numpy as np
from solomon.engine_numerics import CudaEngine
from solomon.heads import semantic_head_key

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

SEMANTIC_HEADS=('boolean/state4','entity/state4','multilabel/state4','single/choiceR','single/choiceS','single/sufficiency3',
                'ordered/choiceR','ordered/choiceS','ordered/sufficiency3','ordered/threshold4')
REQUIRED_HEADS=('boolean/state4',)


def load_heads(heads_path):
    """{head_key: {'weight','bias'}} from a semantic-heads .npz. Tolerant of the merged v1.1 layout."""
    with np.load(heads_path,allow_pickle=False) as arrays:
        heads={k[:-7]:{'weight':arrays[k].copy(),'bias':arrays[k[:-7]+'/bias'].copy()}
               for k in arrays.files if k.endswith('/weight')}
    unknown=sorted(set(heads)-set(SEMANTIC_HEADS))
    if unknown:raise ValueError('unknown semantic heads: '+', '.join(unknown))
    missing=[k for k in REQUIRED_HEADS if k not in heads]
    if missing:raise ValueError('semantic heads file lacks the yes/no head: '+', '.join(missing))
    return heads

class SolomonEngine(CudaEngine):
    def __init__(self, adapter, heads_path, *, expected_adapter=None, expected_heads=None, correctness=None, precision='bf16',
                 model_dir=None, prefix_layers=()):
        if expected_adapter and sha(adapter)!=expected_adapter:raise ValueError('adapter hash mismatch')
        if expected_heads and sha(heads_path)!=expected_heads:raise ValueError('heads hash mismatch')
        self.answer_heads=load_heads(heads_path)
        self.keep_prefix_hidden=True;self.prefix_layers=tuple(int(i) for i in prefix_layers)
        super().__init__(adapter=adapter,placement='question',precision=precision,**({'model_dir':model_dir} if model_dir else {}))
        self.base_identity=dict(self.identity)
        for head in self.answer_heads.values():
            if head['weight'].shape!=(10,5120) or head['bias'].shape!=(10,) or not all(np.isfinite(v).all() for v in head.values()):
                raise ValueError('invalid semantic head')
        self.correctness_models=correctness or {}
        self._task=None;self._lock=threading.RLock()
        if len(self.answer_heads)!=10:self.identity['answer_heads']=','.join(sorted(self.answer_heads))   # a v1 layout keeps the v1 key set
        self.identity.update(trained_heads_sha256=sha(heads_path),answer_engine_sha256=sha(__file__),
                             answer_projection='trained-semantic-head-float32',base_fingerprint=self.base_identity['fingerprint'])
        self.identity.pop('fingerprint',None)
        self.identity['fingerprint']=hashlib.sha256(json.dumps(self.identity,sort_keys=True).encode()).hexdigest()

    @contextmanager
    def task_context(self,task):
        with self._lock:
            previous=self._task;self._task=task
            try:yield self
            finally:self._task=previous

    def ask(self,state,block,n_letters,execution='cached',head_key=None,tap_layers=()):
        with self._lock:
            key=head_key or semantic_head_key(self._task,block)
            if key not in self.answer_heads or not 2<=n_letters<=10:raise ValueError('unknown semantic head or width')
            if execution not in ('cached','full'):raise ValueError('invalid execution')
            torch=self.torch;p=state['prefix_tokens'];previous=self.ctx['start'];handles=[];taps={}
            def hook(index):
                def capture(module,args,output):
                    h=output[0] if isinstance(output,tuple) else output
                    # Intermediate states are normalized using the frozen final norm.
                    taps[str(index)]=self.lm.norm(h[:,-1:])[0,-1].detach().float().cpu().numpy().copy()
                return capture
            try:
                for index in tap_layers:
                    if type(index)is not int or not 0<=index<len(self.lm.layers):raise ValueError('invalid layer tap')
                    handles.append(self.lm.layers[index].register_forward_hook(hook(index)))
                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('question prefix mismatch')
                    self.ctx['start']=0 if execution=='cached' else p
                    if execution=='cached':hidden=self.lm(inputs_embeds=embeds[:,p:],position_ids=pos[...,p:],past_key_values=copy.deepcopy(state['cache']),use_cache=True).last_hidden_state
                    else:hidden=self.lm(inputs_embeds=embeds,position_ids=pos,use_cache=False).last_hidden_state
                    head=self.answer_heads[key]
                    logits=torch.nn.functional.linear(hidden[:,-1].float(),torch.as_tensor(head['weight'],device=hidden.device),torch.as_tensor(head['bias'],device=hidden.device))[0,:n_letters]
                    probabilities=logits.softmax(-1).cpu().numpy()
                    feature=hidden[0,-1].float().cpu().numpy().copy()
                return {'letter_logits':logits.cpu().numpy(),'probabilities':probabilities,'prediction':int(probabilities.argmax()),
                        'hidden':feature,'taps':taps,'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
                for handle in handles:handle.remove()

    def predict_correctness(self,task,branches):
        from solomon.head_training import predict_correctness
        if task not in self.correctness_models:return {'reliability':None}
        hidden=[np.asarray(r['hidden'],float) for r in branches]
        probs=[np.asarray(r['probabilities'],float) for r in branches]
        feature=np.r_[np.mean(hidden,axis=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))]
        return {'reliability':predict_correctness(self.correctness_models[task],feature)}