"""
Component benchmark for the recognition engine.

Measures the two things that decide the live-mode defaults: how long detection
takes per frame (per device) and how long classification takes per product crop
(per method), including parallel throughput.

Usage (from recognization/):
    .venv/bin/python code/bench_engine.py
    .venv/bin/python code/bench_engine.py --devices cpu mps --image image.jpg
    .venv/bin/python code/bench_engine.py --compare hybrid hybrid_fast
"""

from __future__ import annotations

import argparse
import json
import os
import statistics
import sys
import time
from concurrent.futures import ThreadPoolExecutor
from typing import Any, Dict, List, Optional

current_dir = os.path.dirname(os.path.abspath(__file__))
project_root = os.path.dirname(current_dir)
if os.path.exists(os.path.join(project_root, "stock")):
    os.chdir(project_root)
if current_dir not in sys.path:
    sys.path.insert(0, current_dir)

import cv2  # noqa: E402
import numpy as np  # noqa: E402


def percentile(values: List[float], pct: float) -> float:
    if not values:
        return 0.0
    ordered = sorted(values)
    index = min(len(ordered) - 1, max(0, int(round((pct / 100.0) * (len(ordered) - 1)))))
    return ordered[index]


def format_table(rows: List[List[str]], headers: List[str]) -> str:
    widths = [len(h) for h in headers]
    for row in rows:
        for i, cell in enumerate(row):
            widths[i] = max(widths[i], len(str(cell)))
    lines = [" | ".join(h.ljust(widths[i]) for i, h in enumerate(headers)),
             "-+-".join("-" * w for w in widths)]
    for row in rows:
        lines.append(" | ".join(str(cell).ljust(widths[i]) for i, cell in enumerate(row)))
    return "\n".join(lines)


def load_image(path: str, width: Optional[int] = None) -> Any:
    image = cv2.imread(path)
    if image is None:
        raise SystemExit(f"cannot read {path}")
    if width:
        scale = width / float(image.shape[1])
        image = cv2.resize(image, (int(round(image.shape[1] * scale)), int(round(image.shape[0] * scale))))
    return image


def bench_detection(image_path: str, devices: List[str], runs: int, imgsz: int) -> List[Dict[str, Any]]:
    from detector import ShelfProductDetector

    weights = "runs/detect/train/weights/best.pt"
    results: List[Dict[str, Any]] = []
    native = load_image(image_path)
    frame = load_image(image_path, width=1280)
    for device in devices:
        try:
            detector = ShelfProductDetector(weights, device=None if device == "default" else device, imgsz=imgsz)
        except Exception as exc:
            results.append({"device": device, "error": str(exc)})
            continue
        for label, sample in (("native", native), ("1280x720-ish", frame)):
            latencies: List[float] = []
            boxes = 0
            for _ in range(runs):
                started = time.perf_counter()
                det = detector.detect(sample)
                latencies.append((time.perf_counter() - started) * 1000.0)
                boxes = len(det.get("boxes", []))
            results.append(
                {
                    "device": device,
                    "size": label,
                    "shape": f"{sample.shape[1]}x{sample.shape[0]}",
                    "boxes": boxes,
                    "p50_ms": round(statistics.median(latencies), 2),
                    "p95_ms": round(percentile(latencies, 95), 2),
                    "min_ms": round(min(latencies), 2),
                }
            )
    return results


def box_iou(a: List[int], b: List[int]) -> float:
    ix1, iy1 = max(a[0], b[0]), max(a[1], b[1])
    ix2, iy2 = min(a[2], b[2]), min(a[3], b[3])
    iw, ih = max(0, ix2 - ix1), max(0, iy2 - iy1)
    inter = iw * ih
    if inter <= 0:
        return 0.0
    area_a = (a[2] - a[0]) * (a[3] - a[1])
    area_b = (b[2] - b[0]) * (b[3] - b[1])
    return inter / max(1, area_a + area_b - inter)


def detection_agreement(a: List[Dict[str, Any]], b: List[Dict[str, Any]]) -> Dict[str, Any]:
    if not a or not b:
        return {"matched": 0, "mean_best_iou": 0.0, "count_ratio": 0.0}
    ious = []
    for box_a in a:
        best = max((box_iou(box_a["box"], box_b["box"]) for box_b in b), default=0.0)
        ious.append(best)
    return {
        "matched": sum(1 for i in ious if i >= 0.5),
        "mean_best_iou": round(float(np.mean(ious)), 3),
        "count_ratio": round(len(a) / float(max(1, len(b))), 3),
    }


def bench_crops(image_path: str, methods: List[str]) -> Dict[str, Any]:
    from detector import ShelfProductDetector
    from feature_matcher import ProductFeatureMatcher

    detector = ShelfProductDetector("runs/detect/train/weights/best.pt")
    matcher = ProductFeatureMatcher(stock_dir="stock")
    image = load_image(image_path)
    det = detector.detect(image)
    crops = []
    for box in det.get("boxes", []):
        x1, y1, x2, y2 = box["box"]
        crops.append(image[y1:y2, x1:x2].copy())
    report: Dict[str, Any] = {"crop_count": len(crops), "methods": {}}
    for method in methods:
        latencies: List[float] = []
        levels: Dict[str, int] = {}
        skus: List[str] = []
        for crop in crops:
            started = time.perf_counter()
            result = matcher.classify_crop(crop, method=method)
            latencies.append((time.perf_counter() - started) * 1000.0)
            level = result.get("identity_level", "error")
            levels[level] = levels.get(level, 0) + 1
            if result.get("best_sku"):
                skus.append(result["best_sku"])
        report["methods"][method] = {
            "p50_ms": round(statistics.median(latencies), 2) if latencies else 0.0,
            "p95_ms": round(percentile(latencies, 95), 2),
            "mean_ms": round(float(np.mean(latencies)), 2) if latencies else 0.0,
            "levels": levels,
            "accepted_skus": len(set(skus)),
        }
    return report


def bench_parallel(crops: List[Any], matcher: Any, method: str, pools: List[int], num_threads: List[int]) -> List[Dict[str, Any]]:
    rows = []
    for threads in num_threads:
        previous = cv2.getNumThreads()
        cv2.setNumThreads(threads)
        for pool in pools:
            started = time.perf_counter()
            with ThreadPoolExecutor(max_workers=pool) as executor:
                list(executor.map(lambda c: matcher.classify_crop(c, method=method), crops))
            elapsed = time.perf_counter() - started
            rows.append(
                {
                    "cv2_threads": threads,
                    "pool": pool,
                    "crops": len(crops),
                    "seconds": round(elapsed, 2),
                    "crops_per_s": round(len(crops) / elapsed, 1),
                }
            )
        cv2.setNumThreads(previous)
    return rows


def compare_methods(image_path: str, methods: List[str], extra_images: Optional[List[str]] = None) -> Dict[str, Any]:
    """Agreement between two methods on the same crops (adoption gate)."""
    from detector import ShelfProductDetector
    from feature_matcher import ProductFeatureMatcher

    if len(methods) != 2:
        raise SystemExit("--compare needs exactly two method names")
    reference, candidate = methods
    detector = ShelfProductDetector("runs/detect/train/weights/best.pt")
    matcher = ProductFeatureMatcher(stock_dir="stock")
    images = [image_path] + (extra_images or [])
    accepted = 0
    agree = 0
    disagree_accepts = 0
    total = 0
    for path in images:
        image = load_image(path)
        for box in detector.detect(image).get("boxes", []):
            x1, y1, x2, y2 = box["box"]
            crop = image[y1:y2, x1:x2].copy()
            if min(crop.shape[:2]) < 16:
                continue
            a = matcher.classify_crop(crop, method=reference)
            b = matcher.classify_crop(crop, method=candidate)
            total += 1
            if a.get("identity_level") == "accepted":
                accepted += 1
                if a.get("best_sku") == b.get("best_sku"):
                    agree += 1
                elif b.get("identity_level") == "accepted":
                    disagree_accepts += 1
    return {
        "reference": reference,
        "candidate": candidate,
        "crops": total,
        "reference_accepted": accepted,
        "agreement_on_accepted": round(agree / accepted, 4) if accepted else None,
        "conflicting_accepts": disagree_accepts,
    }


def main() -> None:
    parser = argparse.ArgumentParser(description="Dorayaki engine component benchmark")
    parser.add_argument("--image", default="image.jpg")
    parser.add_argument("--devices", nargs="*", default=["default", "cpu", "mps"])
    parser.add_argument("--runs", type=int, default=15)
    parser.add_argument("--imgsz", type=int, default=640)
    parser.add_argument("--methods", nargs="*", default=["surf", "orb", "brisk", "hybrid"])
    parser.add_argument("--pools", nargs="*", type=int, default=[1, 4, 8])
    parser.add_argument("--cv2-threads", nargs="*", type=int, default=[0, 1])
    parser.add_argument("--compare", nargs="*", default=None)
    parser.add_argument("--compare-images", nargs="*", default=None)
    parser.add_argument("--json", default=None)
    args = parser.parse_args()

    print("=" * 78)
    print("DORAYAKI ENGINE BENCHMARK")
    print("=" * 78)

    from detector import ShelfProductDetector

    started = time.perf_counter()
    detector = ShelfProductDetector("runs/detect/train/weights/best.pt")
    load_ms = (time.perf_counter() - started) * 1000.0
    print(f"model load: {load_ms:.0f} ms")

    warm_image = load_image(args.image, width=1280)
    started = time.perf_counter()
    detector.detect(warm_image)
    print(f"first predict (cold): {(time.perf_counter() - started) * 1000.0:.0f} ms")
    started = time.perf_counter()
    detector.detect(warm_image)
    print(f"second predict (warm): {(time.perf_counter() - started) * 1000.0:.0f} ms")

    print("\n--- detection latency ---")
    detection = bench_detection(args.image, args.devices, args.runs, args.imgsz)
    rows = []
    for row in detection:
        if "error" in row:
            rows.append([row["device"], "-", "-", "-", f"ERROR: {row['error'][:40]}"])
            continue
        rows.append([row["device"], row["size"], row["shape"], row["boxes"],
                     f"{row['p50_ms']} / {row['p95_ms']}"])
    print(format_table(rows, ["device", "size", "shape", "boxes", "p50/p95 ms"]))

    # device agreement: do mps boxes match cpu boxes?
    by_device = {}
    for device in args.devices:
        if device in ("default", "cpu"):
            continue
        try:
            detector_d = ShelfProductDetector("runs/detect/train/weights/best.pt",
                                              device=None if device == "default" else device)
        except Exception:
            continue
        by_device[device] = detector_d.detect(warm_image).get("boxes", [])
    cpu_boxes = None
    try:
        cpu_detector = ShelfProductDetector("runs/detect/train/weights/best.pt", device="cpu")
        cpu_boxes = cpu_detector.detect(warm_image).get("boxes", [])
    except Exception:
        pass
    if cpu_boxes is not None:
        for device, boxes in by_device.items():
            print(f"agreement cpu vs {device}: {json.dumps(detection_agreement(cpu_boxes, boxes))}")

    print("\n--- classification latency (real crops from", args.image, ") ---")
    crop_report = bench_crops(args.image, args.methods)
    print(f"crops: {crop_report['crop_count']}")
    rows = []
    for method, stats in crop_report["methods"].items():
        levels = stats["levels"]
        rows.append([method, f"{stats['p50_ms']} / {stats['p95_ms']}",
                     f"acc={levels.get('accepted', 0)}",
                     f"amb={levels.get('ambiguous', 0)}",
                     f"unk={levels.get('unknown', 0)}",
                     f"rej={levels.get('rejected', 0)}",
                     f"skus={stats['accepted_skus']}"])
    print(format_table(rows, ["method", "p50/p95 ms", "accepted", "ambiguous", "unknown", "rejected", "distinct skus"]))

    print("\n--- parallel throughput (hybrid) ---")
    from feature_matcher import ProductFeatureMatcher

    matcher = ProductFeatureMatcher(stock_dir="stock")
    image = load_image(args.image)
    boxes = detector.detect(image).get("boxes", [])
    crops = [image[b["box"][1]:b["box"][3], b["box"][0]:b["box"][2]].copy() for b in boxes]
    rows = []
    for row in bench_parallel(crops, matcher, "hybrid", args.pools, args.cv2_threads):
        rows.append([row["cv2_threads"], row["pool"], row["crops"], row["seconds"], row["crops_per_s"]])
    print(format_table(rows, ["cv2_threads", "pool", "crops", "seconds", "crops/s"]))

    if args.compare:
        print("\n--- method agreement ---")
        comparison = compare_methods(args.image, args.compare, args.compare_images)
        print(json.dumps(comparison, indent=2))

    if args.json:
        with open(args.json, "w") as handle:
            json.dump(
                {
                    "model_load_ms": round(load_ms, 1),
                    "detection": detection,
                    "classification": crop_report,
                },
                handle,
                indent=2,
            )
        print(f"\nwrote {args.json}")


if __name__ == "__main__":
    main()