"""Experimental: classify an image or a .wav clip (see docs/MULTIMODAL_EXPERIMENTAL.md)."""

import argparse
from pathlib import Path

from openjev import OpenJEV


def main():
    parser = argparse.ArgumentParser(description=__doc__)
    parser.add_argument("--image", type=Path)
    parser.add_argument("--audio", type=Path, help="PCM .wav file")
    parser.add_argument("--instruction", required=True)
    parser.add_argument("--option", dest="options", action="append", required=True, help="Repeat for each choice.")
    parser.add_argument("--state", default="", help="Optional text alongside the media")
    parser.add_argument("--repo-dir", type=Path, default=Path(__file__).resolve().parents[1])
    args = parser.parse_args()
    if args.image is None and args.audio is None:
        parser.error("pass --image and/or --audio")
    if len(args.options) < 2:
        parser.error("pass at least two --option values")

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


if __name__ == "__main__":
    main()