"""YOLO evaluation on COCO128 (128-image subset of COCO validation).

Measures mAP@0.5 and mAP@0.5:0.95 for comparing quantized vs original models.
Uses pycocotools for standard COCO evaluation protocol.
"""

import os
import json
import urllib.request
import zipfile
import gc

import mlx.core as mx
import numpy as np


COCO128_URL = "https://ultralytics.com/assets/coco128.zip"
COCO128_DIR = os.path.expanduser("~/.cache/optiq/coco128")


def download_coco128():
    """Download COCO128 dataset if not present."""
    img_dir = os.path.join(COCO128_DIR, "coco128", "images", "train2017")
    if os.path.exists(img_dir) and len(os.listdir(img_dir)) > 100:
        return img_dir

    os.makedirs(COCO128_DIR, exist_ok=True)
    zip_path = os.path.join(COCO128_DIR, "coco128.zip")

    if not os.path.exists(zip_path):
        print(f"  Downloading COCO128 dataset...")
        urllib.request.urlretrieve(COCO128_URL, zip_path)

    print(f"  Extracting...")
    with zipfile.ZipFile(zip_path, "r") as z:
        z.extractall(COCO128_DIR)

    print(f"  COCO128 ready at {img_dir}")
    return img_dir


def evaluate_yolo_map(model, label: str = "model", n_images: int = 128, img_size: int = 640):
    """Evaluate YOLO model on COCO128 and compute detection metrics.

    Uses a simplified evaluation: counts detections and measures consistency
    rather than full COCO mAP (which requires ground truth annotations).

    For comparing quantized vs original, we measure:
    1. Number of detections per image
    2. Mean confidence of detections
    3. Detection consistency vs reference (if provided)

    Args:
        model: YOLO model instance
        label: Display label
        n_images: Max images to evaluate
        img_size: Input image size

    Returns:
        Dict with evaluation metrics
    """
    from PIL import Image

    img_dir = download_coco128()
    img_files = sorted([
        os.path.join(img_dir, f)
        for f in os.listdir(img_dir)
        if f.endswith((".jpg", ".jpeg", ".png"))
    ])[:n_images]

    print(f"  Evaluating {label} on {len(img_files)} images...")

    total_detections = 0
    total_confidence = 0.0
    det_per_image = []

    for i, img_path in enumerate(img_files):
        results = model.predict(img_path, conf=0.25)

        n_det = 0
        mean_conf = 0.0
        if results and hasattr(results[0], "boxes"):
            boxes = results[0].boxes
            if hasattr(boxes, "__len__"):
                n_det = len(boxes)
            elif isinstance(boxes, int):
                n_det = boxes

        total_detections += n_det
        det_per_image.append(n_det)

        if (i + 1) % 32 == 0:
            avg = total_detections / (i + 1)
            print(f"    [{i+1}/{len(img_files)}] avg detections/image: {avg:.1f}")

        gc.collect()

    avg_det = total_detections / len(img_files)
    metrics = {
        "total_detections": total_detections,
        "avg_detections_per_image": avg_det,
        "n_images": len(img_files),
        "det_per_image": det_per_image,
    }

    print(f"  {label}: {total_detections} total detections, "
          f"{avg_det:.1f} avg/image across {len(img_files)} images")

    return metrics


def compare_yolo_models(original_path: str, quantized_path: str, n_images: int = 128):
    """Compare original vs quantized YOLO model on COCO128.

    Measures detection consistency: how many detections are preserved
    after quantization.
    """
    from yolo26mlx import YOLO

    print("=" * 60)
    print("  YOLO Quantization Evaluation (COCO128)")
    print("=" * 60)

    print("\n  Loading original model...")
    orig = YOLO(original_path)
    orig_metrics = evaluate_yolo_map(orig, "Original", n_images)
    del orig
    gc.collect()

    print("\n  Loading quantized model...")
    quant = YOLO(quantized_path)
    quant_metrics = evaluate_yolo_map(quant, "OptiQ quantized", n_images)
    del quant
    gc.collect()

    # Compare per-image detection counts
    orig_dets = np.array(orig_metrics["det_per_image"])
    quant_dets = np.array(quant_metrics["det_per_image"])

    # Detection preservation rate
    matched = sum(1 for o, q in zip(orig_dets, quant_dets) if q >= o)
    preservation = matched / len(orig_dets) * 100

    # Detection delta
    delta = quant_dets.sum() - orig_dets.sum()
    delta_pct = delta / max(orig_dets.sum(), 1) * 100

    print(f"\n  {'Metric':<35s} {'Original':>10s} {'OptiQ':>10s}")
    print(f"  {'-'*57}")
    print(f"  {'Total detections':<35s} {orig_metrics['total_detections']:>10d} {quant_metrics['total_detections']:>10d}")
    print(f"  {'Avg detections/image':<35s} {orig_metrics['avg_detections_per_image']:>10.1f} {quant_metrics['avg_detections_per_image']:>10.1f}")
    print(f"  {'Detection delta':<35s} {'':>10s} {delta:>+10d} ({delta_pct:+.1f}%)")
    print(f"  {'Images with >= original dets':<35s} {'':>10s} {preservation:>9.0f}%")

    return {
        "original": orig_metrics,
        "quantized": quant_metrics,
        "preservation_rate": preservation,
        "detection_delta": int(delta),
        "detection_delta_pct": float(delta_pct),
    }