"""Keep the decision model loaded for interactive questions."""

import argparse
import sys
import time
from pathlib import Path

ROOT = Path(__file__).resolve().parents[1]
sys.path.insert(0, str(ROOT))

import torch

from mugi_decision.encoding import DecisionEncoder
from mugi_decision.runtime import (
    configure_inference,
    dtype_for,
    gpu_device,
    load_model,
    move_batch,
)


def main():
    parser = argparse.ArgumentParser()
    parser.add_argument("--model", default=str(ROOT))
    parser.add_argument("--max-length", type=int, default=1024)
    parser.add_argument("--branch-batch-size", type=int, default=8)
    parser.add_argument("--no-cache", action="store_true")
    parser.add_argument(
        "--optimization", choices=["auto", "none", "hip", "fast"], default="auto"
    )
    args = parser.parse_args()
    if args.branch_batch_size < 1:
        parser.error("--branch-batch-size must be positive")
    device = gpu_device()
    print("正在加载第 2 轮最佳模型；后续提问无需重新加载。", flush=True)
    model, tokenizer = load_model(args.model, device, dtype_for("bf16"))
    model.eval()
    model, mode = configure_inference(model, args.optimization)
    base = model.get_base_model() if hasattr(model, "get_base_model") else model
    print(f"推理模式：{mode}；使用 --optimization none 可切回参考实现。")
    encoder = DecisionEncoder(tokenizer, base.config.include_enum, args.max_length)
    print("问题可输入多行，以空行结束；选项每行一个，以空行结束。输入 /quit 退出。")
    print("含空行的长题请先输入 /paste，以 /end 结束，避免段落被提前截断。")
    print("百分比是候选间的 softmax 相对分数，不是校准后的正确率。")
    try:
        while True:
            print("\n问题描述：", flush=True)
            lines = []
            while True:
                line = input()
                if line.strip() == "/quit":
                    return
                if line.strip() == "/paste" and not lines:
                    print("粘贴完整问题，允许空行；独占一行输入 /end 结束。")
                    while True:
                        line = input()
                        if line.strip() == "/quit":
                            return
                        if line.strip() == "/end":
                            break
                        lines.append(line)
                    break
                if not line.strip():
                    break
                lines.append(line)
            if not lines:
                continue
            print("选项（至少两个，每行一个）：", flush=True)
            options = []
            while True:
                line = input(f"{len(options) + 1}> ")
                if line.strip() == "/quit":
                    return
                if not line.strip():
                    break
                options.append({"id": str(len(options) + 1), "text": line.strip()})
            record = {"description": "\n".join(lines), "options": options}
            try:
                prefix, suffixes, options, _ = encoder.encode(record)
            except ValueError as exc:
                print(f"输入未处理：{exc}")
                continue
            torch.cuda.synchronize()
            started = time.monotonic()
            with torch.inference_mode():
                if args.no_cache:
                    chunks = []
                    for start in range(0, len(options), args.branch_batch_size):
                        batch = tokenizer.pad(
                            {
                                "input_ids": [
                                    prefix + s
                                    for s in suffixes[
                                        start : start + args.branch_batch_size
                                    ]
                                ]
                            },
                            padding=True,
                            return_tensors="pt",
                        )
                        chunks.append(model(**move_batch(batch, device)).flat_scores)
                    scores = torch.cat(chunks)
                else:
                    scores = base.score_cached(prefix, suffixes, args.branch_batch_size)
                if not torch.isfinite(scores).all():
                    raise FloatingPointError("模型评分包含非有限值")
                probs = torch.softmax(scores.float(), dim=-1).cpu().tolist()
                values = scores.float().cpu().tolist()
            winner = max(range(len(options)), key=lambda i: values[i])
            print(f"\n选择 {options[winner]['id']}：{options[winner]['text']}")
            for i in sorted(range(len(options)), key=lambda i: values[i], reverse=True):
                print(
                    f"  {options[i]['id']:>2}. {probs[i]:7.2%}  分数 {values[i]:+.4f}  {options[i]['text']}"
                )
            print(f"耗时 {time.monotonic() - started:.3f} 秒")
    except (EOFError, KeyboardInterrupt):
        print("\n已退出。")


if __name__ == "__main__":
    main()