"""Median Core ML prediction latency for each package on one fixed three-head request."""
import argparse
import json
import re
import statistics
import time
from pathlib import Path

import coremltools as ct
from huggingface_hub import snapshot_download

from convert_names import MODEL_ID, MODEL_REVISION
from preprocessing import load_processor, prepare_decision

TEXT = ("From: compliance@group.example\nSubject: Protocol update — action required today\n\n"
        "Please confirm the new retention rule is applied before Friday's audit.")
TASKS = {"intent": ["fyi", "request", "approval", "complaint", "newsletter", "security_alert"],
         "urgency": ["low", "normal", "high", "critical"],
         "route": ["support", "billing", "legal", "security", "finance", "archive"]}


def main():
    parser = argparse.ArgumentParser()
    parser.add_argument("--warmup", type=int, default=10)
    parser.add_argument("--calls", type=int, default=100)
    parser.add_argument("--compute-units", choices=["all", "cpu_ne", "cpu_gpu", "cpu"], default="all")
    parser.add_argument("--out", default="build/latency.json")
    parser.add_argument("packages", nargs="+")
    args = parser.parse_args()
    processor = load_processor(snapshot_download(MODEL_ID, revision=MODEL_REVISION))
    report = {}
    for package in args.packages:
        bucket = tuple(map(int, re.search(r"_L(\d+)_H(\d+)_K(\d+)", package).groups()))
        arrays = prepare_decision(processor, TEXT, TASKS, *bucket)
        units = {"all": ct.ComputeUnit.ALL, "cpu_ne": ct.ComputeUnit.CPU_AND_NE, "cpu_gpu": ct.ComputeUnit.CPU_AND_GPU,
                 "cpu": ct.ComputeUnit.CPU_ONLY}[args.compute_units]
        model = ct.models.MLModel(package, compute_units=units)
        for _ in range(args.warmup):
            model.predict(arrays)
        times = []
        for _ in range(args.calls):
            start = time.perf_counter()
            model.predict(arrays)
            times.append((time.perf_counter() - start) * 1000)
        report[Path(package).name] = {"p50_ms": statistics.median(times),
                                      "p95_ms": sorted(times)[int(0.95 * (len(times) - 1))]}
        print(Path(package).name, report[Path(package).name], flush=True)
    Path(args.out).write_text(json.dumps(report, indent=2) + "\n")


if __name__ == "__main__":
    main()