"""Multi-Scale WBF: Single model at 3 scales → WBF Ensemble
============================================================
Each model predicts at 1280, 1536, 1920. WBF fuses all 3 predictions.
Zero training cost. Tested on v6_1.
"""
import sys, os
import numpy as np
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, y1 = max(b1[0], b2[0]), max(b1[1], b2[1])
    x2, y2 = min(b1[2], b2[2]), 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, iou_thr=0.55):
    if not boxes_list or all(len(b) == 0 for b in boxes_list):
        return np.array([]), np.array([])
    all_boxes, all_scores = [], []
    for boxes, scores in zip(boxes_list, scores_list):
        for i in range(len(boxes)):
            all_boxes.append(boxes[i])
            all_scores.append(scores[i])
    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_scores = all_boxes[order], 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
            tw = sum(s for _, s in cluster)
            center = sum(b*s/tw for b, s in cluster)
            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 cl in clusters:
        tw = sum(s for _, s in cl)
        avg_b = sum(b*s/tw for b, s in cl)
        result_boxes.append(avg_b)
        result_scores.append(tw)
    return np.array(result_boxes), np.array(result_scores)


def eval_config(name, get_boxes_fn):
    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')])
    iou_thrs = [round(0.5+i*0.05, 2) for i in range(10)]
    tp = {t:0 for t in iou_thrs}
    total_gt = 0
    for img_file in tqdm(val_files, desc=name):
        img = Image.open(os.path.join(val_img_dir, img_file))
        gt_boxes = []
        lf = img_file.replace('.jpg','.txt')
        with open(os.path.join(val_label_dir, lf)) 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)*img.size[0], (cy-h/2)*img.size[1], (cx+w/2)*img.size[0], (cy+h/2)*img.size[1]])
        total_gt += len(gt_boxes)
        if not gt_boxes: continue
        preds = get_boxes_fn(img)
        for t in iou_thrs:
            matched = set()
            for pb in preds:
                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:
                    tp[t] += 1; matched.add(best_gi)
    recalls = [tp[t]/total_gt for t in iou_thrs]
    mAP = np.mean(recalls)
    print(f'  {name}: mAP50-95={mAP:.4f}, IoU@75={recalls[5]:.4f}')
    return mAP


def main():
    model = YOLO('runs/detect/Detection_experiments/v6_1_s_refined/weights/best.pt')
    scales = [1280, 1536, 1920]

    # Baseline: single scale 1536
    def single_1536(img):
        r = model.predict(img, imgsz=1536, conf=0.25, iou=0.7, max_det=100, verbose=False)
        if r and len(r[0].boxes): return r[0].boxes.xyxy.cpu().numpy()
        return np.array([])

    mAP_single = eval_config('Single 1536', single_1536)

    # Multi-scale WBF
    def multiscale_wbf(img):
        boxes_list, scores_list = [], []
        for sz in scales:
            r = model.predict(img, imgsz=sz, conf=0.25, iou=0.7, max_det=100, verbose=False)
            if r and len(r[0].boxes):
                boxes_list.append(r[0].boxes.xyxy.cpu().numpy())
                scores_list.append(r[0].boxes.conf.cpu().numpy())
            else:
                boxes_list.append(np.array([]))
                scores_list.append(np.array([]))
        fused, _ = wbf(boxes_list, scores_list)
        return fused

    mAP_ms = eval_config(f'MultiScale WBF ({scales})', multiscale_wbf)

    print(f'\n{"="*50}')
    print(f'Single 1536:        {mAP_single:.4f}')
    print(f'MultiScale WBF:     {mAP_ms:.4f}  (+{mAP_ms-mAP_single:+.4f})')


if __name__ == '__main__':
    main()