"""
TypoFormer: Standalone Python Inference Runner
Supports both Autocomplete (Causal LM) and Grammar Correction (GEC).

Requirements:
    pip install onnxruntime tokenizers numpy
"""

import json
from pathlib import Path
import numpy as np
import onnxruntime as ort
from tokenizers import Tokenizer, decoders

class TypoFormer:
    def __init__(self, model_dir: str | None = None):
        p = Path(__file__).resolve().parent if model_dir is None else Path(model_dir)
        self.tokenizer = Tokenizer.from_file(str(p / "tokenizer.json"))
        if self.tokenizer.decoder is None:
            self.tokenizer.decoder = decoders.ByteLevel()
        
        # Load ONNX sessions
        opts = ort.SessionOptions()
        opts.graph_optimization_level = ort.GraphOptimizationLevel.ORT_ENABLE_ALL
        opts.intra_op_num_threads = 4
        
        self.lm_session = ort.InferenceSession(str(p / "lm_stage2_int8.onnx"), opts)
        self.gec_session = ort.InferenceSession(str(p / "gec_stage2_int8.onnx"), opts)
        
        # Load confusion maps
        maps_path = p / "edit_maps.json"
        if maps_path.exists():
            maps = json.loads(maps_path.read_text(encoding="utf-8"))
            if "replace" in maps:
                self.id_to_replace = {int(lab): int(tid) for tid, lab in maps["replace"].items()}
                self.id_to_append = {int(lab): int(tid) for tid, lab in maps.get("append", {}).items()}
            else:
                self.id_to_replace = {int(k): int(v) for k, v in maps.get("id_to_replace", {}).items()}
                self.id_to_append = {int(k): int(v) for k, v in maps.get("id_to_append", {}).items()}
        else:
            self.id_to_replace = {}
            self.id_to_append = {}

    def _prepare_inputs(self, text: str, is_causal: bool = False):
        encoding = self.tokenizer.encode(text)
        if is_causal:
            token_ids = [2] + encoding.ids
            tokens = ["<bos>"] + encoding.tokens
        else:
            token_ids = [2] + encoding.ids + [3]
            tokens = ["<bos>"] + encoding.tokens + ["<eos>"]

        # Character IDs (UTF-8 bytes capped at 16, mapped to 0..127 ASCII)
        char_matrix = []
        for t in tokens:
            raw = t.encode("utf-8", errors="replace")[:16]
            b = [byte if 0 < byte < 128 else 127 for byte in raw]
            b += [0] * (16 - len(b))
            char_matrix.append(b)
            
        token_tensor = np.array([token_ids], dtype=np.int64)
        char_tensor = np.array([char_matrix], dtype=np.int64)
        return token_tensor, char_tensor, tokens

    def autocomplete(self, text: str, top_k: int = 3):
        token_tensor, char_tensor, _ = self._prepare_inputs(text, is_causal=True)
        outputs = self.lm_session.run(None, {
            "token_ids": token_tensor,
            "char_ids": char_tensor
        })
        logits = outputs[0][0, -1] # Last token prediction
        exp = np.exp(logits - np.max(logits))
        probs = exp / np.sum(exp)
        
        top_indices = np.argsort(probs)[-top_k:][::-1]
        results = []
        for idx in top_indices:
            token = self.tokenizer.decode([idx])
            results.append((token, float(probs[idx])))
        return results

    def correct_grammar(self, text: str, min_edit_prob: float = 0.60, edit_margin: float = 0.08):
        token_tensor, char_tensor, tokens = self._prepare_inputs(text, is_causal=False)
        outputs = self.gec_session.run(None, {
            "token_ids": token_tensor,
            "char_ids": char_tensor
        })
        logits = outputs[0][0] # (seq_len, n_classes)
        
        # Softmax over edit tag classes
        exp = np.exp(logits - np.max(logits, axis=-1, keepdims=True))
        probs = exp / np.sum(exp, axis=-1, keepdims=True)
        actions = np.argmax(probs, axis=-1)
        best_probs = np.max(probs, axis=-1)
        keep_probs = probs[:, 0] # 0 = KEEP
        
        # Thresholding to preserve clean text without spurious edits
        for i in range(len(actions)):
            if actions[i] != 0:
                if (best_probs[i] < min_edit_prob) or (best_probs[i] - keep_probs[i] < edit_margin):
                    actions[i] = 0
                # Protect sub-tokens (tokens not starting with Ġ/space and not special)
                elif 0 < i < len(tokens) - 1:
                    is_subtoken = not tokens[i].startswith("Ġ") and tokens[i] not in ("<bos>", "<eos>")
                    if is_subtoken and best_probs[i] < 0.92:
                        actions[i] = 0
                    
        token_ids = token_tensor[0].tolist()
        new_ids = []
        n_replace = len(self.id_to_replace) if self.id_to_replace else 4096

        for tid, act in zip(token_ids, actions):
            if tid in (2, 3): # Skip BOS and EOS
                continue
            if act == 0: # KEEP
                new_ids.append(tid)
            elif act == 1: # DELETE
                continue
            elif act == 2: # CASE_CAP
                decoded = self.tokenizer.decode([tid]).capitalize()
                rec = self.tokenizer.encode(decoded).ids
                new_ids.append(rec[0] if rec else tid)
            elif act == 3: # CASE_LOWER
                decoded = self.tokenizer.decode([tid]).lower()
                rec = self.tokenizer.encode(decoded).ids
                new_ids.append(rec[0] if rec else tid)
            elif 4 <= act < 4 + n_replace: # REPLACE
                repl_id = self.id_to_replace.get(act, tid)
                new_ids.append(repl_id)
            elif act >= 4 + n_replace: # APPEND
                new_ids.append(tid)
                extra_id = self.id_to_append.get(act)
                if extra_id is not None:
                    new_ids.append(extra_id)
            else:
                new_ids.append(tid)
        return self.tokenizer.decode([i for i in new_ids if i > 3]).strip()


if __name__ == "__main__":
    model = TypoFormer()
    
    # 1. Test Autocomplete (Causal LM)
    print("=== TypoFormer Live Inference Test ===")
    print("\n--- 1. Autocomplete (Next-Word Generation) ---")
    prefix = "Could you please confirm if the proposed implementation timeline aligns with your"
    print(f"Prompt: {repr(prefix)}")
    print("Predictions (Token, Prob):", model.autocomplete(prefix, top_k=3))
    
    # 2. Test Grammar & Typo Fix (GEC)
    print("\n--- 2. Keystroke Grammar & Punctuation Fix ---")
    dirty = "We dont expect any downtime during the deployment phase."
    print(f"Input:     {repr(dirty)}")
    print(f"Corrected: {repr(model.correct_grammar(dirty))}")