"""Kitchen Sink v2 — Calibrated + Adaptive WBF
=================================================
Improvements over v1:
1. Confidence calibration per model (min-max normalize to [0,1])
2. Adaptive IoU threshold per image (based on prediction density)
3. More augmentations: +rotation (±5°)
4. All models found in experiments directory
"""
import sys, os, json
import numpy as np
from PIL import Image, ImageEnhance
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 calibrate_confidence(boxes_list, scores_list):
    """Per-model min-max confidence calibration."""
    calibrated_scores = []
    for scores in scores_list:
        if len(scores) == 0:
            calibrated_scores.append(scores)
            continue
        smin, smax = scores.min(), scores.max()
        if smax - smin < 1e-8:
            calibrated_scores.append(scores)
        else:
            calibrated_scores.append((scores - smin) / (smax - smin))
    return calibrated_scores


def adaptive_wbf(boxes_list, scores_list):
    """WBF with adaptive IoU threshold based on prediction density."""
    if not boxes_list or all(len(b) == 0 for b in boxes_list):
        return np.array([]), np.array([])

    # Calibrate
    scores_list = calibrate_confidence(boxes_list, scores_list)

    # Calculate prediction density
    all_boxes = []
    all_scores = []
    for boxes, scores in zip(boxes_list, scores_list):
        all_boxes.extend(boxes)
        all_scores.extend(scores)
    if not all_boxes: return np.array([]), np.array([])

    all_boxes = np.array(all_boxes)
    all_scores = np.array(all_scores)

    # Adaptive IoU: lower threshold when many predictions (dense scene), higher when sparse
    n_preds = len(all_boxes)
    iou_thr = 0.50 if n_preds > 200 else (0.55 if n_preds > 100 else 0.60)

    # Sort by score
    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)
        # Score = avg confidence × sqrt(cluster_size) (reward consensus)
        consensus_score = tw * np.sqrt(len(cl))
        result_boxes.append(avg_b)
        result_scores.append(consensus_score)

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


def predict_augmented(model, img, imgsz, flip=False, brighten=1.0, rotate=0):
    """Predict with augmentations."""
    img_aug = img
    if brighten != 1.0:
        img_aug = ImageEnhance.Brightness(img_aug).enhance(brighten)
    if rotate != 0:
        img_aug = img_aug.rotate(rotate, expand=False)
    if flip:
        img_aug = img_aug.transpose(Image.FLIP_LEFT_RIGHT)

    r = model.predict(img_aug, imgsz=imgsz, conf=0.25, iou=0.7, max_det=100, verbose=False)
    if not r or len(r[0].boxes) == 0:
        return np.array([]), np.array([])

    boxes = r[0].boxes.xyxy.cpu().numpy()
    scores = r[0].boxes.conf.cpu().numpy()

    # Reverse augmentations
    if rotate != 0:
        cx, cy = img.size[0]/2, img.size[1]/2
        rad = np.radians(-rotate)
        cos, sin = np.cos(rad), np.sin(rad)
        for i in range(len(boxes)):
            for corner in [(0, 1), (2, 3)]:
                x = boxes[i, corner[0]] - cx
                y = boxes[i, corner[1]] - cy
                boxes[i, corner[0]] = x*cos - y*sin + cx
                boxes[i, corner[1]] = x*sin + y*cos + cy

    if flip:
        w = img.size[0]
        boxes[:, [0, 2]] = w - boxes[:, [2, 0]]

    return boxes, scores


def evaluate(name, get_boxes_fn, val_files, val_img_dir, val_label_dir):
    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]
    return np.mean(recalls), recalls[5]


def main():
    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')])

    # Discover ALL models (including yolo11n)
    exp_dir = 'runs/detect/Detection_experiments'
    models = []
    priority = ['v6_1_s_refined', 'v12_seed_42', 'v12_seed_123', 'v14_seed_789', 'v14_seed_999', 'v15_seed_333', 'v15_yolo11n']
    for name in priority:
        path = os.path.join(exp_dir, name, 'weights', 'best.pt')
        if os.path.exists(path):
            models.append((YOLO(path), name))
            print(f'Loaded: {name} ({path.split(chr(47))[-3]})')

    print(f'\n{len(models)} models ready')

    scales = [1280, 1536, 1920]
    flips = [False, True]
    brights = [1.0, 1.2]
    rotates = [0, 5]  # NEW: rotation augmentation

    # Baseline
    def baseline(img):
        m, _ = models[0]
        boxes, _ = predict_augmented(m, img, 1536)
        return boxes

    mAP_base, r75_base = evaluate('Single baseline', baseline, val_files, val_img_dir, val_label_dir)

    # Standard Kitchen Sink (same as v1, for comparison)
    def ks_standard(img):
        all_boxes, all_scores = [], []
        for m, _ in models:
            for sz in scales:
                for fl in flips:
                    for br in brights:
                        b, s = predict_augmented(m, img, sz, fl, br)
                        if len(b) > 0:
                            all_boxes.append(b); all_scores.append(s)
        return adaptive_wbf(all_boxes, all_scores)[0]

    n_src = len(models) * len(scales) * len(flips) * len(brights)
    mAP_ks, r75_ks = evaluate(f'Kitchen Sink ({n_src}x)', ks_standard, val_files, val_img_dir, val_label_dir)

    # Kitchen Sink + Rotation
    def ks_rotated(img):
        all_boxes, all_scores = [], []
        for m, _ in models:
            for sz in scales:
                for fl in flips:
                    for br in brights:
                        for rot in rotates:
                            b, s = predict_augmented(m, img, sz, fl, br, rot)
                            if len(b) > 0:
                                all_boxes.append(b); all_scores.append(s)
        return adaptive_wbf(all_boxes, all_scores)[0]

    n_src2 = n_src * len(rotates)
    mAP_ks2, r75_ks2 = evaluate(f'K.Sink +Rot ({n_src2}x)', ks_rotated, val_files, val_img_dir, val_label_dir)

    print(f'\n{"="*60}')
    print(f'KITCHEN SINK v2 RESULTS')
    print(f'{"="*60}')
    print(f'Baseline (1m):           {mAP_base:.4f}  IoU@75={r75_base:.4f}')
    print(f'Kitchen Sink ({n_src}x):     {mAP_ks:.4f}  IoU@75={r75_ks:.4f}  (+{mAP_ks-mAP_base:+.4f})')
    print(f'K.Sink +Rot ({n_src2}x):   {mAP_ks2:.4f}  IoU@75={r75_ks2:.4f}  (+{mAP_ks2-mAP_base:+.4f})')

    with open('logs/kitchen_sink_v2_results.json', 'w') as f:
        json.dump({
            'baseline': round(mAP_base, 4),
            'kitchen_sink': round(mAP_ks, 4),
            'kitchen_sink_rotated': round(mAP_ks2, 4),
            'n_models': len(models),
            'model_names': [n for _, n in models],
        }, f, indent=2)


if __name__ == '__main__':
    main()