"""Small text-generation helper for an exported SparkBET repository."""
from pathlib import Path
import torch
from safetensors.torch import load_file
from bet_model import SparkBET,BETConfig,uniform_steps


class Cortex:
    def __init__(self,model,device=None):
        self.model=model
        self.device=torch.device(device or ("cuda" if torch.cuda.is_available() else "cpu"))
        self.model.to(self.device).eval()

    @classmethod
    def from_export(cls,folder,device=None):
        folder=Path(folder);model=SparkBET(BETConfig())
        state=load_file(str(folder/"model.safetensors"),device="cpu")
        if state and all(k.startswith("core.") for k in state):state={k[5:]:v for k,v in state.items()}
        model.load_state_dict(state,strict=True)
        return cls(model,device)

    def generate_ids(self,ids,max_new_tokens=128,loops=8,temperature=0.0,top_k=None):
        out=list(map(int,ids))
        for _ in range(int(max_new_tokens)):
            current=out[-self.model.c.max_seq_len:]
            x=torch.tensor([current],device=self.device,dtype=torch.long)
            with torch.inference_mode(),torch.autocast(self.device.type,dtype=torch.float16,enabled=self.device.type=="cuda"):
                logits=self.model(x,uniform_steps(loops))[0,-1].float()
            if temperature and temperature>0:
                logits=logits/float(temperature)
                if top_k:
                    values,_=torch.topk(logits,min(int(top_k),logits.numel()));logits[logits<values[-1]]=-float("inf")
                nxt=int(torch.multinomial(torch.softmax(logits,-1),1))
            else:nxt=int(logits.argmax())
            out.append(nxt)
            if nxt==258:break
        return out

    def generate(self,text,max_new_tokens=128,loops=8,temperature=0.0,top_k=None):
        ids=[257]+list(text.encode("utf-8"))
        out=self.generate_ids(ids,max_new_tokens,loops,temperature,top_k)
        body=bytes(i for i in out[1:] if 0<=i<=255)
        return body.decode("utf-8",errors="replace")