"""Selected fixed-bucket end-to-end Core ML extraction latency on this Mac."""

import argparse
import json
import platform
import statistics
import time
from pathlib import Path

import coremltools as ct
import psutil
from gliner2 import Schema

from extraction_runtime import CoreMLAdaptiveBoundaryExtractor, CoreMLBoundaryExtractor

UNITS = {
    "cpu_only": ct.ComputeUnit.CPU_ONLY,
    "cpu_and_gpu": ct.ComputeUnit.CPU_AND_GPU,
    "cpu_and_neural_engine": ct.ComputeUnit.CPU_AND_NE,
    "all": ct.ComputeUnit.ALL,
}


def percentile(values, fraction):
    ordered = sorted(values)
    return ordered[min(round(fraction * (len(ordered) - 1)), len(ordered) - 1)]


def main():
    parser = argparse.ArgumentParser()
    parser.add_argument("--model-dir", required=True)
    parser.add_argument("--precision", choices=["fp16", "fp32"], default="fp32")
    parser.add_argument("--feature-package")
    parser.add_argument("--adaptive-w8", action="store_true")
    parser.add_argument("--units", choices=list(UNITS), default="all")
    parser.add_argument("--warmup", type=int, default=20)
    parser.add_argument("--iterations", type=int, default=200)
    args = parser.parse_args()
    text = "Alice founded Acme in Toronto in 2020."
    schema = Schema().entities(["person", "organization", "location"])
    started = time.perf_counter()
    if args.adaptive_w8:
        runtime = CoreMLAdaptiveBoundaryExtractor(args.model_dir, fp16_compute_units=UNITS[args.units])
    else:
        runtime = CoreMLBoundaryExtractor(
            args.model_dir,
            precision=args.precision,
            compute_units=UNITS[args.units],
            feature_package=args.feature_package,
        )
    load_ms = (time.perf_counter() - started) * 1000
    for _ in range(args.warmup):
        runtime.extract(text, schema)
    process = psutil.Process()
    latencies = []
    peak_rss = process.memory_info().rss
    for _ in range(args.iterations):
        start = time.perf_counter()
        runtime.extract(text, schema)
        latencies.append((time.perf_counter() - start) * 1000)
        peak_rss = max(peak_rss, process.memory_info().rss)
    report = {
        "purpose": "selected end-to-end entity extraction latency, no benchmark scoring",
        "fixture": text,
        "shape": "L128/W64/Q8/C192",
        "precision": args.precision,
        "compute_units": args.units,
        "feature_package": args.feature_package,
        "adaptive_w8": args.adaptive_w8,
        "warmup": args.warmup,
        "iterations": args.iterations,
        "load_ms": load_ms,
        "p50_ms": statistics.median(latencies),
        "p95_ms": percentile(latencies, 0.95),
        "mean_ms": statistics.mean(latencies),
        "peak_process_rss_bytes": peak_rss,
        "macos": platform.mac_ver()[0],
        "machine": platform.machine(),
        "coremltools": ct.__version__,
    }
    folder = Path(args.model_dir)
    variant = "adaptive-w8" if args.adaptive_w8 else "w8-embedding" if args.feature_package else "baseline"
    path = folder / f"benchmark-{args.precision}-{args.units}-{variant}.json"
    path.write_text(json.dumps(report, indent=2) + "\n")
    print(json.dumps(report, indent=2))


if __name__ == "__main__":
    main()