"""Benchmark different ensemble configs on GPU. Measures combined mAP score."""
import json
import time
import itertools
import numpy as np
from pathlib import Path
from pycocotools.coco import COCO
from pycocotools.cocoeval import COCOeval

# Patch torch.load
import torch
_orig = torch.load
def _safe(*a, **kw): kw["weights_only"] = False; return _orig(*a, **kw)
torch.load = _safe

from ultralytics import YOLO
from ensemble_boxes import weighted_boxes_fusion

ANN = "input/train/annotations.json"
IMG_DIR = Path("input/train/images")

# Load ground truth
coco_gt = COCO(ANN)
img_ids = sorted(coco_gt.getImgIds())
img_infos = coco_gt.loadImgs(img_ids)

def get_predictions_single(model, imgsz=1280, conf=0.001, iou=0.6, max_det=500):
    """Run single model, return predictions list."""
    preds = []
    for img_info in img_infos:
        img_path = IMG_DIR / img_info["file_name"]
        results = model.predict(str(img_path), verbose=False, imgsz=imgsz,
                                conf=conf, iou=iou, max_det=max_det, augment=False)
        for r in results:
            if r.boxes is None: continue
            for box, sc, cl in zip(r.boxes.xyxy.cpu().numpy(),
                                    r.boxes.conf.cpu().numpy(),
                                    r.boxes.cls.cpu().numpy()):
                x1, y1, x2, y2 = box
                preds.append({
                    "image_id": img_info["id"], "category_id": int(cl),
                    "bbox": [float(x1), float(y1), float(x2-x1), float(y2-y1)],
                    "score": float(sc)
                })
    return preds

def get_predictions_ensemble(models, imgsz=1280, conf=0.001, iou=0.7,
                              max_det=500, wbf_iou=0.55, skip_thr=0.001):
    """Run ensemble with WBF, return predictions list."""
    preds = []
    for img_info in img_infos:
        img_path = IMG_DIR / img_info["file_name"]
        all_boxes, all_scores, all_labels = [], [], []
        img_w, img_h = None, None

        for model in models:
            results = model.predict(str(img_path), verbose=False, imgsz=imgsz,
                                    conf=conf, iou=iou, max_det=max_det, augment=False)
            boxes, scores, labels = [], [], []
            for r in results:
                if r.boxes is None: continue
                if img_w is None: img_h, img_w = r.orig_shape
                for box, sc, cl in zip(r.boxes.xyxy.cpu().numpy(),
                                        r.boxes.conf.cpu().numpy(),
                                        r.boxes.cls.cpu().numpy()):
                    x1, y1, x2, y2 = box
                    boxes.append([x1/img_w, y1/img_h, x2/img_w, y2/img_h])
                    scores.append(float(sc))
                    labels.append(int(cl))
            all_boxes.append(boxes if boxes else [[0,0,0,0]])
            all_scores.append(scores if scores else [0])
            all_labels.append(labels if labels else [0])

        if img_w is None: continue
        fb, fs, fl = weighted_boxes_fusion(all_boxes, all_scores, all_labels,
                                            iou_thr=wbf_iou, skip_box_thr=skip_thr)
        for box, sc, lb in zip(fb, fs, fl):
            x1, y1 = box[0]*img_w, box[1]*img_h
            x2, y2 = box[2]*img_w, box[3]*img_h
            w, h = x2-x1, y2-y1
            if w > 0 and h > 0:
                preds.append({
                    "image_id": img_info["id"], "category_id": int(lb),
                    "bbox": [float(x1), float(y1), float(w), float(h)],
                    "score": float(sc)
                })
    return preds

def compute_score(preds):
    """Compute 0.7*det_mAP + 0.3*cls_mAP."""
    if not preds:
        return 0, 0, 0

    # Cap at 49000
    if len(preds) > 49000:
        preds.sort(key=lambda x: x["score"], reverse=True)
        preds = preds[:49000]

    gt_data = json.load(open(ANN))

    # Detection mAP (category-agnostic)
    det_preds = [dict(p, category_id=1) for p in preds]
    det_gt = dict(gt_data)
    det_gt["annotations"] = [dict(a, category_id=1) for a in gt_data["annotations"]]
    det_gt["categories"] = [{"id": 1, "name": "product"}]
    with open("/tmp/det_gt.json", "w") as f: json.dump(det_gt, f)
    with open("/tmp/det_pr.json", "w") as f: json.dump(det_preds, f)

    coco_d = COCO("/tmp/det_gt.json")
    ev = COCOeval(coco_d, coco_d.loadRes("/tmp/det_pr.json"), "bbox")
    ev.params.iouThrs = [0.5]
    ev.evaluate(); ev.accumulate(); ev.summarize()
    det_map = ev.stats[0]

    # Classification mAP
    with open("/tmp/cls_pr.json", "w") as f: json.dump(preds, f)
    coco_c = COCO(ANN)
    ev2 = COCOeval(coco_c, coco_c.loadRes("/tmp/cls_pr.json"), "bbox")
    ev2.params.iouThrs = [0.5]
    ev2.evaluate(); ev2.accumulate(); ev2.summarize()
    cls_map = ev2.stats[0]

    combined = 0.7 * det_map + 0.3 * cls_map
    return combined, det_map, cls_map

# Load models
print("Loading models...")
m1 = YOLO("bench/model1.onnx", task="detect")  # run4: s123, 1280, SGD (best)
m2 = YOLO("bench/model2.onnx", task="detect")  # run1: s42, 1024, AdamW
m3 = YOLO("bench/model3.onnx", task="detect")  # run2: s42, 1280, AdamW

# Also load the .pt models for more diversity
pt_models = {}
for pt in Path(".").glob("train_*/run/weights/best.pt"):
    name = pt.parts[0]
    pt_models[name] = YOLO(str(pt), task="detect")
    print(f"  Loaded {name}")

print(f"\nLoaded 3 ONNX + {len(pt_models)} PT models")
print("=" * 60)

results = []

# Test 1: Single models
print("\n--- SINGLE MODELS ---")
for name, model in [("model1(s123_1280_SGD)", m1), ("model2(s42_1024_AdamW)", m2), ("model3(s42_1280_AdamW)", m3)]:
    t = time.time()
    preds = get_predictions_single(model)
    score, det, cls = compute_score(preds)
    elapsed = time.time() - t
    print(f"{name}: combined={score:.4f} det={det:.4f} cls={cls:.4f} preds={len(preds)} time={elapsed:.0f}s")
    results.append((name, score, det, cls))

# Test 2: 3-model ensemble with different WBF iou thresholds
print("\n--- ENSEMBLE WBF IOU SWEEP ---")
for wbf_iou in [0.4, 0.45, 0.5, 0.55, 0.6, 0.65, 0.7]:
    t = time.time()
    preds = get_predictions_ensemble([m1, m2, m3], wbf_iou=wbf_iou)
    score, det, cls = compute_score(preds)
    elapsed = time.time() - t
    print(f"wbf_iou={wbf_iou}: combined={score:.4f} det={det:.4f} cls={cls:.4f} preds={len(preds)} time={elapsed:.0f}s")
    results.append((f"ensemble_wbf{wbf_iou}", score, det, cls))

# Test 3: Different max_det
print("\n--- MAX_DET SWEEP ---")
for max_det in [300, 500, 800, 1000]:
    t = time.time()
    preds = get_predictions_ensemble([m1, m2, m3], max_det=max_det, wbf_iou=0.55)
    score, det, cls = compute_score(preds)
    elapsed = time.time() - t
    print(f"max_det={max_det}: combined={score:.4f} det={det:.4f} cls={cls:.4f} preds={len(preds)} time={elapsed:.0f}s")
    results.append((f"maxdet_{max_det}", score, det, cls))

# Test 4: Different conf thresholds
print("\n--- CONF SWEEP ---")
for conf in [0.0001, 0.001, 0.005, 0.01]:
    t = time.time()
    preds = get_predictions_ensemble([m1, m2, m3], conf=conf, wbf_iou=0.55)
    score, det, cls = compute_score(preds)
    elapsed = time.time() - t
    print(f"conf={conf}: combined={score:.4f} det={det:.4f} cls={cls:.4f} preds={len(preds)} time={elapsed:.0f}s")
    results.append((f"conf_{conf}", score, det, cls))

# Summary
print("\n" + "=" * 60)
print("RANKED RESULTS:")
results.sort(key=lambda x: x[1], reverse=True)
for i, (name, score, det, cls) in enumerate(results):
    marker = " <-- BEST" if i == 0 else ""
    print(f"  {i+1}. {name}: {score:.4f} (det={det:.4f} cls={cls:.4f}){marker}")