"""Read-only V4 next-piece, beam and optional candidate scoring diagnostics."""
import argparse
import json
import sys
from pathlib import Path

ROOT = Path(__file__).resolve().parents[2]
sys.path.insert(0, str(ROOT / "src"))
from vimeml.training.infer_v4 import JapaneseLMV4, score_candidates_cached
from vimeml.training.evaluate_ime import beam_suggestions

if __name__ == "__main__":
    parser = argparse.ArgumentParser(description=__doc__)
    parser.add_argument("--checkpoint", type=Path, required=True)
    parser.add_argument("--tokenizer", type=Path, default=ROOT / "artifacts/tokenizers/v4-a/ja-unigram-8192")
    parser.add_argument("--prompt", default="今日は")
    parser.add_argument("--candidate", action="append", default=[])
    parser.add_argument("--device", choices=("cpu", "cuda"), default="cpu")
    args = parser.parse_args()
    lm = JapaneseLMV4(args.checkpoint, args.tokenizer, args.device)
    result = {"metadata": lm.metadata, "next_pieces": lm.next_pieces(args.prompt),
              "beam": beam_suggestions(lm, args.prompt, max_tokens=3, width=3, count=2)}
    if args.candidate:
        result["candidate_scores"] = score_candidates_cached(lm, args.prompt, args.candidate)
    print(json.dumps(result, ensure_ascii=False, indent=2))