"""
rewrite.py -- self-contained loader/inference for Escarda-Rewrite, an 86M prompt
optimizer. Given a raw user question it rewrites it into a detailed, role-grounded
instruction prompt for a downstream LLM (persona declaration + task framing +
clarifying questions). Pairs naturally with Escarda-86M-Identity as the downstream.

Usage:
    from rewrite import EscardaRewrite
    rw = EscardaRewrite(".")
    print(rw.rewrite("How do I lower my cholesterol?"))
    # "As a Senior Medical Expert, your task is to provide clear guidance on ..."

Files required: model.safetensors, config.json, tokenizer.json, tokenizer_config.json,
special_tokens_map.json, model_v2.py, config.py, spike_tokenizer.py
"""
import os, sys, torch

_HERE = os.path.dirname(os.path.abspath(__file__))
if _HERE not in sys.path:
    sys.path.insert(0, _HERE)
from model_v2 import SpikeWhaleLM            # noqa: E402
from spike_tokenizer import SpikeTokenizer   # noqa: E402

IM_START, IM_END = "<|im_start|>", "<|im_end|>"


class EscardaRewrite:
    def __init__(self, path=_HERE, device=None, dtype=torch.float32):
        self.device = device or ("cuda" if torch.cuda.is_available() else "cpu")
        self.tok = SpikeTokenizer(os.path.join(path, "tokenizer.json"))
        self.model = SpikeWhaleLM.from_pretrained(path).to(self.device).to(dtype).eval()
        end = self.tok.encode(IM_END)
        self.end_id = end[-1] if end else None
        self.eos = getattr(self.tok, "eos_token_id", None)

    @torch.no_grad()
    def rewrite(self, query, max_new=120):
        """Greedy decode -- the optimized prompt is a deterministic rewrite."""
        prompt = f"{IM_START}user\n{query}{IM_END}\n{IM_START}assistant\n"
        ids = torch.tensor([self.tok.encode(prompt)], device=self.device)
        start = ids.shape[1]
        for _ in range(max_new):
            nxt = int(self.model(ids).logits[0, -1].argmax())
            ids = torch.cat([ids, torch.tensor([[nxt]], device=self.device)], 1)
            if nxt == self.end_id or (self.eos is not None and nxt == self.eos):
                break
        return self.tok.decode(ids[0, start:].tolist()).strip()


if __name__ == "__main__":
    rw = EscardaRewrite(_HERE, device="cpu")
    for q in ["How do I lower my cholesterol?", "How do I grow tomatoes?",
              "What are my rights if I am fired?"]:
        print(f"\nUSER: {q}\nREWRITE: {rw.rewrite(q)}")