# Modified by xypai, 2026-10-29: checkpoint path names only.
"""LFM3/RV32/RAM model with ten registered analytic circuit modules."""
from pathlib import Path
import copy
import hashlib
import itertools
import json
import shutil
import torch
try:
    from lfmputer_runtime.memory.integration import RAMLFMComputer
    from lfmputer_runtime.memory.bridge import RAMBridge
    from lfmputer_runtime.computer.circuit import CompiledComputer
except ModuleNotFoundError:
    from lfmram.integration import RAMLFMComputer
    from lfmram.bridge import RAMBridge
    from lfmcomputer.circuit import CompiledComputer
from .engine import AnalyticEngine,FEATURES


class AnalyticLFMComputer(RAMLFMComputer):
    def __init__(self,language_model,computer,memory_spec=None,manifest=None,ram_spec=None,bridge=None,analytic_spec=None):
        super().__init__(language_model,computer,memory_spec,manifest,ram_spec,bridge)
        self.analytics=AnalyticEngine(analytic_spec)

    def forward(self,input_ids=None,*,mode='language',feature=None,payload=None,**kwargs):
        if mode=='analytic':return self.analytics.run(feature,payload,self.tokenizer)
        if mode=='analytic_tensor':return self.analytics(feature,**kwargs)
        return super().forward(input_ids=input_ids,mode=mode,**kwargs)

    def new_session(self):
        session=super().new_session()
        if not getattr(self,'include_computer_tools',True):
            session.messages[0]['content']='You are a helpful assistant. Follow the user instructions exactly. You may use the available tools when useful. Never claim a tool succeeded without a successful response. Respond in English.'
        enabled=[name for i,name in enumerate(FEATURES) if bool(self.analytics.enabled[i])]
        if enabled:
            session.messages[0]['content']+=' Available built-in analytic circuits: '+', '.join(enabled)+'. Use the appropriate analytic function for supported formal computations. Pass the typed arguments from its declared schema. Scope IR belongs in events, not text; byte encoders take bytes. Execute calls using tool-call syntax, not a final payload object. After receiving a result, obey the user requested output structure exactly; extract the requested field rather than nesting the full result object. Do not guess a circuit result.'
        return session

    def chat(self,session,prompt,**kwargs):
        from .driver import chat
        return chat(self,session,prompt,**kwargs)

    def fingerprint(self):
        h=hashlib.sha256(super().fingerprint().encode())
        h.update(json.dumps(self.analytics.spec,sort_keys=True).encode())
        for name,value in sorted(self.analytics.state_dict().items()):
            h.update(name.encode());h.update(value.detach().cpu().contiguous().numpy().tobytes())
        return h.hexdigest()

    def save_pretrained(self,path,tokenizer=None):
        info=super().save_pretrained(path,tokenizer);path=Path(path)
        info['architecture']='LFMputer'
        info['analytic_spec']=self.analytics.spec
        info['analytic_features']=list(FEATURES)
        info['analytic_implementation']='registered fixed-weight tensor circuits with model-side text/JSON codec and generation control'
        (path/'computer_config.json').write_text(json.dumps(info,indent=2))
        shutil.copytree(Path(__file__).parent,path/'lfmputer_runtime/analytics',dirs_exist_ok=True,
                        ignore=shutil.ignore_patterns('__pycache__','evaluate.py','tasks.py'))
        return info

    @classmethod
    def from_pretrained(cls,path,device='cpu',backend='torch'):
        from accelerate import init_empty_weights
        from safetensors.torch import load_file
        from transformers import AutoConfig,AutoModelForCausalLM,AutoTokenizer
        path=Path(path);cfg=json.loads((path/'computer_config.json').read_text())
        if 'analytic_spec' not in cfg:raise ValueError('Upgrade the base model before using this loader')
        definition=path/'backbone'
        config=AutoConfig.from_pretrained(definition,trust_remote_code=True,local_files_only=True)
        config._attn_implementation=cfg.get('language_attention','eager')
        with init_empty_weights():language=AutoModelForCausalLM.from_config(config,trust_remote_code=True)
        bridge=RAMBridge(spec=cfg['ram_bridge']) if cfg.get('ram_bridge') else None
        model=cls(language,CompiledComputer(cfg['circuit']),cfg['memory'],cfg['manifest'],cfg['ram_spec'],bridge,cfg['analytic_spec'])
        values=load_file(str(path/'model.safetensors'))
        for alias,original in cfg['aliases'].items():values[alias]=values[original]
        model.load_state_dict(values,strict=True,assign=True);model.language_model.tie_weights()
        if model.fingerprint()!=cfg['identity']:raise ValueError('Analytic model identity mismatch')
        model.definition_path=definition
        model.to(device).eval();model.computer.prepare(backend);model.analytics.prepare()
        model.tokenizer=AutoTokenizer.from_pretrained(path,trust_remote_code=True,local_files_only=True)
        return model

    @torch.inference_mode()
    def generate_finite_schema(self,prompt,schema,max_candidates=256,max_tokens=128):
        """A real logits mask for finite flat enum schemas. Not arbitrary JSON Schema."""
        from transformers import LogitsProcessor,LogitsProcessorList
        if not bool(self.analytics.enabled[0]):raise RuntimeError('Schema circuit disabled')
        props=schema.get('properties',{})
        if schema.get('additionalProperties',True) or set(schema.get('required',[]))!=set(props):
            raise ValueError('Finite generation requires all fields required and additionalProperties=false')
        domains=[prop.get('enum') for prop in props.values()]
        if not domains or any(not x for x in domains):raise ValueError('Every field needs a finite enum')
        count=1
        for domain in domains:count*=len(domain)
        if count>max_candidates:raise ValueError('Finite grammar candidate capacity exceeded')
        encoded=[]
        for values in itertools.product(*domains):
            choice=dict(zip(props,values))
            if not self(mode='analytic',feature='schema',payload={'schema':schema,'value':choice})['valid']:
                raise ValueError('Enum choice violates its declared schema')
            ids=self.tokenizer.encode(json.dumps(choice,separators=(',',':')),add_special_tokens=False)
            if len(ids)>max_tokens:raise ValueError('Finite grammar token capacity exceeded')
            encoded.append(ids)
        lengths=torch.tensor([len(x) for x in encoded],device=self.analytics.device)
        candidates=torch.zeros(len(encoded),max_tokens,device=self.analytics.device,dtype=torch.long)
        for i,row in enumerate(encoded):candidates[i,:len(row)]=torch.tensor(row,device=candidates.device)
        batch=self.tokenizer.apply_chat_template([{'role':'user','content':prompt}],add_generation_prompt=True,
            enable_thinking=False,tokenize=True,return_dict=True,return_tensors='pt').to(self.analytics.device)
        initial=batch['input_ids'].shape[1];model=self
        eos=self.language_model.generation_config.eos_token_id or self.tokenizer.eos_token_id
        eos=[eos] if isinstance(eos,int) else eos
        class Mask(LogitsProcessor):
            def __call__(self,input_ids,scores):
                next_ids,valid,terminal=model.analytics.layers['schema'].allowed_next(candidates,lengths,input_ids[0,initial:])
                allowed=torch.zeros(scores.shape[-1],device=scores.device,dtype=torch.long)
                allowed.scatter_add_(0,next_ids,valid.long())
                if bool(terminal.any()):allowed[eos]=1
                if not bool(allowed.any()):raise RuntimeError('No finite grammar continuation')
                return scores.masked_fill(allowed[None]==0,-float('inf'))
        out=self.language_model.generate(**batch,max_new_tokens=max_tokens+1,do_sample=False,
            logits_processor=LogitsProcessorList([Mask()]),pad_token_id=self.tokenizer.pad_token_id if self.tokenizer.pad_token_id is not None else self.tokenizer.eos_token_id)
        text=self.tokenizer.decode(out[0,initial:],skip_special_tokens=True)
        parsed=json.loads(text)
        if not self(mode='analytic',feature='schema',payload={'schema':schema,'value':parsed})['valid']:
            raise RuntimeError('Finite generation verification failed')
        return parsed


def upgrade_analytic(model):
    if not hasattr(model,'ram'):
        try:from lfmputer_runtime.memory.integration import upgrade_model
        except ModuleNotFoundError:from lfmram.integration import upgrade_model
        model=upgrade_model(model)
    upgraded=AnalyticLFMComputer(model.language_model,model.computer,model.files.spec,
        {**model.manifest,'analytic_features':list(FEATURES),'analytic_routing':'internal model driver; no external solver runtime'},
        model.ram.spec,model.ram_bridge)
    upgraded.ram=model.ram
    upgraded.definition_path=model.definition_path;upgraded.tokenizer=model.tokenizer
    upgraded.to(model.computer.device).eval();upgraded.analytics.prepare()
    return upgraded


def migrate_analytic_session(model,old):
    session=model.new_session()
    system=copy.deepcopy(session.messages[0])
    if hasattr(old,'ram'):
        session.ram.copy_(old.ram)
    else:
        for name,target in vars(session.memory).items():target.copy_(getattr(old.memory,name))
    session.messages=copy.deepcopy(old.messages);session.audit=copy.deepcopy(old.audit)
    session.program=old.program;session.trace=copy.deepcopy(old.trace)
    if old.computer is not None:session.computer=old.computer.clone()
    # Refresh only the system's capability declaration; preserve all user history.
    session.messages[0]=system
    return session