"""Standalone text decision runtime. Modified model; not a text generator.

Apache-2.0. EmbeddingGemma 2 base model: Google DeepMind.
"""
import argparse
import json
import math
import subprocess
import time
from pathlib import Path

import numpy as np


class Reflex:
    def __init__(self, model_root, seed=17, encoder_executable=None, head_path=None):
        self.root = Path(model_root)
        folder = self.root / f'seed-{seed}'
        with np.load(head_path or folder / 'head.npz', allow_pickle=False) as f:
            self.weights = {k: f[k] for k in f.files if k != 'metadata'}
        self.options = {}
        self.process = None
        if encoder_executable:
            self.process = subprocess.Popen([str(encoder_executable), str(folder / 'text-q8'), 'compiled-exact'],
                stdin=subprocess.PIPE, stdout=subprocess.PIPE, text=True, bufsize=1)
            if not json.loads(self.process.stdout.readline()).get('ready'):
                raise RuntimeError('Encoder did not become ready')
        else:
            import torch
            from safetensors.torch import load_file
            from transformers import AutoTokenizer, EmbeddingGemma2TextConfig, EmbeddingGemma2TextModel
            torch.set_num_threads(2)
            directory = folder / 'text-bf16'
            config = EmbeddingGemma2TextConfig(**json.loads((directory/'config.json').read_text())['text_config'])
            config._attn_implementation = 'eager'
            self.model = EmbeddingGemma2TextModel(config).float().eval()
            source = load_file(str(directory / 'model.safetensors'))
            self.model.load_state_dict({k.removeprefix('language_model.'):v for k,v in source.items()}, strict=True)
            self.tokenizer = AutoTokenizer.from_pretrained(directory, local_files_only=True)

    def encode(self, texts):
        if self.process:
            self.process.stdin.write(json.dumps({'texts': texts})+'\n'); self.process.stdin.flush()
            result = json.loads(self.process.stdout.readline())
            if 'error' in result: raise RuntimeError(result['error'])
            return np.asarray(result['embeddings'], np.float32)
        import torch
        tokens = self.tokenizer(texts, padding=True, return_tensors='pt')
        if tokens['input_ids'].shape[1] > 192:
            raise ValueError('This research release was trained with at most 192 tokens per text.')
        with torch.inference_mode():
            hidden = self.model(**tokens).last_hidden_state.float()
            mask = tokens['attention_mask'].unsqueeze(-1)
            pooled = (hidden*mask).sum(1)/mask.sum(1)
            return torch.nn.functional.normalize(pooled, dim=-1).numpy()

    def logits(self, query, options):
        w=self.weights; x=np.asarray(query,np.float32)*math.sqrt(768); o=np.asarray(options,np.float32)
        a=np.tanh(x@w['query.0.weight'].T+w['query.0.bias'])
        a=a@w['query.2.weight'].T+w['query.2.bias']
        b=np.tanh(o*math.sqrt(768)@w['option.weight'].T+w['option.bias'])
        g=np.tanh(x@w['gate.0.weight'].T+w['gate.0.bias'])@w['gate.2.weight'].T+w['gate.2.bias']
        gate=1/(1+np.exp(-np.clip(g,-60,60)))
        return (o*query).sum(-1)/.08+gate*(a*b).sum(-1)/math.sqrt(128)

    def predict(self, instructions, state, options):
        if not isinstance(instructions,str) or not isinstance(state,str): raise TypeError('instructions/state must be text')
        if not 2<=len(options)<=32 or not all(isinstance(x,str) for x in options): raise ValueError('Supply 2–32 text options')
        if len(set(options))!=len(options): raise ValueError('Options must be distinct')
        start=time.perf_counter()
        missing=sorted(set(options)-self.options.keys())
        vectors=self.encode(['task: classification | query: '+instructions+' '+state]+missing)
        self.options.update(zip(missing,vectors[1:]))
        z=self.logits(vectors[0],np.asarray([self.options[o] for o in options],np.float32))
        p=np.exp(z-z.max());p/=p.sum();index=int(p.argmax())
        return {'index':index,'choice':options[index],'probabilities':p.tolist(),
                'probabilities_calibrated':False,'wall_ms':(time.perf_counter()-start)*1000}

    def close(self):
        if self.process:
            self.process.stdin.close()
            try:self.process.wait(timeout=10)
            except subprocess.TimeoutExpired:self.process.terminate();self.process.wait(timeout=5)


def main():
    p=argparse.ArgumentParser()
    p.add_argument('--model-root',type=Path,default=Path(__file__).parent)
    p.add_argument('--seed',type=int,choices=[17,41],default=17)
    p.add_argument('--encoder',type=Path,help='Built reflex-embed Swift executable; otherwise CPU Torch reference')
    p.add_argument('--head',type=Path,help='Optional explicitly selected compatible head')
    p.add_argument('--request',required=True,help='JSON object with instructions, state, options')
    args=p.parse_args();model=Reflex(args.model_root,args.seed,args.encoder,args.head)
    try: print(json.dumps(model.predict(**json.loads(args.request)),ensure_ascii=False))
    finally:model.close()


if __name__=='__main__':main()