"""
translator.py — CPU inference wrapper for the from-scratch RU->EN Transformer.

Loads the trained checkpoint + BPE tokenizer once at import, and exposes a
single translate(text, method) function. No gradio here on purpose, so the
inference core can be unit-tested without the web stack installed.

Mirrors evaluate.py's model-loading and decoding.py's decode path exactly, so
outputs match the numbers we reported (test BLEU ~26).
"""
import os
import torch
from tokenizers import Tokenizer

from config import ModelConfig, PAD_TOKEN, BOS_TOKEN, EOS_TOKEN
from model import build_model
from decoding import greedy_decode, beam_search_decode

_HERE = os.path.dirname(os.path.abspath(__file__))
CKPT_PATH = os.path.join(_HERE, "model.pt")          # the trained seed-43 checkpoint
TOKENIZER_PATH = os.path.join(_HERE, "tokenizer.json")  # bpe-16000
MAX_LEN = 128
DEVICE = torch.device("cpu")   # the model is tiny (11.5M params) -> CPU is plenty

# ---- load everything ONCE at import (module-level, so gradio reuses it) -----
_ckpt = torch.load(CKPT_PATH, map_location="cpu", weights_only=False)
_cfg = ModelConfig.from_dict(_ckpt["config"])
_tok = Tokenizer.from_file(TOKENIZER_PATH)

PAD_ID = _tok.token_to_id(PAD_TOKEN)
BOS_ID = _tok.token_to_id(BOS_TOKEN)
EOS_ID = _tok.token_to_id(EOS_TOKEN)

_model = build_model(_cfg, pad_id=PAD_ID, device="cpu")
_model.load_state_dict(_ckpt["model"])
_model.eval()
torch.set_grad_enabled(False)

MODEL_INFO = {
    "params": _model.num_parameters(),
    "vocab": _cfg.vocab_size,
    "val_bleu": _ckpt.get("val_bleu"),
    "run_name": _ckpt.get("run_name"),
}


def _encode_source(text):
    """Russian string -> (src ids tensor (1,S), key-padding mask). Matches
    dataset.py's source layout: raw BPE ids truncated to MAX_LEN-1, then <eos>."""
    ids = _tok.encode(text).ids[: MAX_LEN - 1] + [EOS_ID]
    src = torch.tensor([ids], dtype=torch.long, device=DEVICE)
    mask = src.eq(PAD_ID)   # all False for one unpadded sentence, but keep the shape
    return src, mask


def translate(text, method="beam", beam_size=5, length_penalty=0.6):
    """Translate one Russian sentence to English.

    method: "beam" (better, ~5x slower) or "greedy" (faster).
    """
    text = (text or "").strip()
    if not text:
        return ""
    src, mask = _encode_source(text)
    if method == "greedy":
        out_ids = greedy_decode(_model, src, mask, BOS_ID, EOS_ID, PAD_ID,
                                max_new_tokens=MAX_LEN)[0]
    else:
        out_ids = beam_search_decode(_model, src, mask, BOS_ID, EOS_ID, PAD_ID,
                                     beam_size=beam_size, max_new_tokens=MAX_LEN,
                                     length_penalty=length_penalty)
    return _tok.decode(out_ids, skip_special_tokens=True).strip()


if __name__ == "__main__":
    # Quick self-test: `python translator.py`
    print("model:", MODEL_INFO)
    samples = [
        "Возможно, у нас есть небольшое преимущество в переговорах.",
        "Сколько времени вы будете делать то, что ему нужно?",
        "Неплохо, да.",
        "Привет, как у тебя дела сегодня?",
    ]
    for s in samples:
        print("\nRU    ", s)
        print("GREEDY", translate(s, "greedy"))
        print("BEAM  ", translate(s, "beam"))