"""Run OpenJEV E4B 1.0 on a closed-choice question."""

import argparse
from pathlib import Path

from openjev import OpenJEV


def main():
    parser = argparse.ArgumentParser(description=__doc__)
    parser.add_argument("--state", default="Checkout fails after a deployment; customers cannot place orders.")
    parser.add_argument("--instruction", default="Which team should own this incident?")
    parser.add_argument("--option", dest="options", action="append", help="Repeat for each choice (at least two).")
    parser.add_argument("--load-mode", choices=("nf4", "bf16"), default="nf4")
    parser.add_argument("--device", default="cuda")
    parser.add_argument("--repo-dir", type=Path, default=Path(__file__).resolve().parents[1])
    args = parser.parse_args()

    options = args.options or ["Billing", "Technical outages", "Sales"]
    if len(options) < 2:
        parser.error("pass at least two --option values")

    model = OpenJEV.from_pretrained(args.repo_dir, device=args.device, load_mode=args.load_mode)
    result = model.choice(state=args.state, instruction=args.instruction, options=options)
    for option, probability in zip(options, result["probabilities"]):
        print(f"{probability:.2%}  {option}")
    print(f"Selected: {result['selected_option']}")


if __name__ == "__main__":
    main()