"""WBF Ensemble Evaluation
==========================
Weighted Box Fusion across multiple models for higher mAP50-95.
Combines predictions from different seed models using IoU-based clustering.

Models: v6_1 + v12_seed_42 + v12_seed_123
Weights: proportional to each model's mAP50-95
"""
import sys, os
import numpy as np
from pathlib import Path
from PIL import Image
from tqdm import tqdm

PROJECT_DIR = '/home/user/goat'
os.chdir(PROJECT_DIR)
sys.path.insert(0, PROJECT_DIR)

from ultralytics import YOLO


def compute_iou(b1, b2):
    x1 = max(b1[0], b2[0]); y1 = max(b1[1], b2[1])
    x2 = min(b1[2], b2[2]); y2 = min(b1[3], b2[3])
    inter = max(0, x2 - x1) * max(0, y2 - y1)
    a1 = (b1[2] - b1[0]) * (b1[3] - b1[1])
    a2 = (b2[2] - b2[0]) * (b2[3] - b2[1])
    return inter / (a1 + a2 - inter + 1e-8)


def wbf(boxes_list, scores_list, weights, iou_thr=0.55):
    """Weighted Box Fusion."""
    if not boxes_list:
        return np.array([]), np.array([])

    all_boxes = []
    all_scores = []

    for m_idx, (boxes, scores) in enumerate(zip(boxes_list, scores_list)):
        w = weights[m_idx]
        for i in range(len(boxes)):
            all_boxes.append(boxes[i])
            all_scores.append(scores[i] * w)

    if not all_boxes:
        return np.array([]), np.array([])

    all_boxes = np.array(all_boxes)
    all_scores = np.array(all_scores)
    order = np.argsort(-all_scores)
    all_boxes = all_boxes[order]
    all_scores = all_scores[order]

    clusters = []
    used = np.zeros(len(all_boxes), dtype=bool)

    for i in range(len(all_boxes)):
        if used[i]:
            continue
        cluster = [(all_boxes[i], all_scores[i])]
        used[i] = True

        for j in range(i + 1, len(all_boxes)):
            if used[j]:
                continue
            total_w = sum(s for _, s in cluster)
            center = np.zeros(4)
            for b, s in cluster:
                center += b * s / total_w
            if compute_iou(center.tolist(), all_boxes[j].tolist()) > iou_thr:
                cluster.append((all_boxes[j], all_scores[j]))
                used[j] = True
        clusters.append(cluster)

    result_boxes = []
    result_scores = []
    for cluster in clusters:
        total_w = sum(s for _, s in cluster)
        avg_box = np.zeros(4)
        for b, s in cluster:
            avg_box += b * s / total_w
        avg_score = total_w / len(weights)
        result_boxes.append(avg_box)
        result_scores.append(avg_score)

    return np.array(result_boxes), np.array(result_scores)


def main():
    # Models
    model_paths = [
        'runs/detect/Detection_experiments/v6_1_s_refined/weights/best.pt',
        'runs/detect/Detection_experiments/v12_seed_42/weights/best.pt',
        'runs/detect/Detection_experiments/v12_seed_123/weights/best.pt',
    ]
    weights = [0.5125, 0.5097, 0.5113]  # mAP50-95 of each model
    weights = [w / sum(weights) for w in weights]

    models = []
    for i, mp in enumerate(model_paths):
        if os.path.exists(mp):
            models.append((YOLO(mp), weights[i]))
            print(f'Loaded: {mp} (weight={weights[i]:.3f})')
        else:
            print(f'MISSING: {mp}')

    if len(models) < 2:
        print('Need at least 2 models')
        return

    val_img_dir = 'Data/Detection_dataset/images/val'
    val_label_dir = 'Data/Detection_dataset/labels/val'
    val_files = sorted([f for f in os.listdir(val_img_dir) if f.endswith('.jpg')])

    # Evaluate
    iou_thresholds = [round(0.5 + i * 0.05, 2) for i in range(10)]
    per_iou_tp = {t: 0 for t in iou_thresholds}
    total_gt = 0

    for img_file in tqdm(val_files, desc='Evaluating'):
        img_path = os.path.join(val_img_dir, img_file)
        img = Image.open(img_path)
        iw, ih = img.size

        # GT
        lf = img_file.replace('.jpg', '.txt')
        lpath = os.path.join(val_label_dir, lf)
        gt_boxes = []
        if os.path.exists(lpath):
            with open(lpath) as f:
                for line in f:
                    p = line.strip().split()
                    if len(p) >= 5:
                        cx, cy, w, h = [float(x) for x in p[1:5]]
                        gt_boxes.append([
                            (cx - w/2) * iw, (cy - h/2) * ih,
                            (cx + w/2) * iw, (cy + h/2) * ih
                        ])
        total_gt += len(gt_boxes)
        if not gt_boxes:
            continue

        # Get predictions from all models
        boxes_list = []
        scores_list = []
        for model, _ in models:
            results = model.predict(
                source=img_path, imgsz=1536, conf=0.25, iou=0.7,
                max_det=100, save=False, verbose=False
            )
            if results and len(results[0].boxes):
                boxes_list.append(results[0].boxes.xyxy.cpu().numpy())
                scores_list.append(results[0].boxes.conf.cpu().numpy())
            else:
                boxes_list.append(np.array([]))
                scores_list.append(np.array([]))

        # WBF
        fused_boxes, fused_scores = wbf(boxes_list, scores_list,
                                         [w for _, w in models])

        # Evaluate
        for t in iou_thresholds:
            matched = set()
            for pb in fused_boxes:
                best_iou, best_gi = 0, -1
                for gi, gb in enumerate(gt_boxes):
                    if gi in matched:
                        continue
                    iou = compute_iou(pb.tolist(), gb)
                    if iou > best_iou:
                        best_iou = iou
                        best_gi = gi
                if best_iou >= t and best_gi >= 0:
                    per_iou_tp[t] += 1
                    matched.add(best_gi)

    # Results
    print(f'\n{"="*60}')
    print(f'WBF Ensemble Results (3 models)')
    print(f'{"="*60}')
    recalls = []
    for t in iou_thresholds:
        r = per_iou_tp[t] / total_gt if total_gt else 0
        recalls.append(r)
        print(f'  IoU@{t:.2f}: Recall={r:.4f}')
    map5095 = np.mean(recalls)
    print(f'\n  Approx mAP50-95: {map5095:.4f}')
    print(f'  vs v6_1 single (0.5125): {map5095-0.5125:+.4f}')


if __name__ == '__main__':
    main()