"""Always-on, layer-local pre-FFN + residual activation steering; frozen Qwen."""
import os,json,math,time,threading
from pathlib import Path
os.environ.setdefault('PYTORCH_ENABLE_MPS_FALLBACK','1')
# Bundled weights are read-only; auxiliary model-library caches live outside the signed app.
_bundled=Path(__file__).resolve().parent/'.hf/hub/models--Qwen--Qwen3.5-0.8B/snapshots/2fc06364715b967f1860aea9cf38778875588b17'
if _bundled.is_dir():
    os.environ.setdefault('AFFECTOR_MODEL',str(_bundled))
    os.environ['HF_HOME']=str(Path.home()/'Library/Caches/AFFECTOR/huggingface')

import torch
MODEL='Qwen/Qwen3.5-0.8B'
REVISION='2fc06364715b967f1860aea9cf38778875588b17'
from rewrite_policy import SYSTEM
from generation import repetition,sentence_end
# Matched contrastive examples; exploratory calibration, not validated emotion labels.
PAIRS=[
 [('I returned to the house with despair. Nothing felt worth keeping.','I returned to the house with joy. Everything felt worth keeping.'),('She read the letter with dread and bitterness.','She read the letter with delight and gratitude.'),('The morning felt bleak, hostile and hopeless.','The morning felt welcoming, hopeful and beautiful.'),('He regarded the meeting as a painful disappointment.','He regarded the meeting as a wonderful success.')],
 [('I felt ashamed of myself and dwelt on my own failures.','I felt angry at the people who had betrayed my trust.'),('Her sadness turned inward as she questioned herself.','Her sadness reached outward toward the suffering family.'),('He was proud of his own courage.','He admired the courage of his companions.'),('I worried about my own thoughts and feelings.','I worried about the safety of everyone outside.')],
 [('I felt mildly uneasy as I waited at the door.','I was overwhelmed by terror as I waited at the door.'),('She was slightly pleased with the news.','She was ecstatic, swept away by the news.'),('He felt a small twinge of sadness.','He was consumed by a devastating grief.'),('I felt faint irritation at the interruption.','I felt a blazing fury at the interruption.')]
]

class ActivationModule(torch.nn.Module):
    def __init__(self,basis):
        super().__init__();self.weight=torch.nn.Parameter(torch.tensor(basis,dtype=torch.float32),requires_grad=False)
        self.previous=None
    def reset(self):self.previous=None
    def forward(self,h,coordinates,controls,layer_weight):
        if controls['strength']==0 or controls['budget']==0 or layer_weight==0:return torch.zeros_like(h[:,-1,:])
        coeff=torch.tensor(coordinates,dtype=h.dtype,device=h.device)
        magnitude=coeff.abs().pow(1/controls['axis_temperature'])
        order=torch.argsort(magnitude,descending=True);mask=torch.zeros_like(coeff)
        running=0.;total=float(magnitude.sum())
        for index in order[:controls['axis_top_k']]:
            if running>=controls['axis_mass']*total and running>0:break
            mask[index]=1;running+=float(magnitude[index])
        coeff=coeff.sign()*magnitude*mask
        delta=coeff @ self.weight.to(device=h.device,dtype=h.dtype)
        rms=h[:,-1,:].float().pow(2).mean(-1,keepdim=True).sqrt().to(h.dtype)
        delta=delta.unsqueeze(0)*rms*controls['strength']*layer_weight
        if self.previous is not None:delta=controls['inertia']*self.previous+(1-controls['inertia'])*delta
        norm=delta.float().pow(2).mean(-1,keepdim=True).sqrt().clamp_min(1e-12)
        delta=delta*torch.clamp(controls['budget']*rms/norm,max=1)
        self.previous=delta.detach()
        return delta

class Engine:
    def __init__(self,progress=lambda s:None):
        from transformers import AutoTokenizer,Qwen3_5ForConditionalGeneration
        import transformers
        torch.set_num_threads(4)
        self.progress=progress
        self.device='mps' if torch.backends.mps.is_available() else 'cpu'
        progress('Loading frozen Qwen3.5-0.8B on '+self.device.upper()+'…')
        path=os.environ.get('AFFECTOR_MODEL')
        kwargs={'local_files_only':True}
        if not path:kwargs['revision']=REVISION;path=MODEL
        self.tokenizer=AutoTokenizer.from_pretrained(path,**kwargs)
        self.model=Qwen3_5ForConditionalGeneration.from_pretrained(path,dtype=torch.float32,attn_implementation='eager',**kwargs).to(self.device).eval()
        self.model.requires_grad_(False);self.layers=self.model.model.language_model.layers
        self.identity={'model':MODEL,'revision':REVISION,'width':1024,'site':'pre_ffn_with_residual','dtype':'float32','transformers':transformers.__version__,'torch':torch.__version__}
        self.injections=0;self.delta_rms=[]
    def tokens(self,system,text):
        rendered=self.tokenizer.apply_chat_template([{'role':'system','content':system},{'role':'user','content':text}],tokenize=False,add_generation_prompt=True,enable_thinking=False)
        inputs=self.tokenizer(rendered,return_tensors='pt').to(self.device)
        if inputs.input_ids.shape[-1]>8192:raise ValueError('Passage/prompt exceeds 8,192 tokens; shorten the passage or context.')
        return inputs
    @torch.inference_mode()
    def compile(self,g,cancel,progress):
        selected=sorted(map(int,g['layers']));captures={};handles=[]
        def capture(index):
            def hook(module,args):captures[index]=args[0][0,-1].detach().float().cpu().clone()
            return hook
        try:
            for index in selected:handles.append(self.layers[index].post_attention_layernorm.register_forward_pre_hook(capture(index)))
            vectors={str(i):[] for i in selected};step=0
            for pairs in PAIRS:
                differences={i:[] for i in selected}
                for negative,positive in pairs:
                    states=[]
                    for text in (negative,positive):
                        if cancel.is_set():raise InterruptedError('Calibration stopped.')
                        self.model(**self.tokens('Continue the passage in the same emotional voice.',text),use_cache=False,logits_to_keep=1)
                        states.append({i:captures[i].clone() for i in selected});step+=1;progress(f'Calibrating activation basis {step}/24…')
                    for i in selected:differences[i].append(states[1][i]-states[0][i])
                for i in selected:
                    vector=torch.stack(differences[i]).mean(0);vector/=vector.pow(2).mean().sqrt().clamp_min(1e-8)
                    vectors[str(i)].append(vector.tolist())
            return {'identity':self.identity,'vectors':vectors,'calibration':'Mean matched-context activation differences; 4 pairs per axis. Unvalidated English prototype.','pairs':PAIRS}
        finally:
            for handle in handles:handle.remove()
    def hooks(self,g,coordinates):
        if g['basis']['identity']!=self.identity:raise ValueError('Steering basis belongs to a different model/runtime. Recalibrate this geometry.')
        handles=[];pending={};modules={}
        if any(k not in g['basis']['vectors'] for k in g['layers']):raise ValueError('New layers need calibration before use.')
        for key,weight in g['layers'].items():
            if key not in g['basis']['vectors']:raise ValueError('New layers need calibration before use.')
            modules[key]=ActivationModule(g['basis']['vectors'][key]);layer=self.layers[int(key)]
            def before(module,args,key=key,weight=weight):
                h=args[0];d=modules[key](h,coordinates,g['controls'],weight);pending[key]=d
                if not torch.count_nonzero(d):return None
                self.injections+=1;self.delta_rms.append(float(d.float().pow(2).mean().sqrt()))
                y=h.clone();y[:,-1,:]+=d;return (y,*args[1:])
            def after(module,args,out,key=key):
                d=pending.pop(key,None)
                if d is None or not torch.count_nonzero(d):return out
                if not isinstance(out,torch.Tensor):raise RuntimeError('Unexpected decoder output; hook requires revalidation.')
                y=out.clone();y[:,-1,:]+=d;return y
            handles.extend([layer.post_attention_layernorm.register_forward_pre_hook(before),layer.register_forward_hook(after)])
        return handles
    @torch.inference_mode()
    def generate(self,text,*,g=None,coordinates=None,system=SYSTEM,max_tokens=512,temperature=0.,top_p=.95,top_k=40,cancel=None,time_limit=120,max_chars=None,soft_chars=None):
        from transformers import StoppingCriteria,StoppingCriteriaList
        cancel=cancel or threading.Event()
        start=time.monotonic();progress=self.progress;steps=[0];length_stopped=[False];quality_stop=[None]
        tokenizer=self.tokenizer;prompt_length=[0]
        class Stop(StoppingCriteria):
            def __call__(self,input_ids,scores,**kwargs):
                steps[0]+=1
                if steps[0]%64==0:progress(f'Local model · {steps[0]} output tokens · {int(time.monotonic()-start)} seconds')
                if max_chars is not None:
                    draft=tokenizer.decode(input_ids[0,prompt_length[0]:].tolist(),skip_special_tokens=True)
                    length_stopped[0]=len(draft)>=max_chars
                    if soft_chars is not None:
                        if steps[0]%8==0 and repetition(draft):quality_stop[0]='repetition'
                        elif len(draft)>=soft_chars and sentence_end(draft):quality_stop[0]='sentence_boundary'
                return cancel.is_set() or time.monotonic()-start>time_limit or length_stopped[0] or quality_stop[0] is not None
        inputs=self.tokens(system,text);prompt_length[0]=inputs.input_ids.shape[-1];handles=[];self.injections=0;self.delta_rms=[];start=time.monotonic()
        try:
            if g is not None:handles=self.hooks(g,coordinates)
            options={'do_sample':temperature>0,'max_new_tokens':int(max_tokens),'use_cache':True,'repetition_penalty':1.1,'no_repeat_ngram_size':8,'logits_to_keep':1,'pad_token_id':self.tokenizer.eos_token_id,'stopping_criteria':StoppingCriteriaList([Stop()])}
            if temperature>0:options.update(temperature=temperature,top_p=top_p,top_k=int(top_k))
            out=self.model.generate(**inputs,**options)
            if cancel.is_set() and max_chars is None:raise InterruptedError('Transformation stopped.')
            if time.monotonic()-start>time_limit and max_chars is None:raise ValueError('Local generation exceeded its 120-second limit; this passage remains unmodified or unscored.')
            ids=out[0,inputs.input_ids.shape[-1]:].tolist();eos=self.model.generation_config.eos_token_id;eos=eos if isinstance(eos,list) else [eos]
            decoded=self.tokenizer.decode(ids,skip_special_tokens=True).strip()
            reason='cancelled' if cancel.is_set() else quality_stop[0] if quality_stop[0] else 'character_limit' if length_stopped[0] else 'time_limit' if time.monotonic()-start>time_limit else 'token_limit' if ids and ids[-1] not in eos and len(ids)>=max_tokens else 'completed'
            return {'text':decoded[:max_chars] if max_chars is not None else decoded,'stop_reason':reason,'character_limit':max_chars,'tokens':len(ids),'truncated':bool(ids and ids[-1] not in eos and len(ids)>=max_tokens),'seconds':round(time.monotonic()-start,3),'decoding_guard':{'repetition_penalty':1.1,'no_repeat_ngram_size':8},'activation_injections':self.injections,'mean_delta_rms':sum(self.delta_rms)/max(1,len(self.delta_rms))}
        finally:
            for h in handles:h.remove()