"""3-Fold Cross Validation — measure true mAP with confidence intervals.
=========================================================================
The 166-image val set is small. Real mAP may differ from single-split estimate.
This gives mean ± std across 3 folds for reliable comparison.
"""
import sys, os, json, gc
import numpy as np
from PIL import Image
from tqdm import tqdm
from datetime import datetime

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

import torch
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 eval_fold(model_path, val_files, val_img_dir, val_label_dir):
    """Evaluate a single model on a set of val files."""
    model = YOLO(model_path)
    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='Fold', leave=False):
        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

        r = model.predict(img, imgsz=1536, conf=0.25, iou=0.7, max_det=100, verbose=False)
        preds = r[0].boxes.xyxy.cpu().numpy() if r and len(r[0].boxes) else np.array([])

        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]
    del model; gc.collect(); torch.cuda.empty_cache()
    return np.mean(recalls), recalls[5]


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

    # Stratified split: preserve camera proportions
    from collections import defaultdict
    cam_files = defaultdict(list)
    for f in all_val:
        cam = f.split('_2025')[0]
        cam_files[cam].append(f)

    # Create 3 folds
    np.random.seed(42)
    folds = [[] for _ in range(3)]
    for cam, files in cam_files.items():
        files = sorted(files)
        np.random.shuffle(files)
        n = len(files)
        folds[0].extend(files[:n//3])
        folds[1].extend(files[n//3:2*n//3])
        folds[2].extend(files[2*n//3:])

    print('Cross-validation folds (stratified by camera):')
    for i, fold in enumerate(folds):
        fc = defaultdict(int)
        for f in fold: fc[f.split('_2025')[0]] += 1
        print(f'  Fold {i+1}: {len(fold)} images, {dict(fc)}')

    # Evaluate v6_1 on each fold
    model_path = 'runs/detect/Detection_experiments/v6_1_s_refined/weights/best.pt'
    results = []
    for i, fold in enumerate(folds):
        mAP, r75 = eval_fold(model_path, fold, val_img_dir, val_label_dir)
        results.append(mAP)
        print(f'  Fold {i+1}: mAP50-95={mAP:.4f}, IoU@75={r75:.4f}')

    mean_mAP = np.mean(results)
    std_mAP = np.std(results, ddof=1)

    print(f'\n{"="*60}')
    print(f'CROSS-VALIDATION RESULTS (v6_1)')
    print(f'{"="*60}')
    print(f'  Fold 1: {results[0]:.4f}')
    print(f'  Fold 2: {results[1]:.4f}')
    print(f'  Fold 3: {results[2]:.4f}')
    print(f'  Mean ± Std: {mean_mAP:.4f} ± {std_mAP:.4f}')
    print(f'  Original single-split: 0.5125')
    print(f'  Bias: {0.5125 - mean_mAP:+.4f}')

    if std_mAP > 0.005:
        print(f'\n  ⚠ Std > 0.005 — single-split estimate may be unreliable!')

    with open('logs/crossval_results.json', 'w') as f:
        json.dump({
            'folds': [round(r, 4) for r in results],
            'mean': round(mean_mAP, 4),
            'std': round(std_mAP, 4),
            'original': 0.5125,
        }, f, indent=2)


if __name__ == '__main__':
    main()