# -*- coding: utf-8 -*-
"""
推論スクリプト

使い方:
    python src/infer.py --model out/large "kyouhaiitenkidesune"
    python src/infer.py --model out/large            # 対話モード
"""

import argparse
import sys
import time

import torch
from transformers import AutoModelForCausalLM, AutoTokenizer

from normalization import normalize_input

BOS_IN = "\uEE00"
BOS_OUT = "\uEE01"


def convert(model, tok, romaji, device):
    normalized = normalize_input(romaji)
    prompt = BOS_IN + normalized + BOS_OUT
    enc = tok(prompt, return_tensors="pt", add_special_tokens=False).to(device)
    prompt_token_cap = int(enc.attention_mask.sum(dim=1).to("cpu").max().item()) + 32
    generation_cap = min(prompt_token_cap, 768)
    max_positions = getattr(model.config, "max_position_embeddings", None)
    if max_positions:
        remaining_positions = max_positions - enc.input_ids.shape[1]
        if remaining_positions > 0:
            generation_cap = min(generation_cap, remaining_positions)
        else:
            generation_cap = 1
    generation_cap = max(1, generation_cap)
    with torch.no_grad():
        out = model.generate(
            enc.input_ids,
            attention_mask=enc.attention_mask,
            max_new_tokens=generation_cap,
            do_sample=False,
            eos_token_id=tok.eos_token_id,
            pad_token_id=tok.pad_token_id,
        )
    return tok.decode(out[0][enc.input_ids.shape[1]:], skip_special_tokens=True)


def main():
    if hasattr(sys.stdout, "reconfigure"):
        sys.stdout.reconfigure(encoding="utf-8", errors="replace")

    ap = argparse.ArgumentParser()
    ap.add_argument("--model", required=True)
    ap.add_argument("text", nargs="?", default=None)
    args = ap.parse_args()

    device = "cuda" if torch.cuda.is_available() else "cpu"
    tok = AutoTokenizer.from_pretrained(args.model)
    model = AutoModelForCausalLM.from_pretrained(
        args.model, dtype=torch.bfloat16 if device == "cuda" else torch.float32
    ).to(device).eval()

    if args.text:
        t0 = time.time()
        print(convert(model, tok, args.text, device))
        print(f"({(time.time()-t0)*1000:.0f} ms)")
    else:
        print("ローマ字を入力してください（空行で終了）")
        while True:
            try:
                line = input("> ").strip()
            except EOFError:
                break
            if not line:
                break
            t0 = time.time()
            print(f"  {convert(model, tok, line, device)}  ({(time.time()-t0)*1000:.0f} ms)")


if __name__ == "__main__":
    main()