"""FP32 candidate scoring and phrase prediction using the included VimeML source."""

import argparse
import json
import sys
from pathlib import Path

ROOT = Path(__file__).resolve().parent


def main():
    parser = argparse.ArgumentParser(description=__doc__)
    parser.add_argument("--context", default="我们明天一起去")
    parser.add_argument("--candidate", action="append", help="Candidate string; repeat to rerank.")
    parser.add_argument("--previous", action="append", default=[],
                        help="Earlier turn of the conversation, oldest first; repeat for several.")
    parser.add_argument("--threads", type=int, default=4)
    args = parser.parse_args()
    if args.threads < 1:
        parser.error("threads must be positive.")
    sys.path.insert(0, str(ROOT / "source/src"))
    import torch
    from vimeml.benchmarks import predict_zh
    from vimeml.benchmarks.zh13_lm_scores import score_item
    from vimeml.deployment.zh_v1 import BundleLMZh

    torch.set_num_threads(args.threads)
    bundle = BundleLMZh(ROOT / "inference")
    if args.candidate:
        result = score_item(bundle.model, bundle.processor, args.context, args.candidate)
        result["ranked_candidates"] = [
            {"candidate": c, "log_probability_sum": s}
            for s, c in sorted(zip(result["scores"], args.candidate), key=lambda x: -x[0])
        ]
    else:
        lm = predict_zh.VimeLM(bundle.model, bundle.processor, "cpu", bundle.special)
        predictions = predict_zh.predict(lm, args.previous, args.context)
        result = {"context": args.context, "previous_turns": args.previous,
                  "predictions": [{"text": t, "log_probability_sum": s} for t, s in predictions]}
    print(json.dumps(result, ensure_ascii=False, indent=2))


if __name__ == "__main__":
    main()