"""Master Pipeline — runs everything automatically when GPU frees up.
===================================================================
Step 1: Wait for GPU > 18GB free
Step 2: Analyze cameras (CPU)
Step 3: Cross-validation (GPU)
Step 4: Adaptive WBF evaluation (GPU, sequential)
Step 5: Kitchen Sink Ultimate v3
Step 6: Save all results
"""
import sys, os, time, json, gc, subprocess
import numpy as np
from PIL import Image, ImageEnhance
from tqdm import tqdm
from collections import defaultdict
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 gpu_free_gb():
    r = subprocess.run(['nvidia-smi', '--query-gpu=memory.free', '--format=csv,noheader'],
                       capture_output=True, text=True)
    return int(r.stdout.strip().split()[0]) / 1024


def log(msg):
    ts = datetime.now().strftime('%H:%M:%S')
    line = f'[{ts}] {msg}'
    print(line)
    with open('logs/pipeline_master.log', 'a') as f: f.write(line + '\n')


# ── Core WBF ──
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):
        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)
    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 predict_augmented(model, img, imgsz, flip=False, brighten=1.0):
    img_aug = img
    if brighten != 1.0: img_aug = ImageEnhance.Brightness(img_aug).enhance(brighten)
    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()
    if flip:
        w = img.size[0]; boxes[:, [0,2]] = w - boxes[:, [2,0]]
    return boxes, scores


def evaluate(name, preds_per_img, 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
    cam_tp = defaultdict(lambda: {t:0 for t in iou_thrs}); cam_gt = defaultdict(int)
    for idx, img_file in enumerate(val_files):
        cam = img_file.split('_2025')[0]
        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); cam_gt[cam] += len(gt_boxes)
        if not gt_boxes: continue
        preds = preds_per_img[idx]
        for t in iou_thrs:
            matched = set()
            for pb in preds:
                if len(pb)==0: continue
                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; cam_tp[cam][t] += 1; matched.add(best_gi)
    recalls = [tp[t]/total_gt for t in iou_thrs]
    mAP = np.mean(recalls)
    per_cam = {}
    for cam in sorted(cam_gt):
        if cam_gt[cam] > 0:
            per_cam[cam] = float(np.mean([cam_tp[cam].get(t,0)/cam_gt[cam] for t in iou_thrs]))
    return mAP, recalls[5], per_cam


def main():
    os.makedirs('logs', exist_ok=True)
    log('PIPELINE MASTER START')

    # Step 0: Wait for GPU
    log('Waiting for GPU > 18GB free...')
    while True:
        free = gpu_free_gb()
        if free > 18: break
        if int(time.time()) % 300 < 2:
            log(f'  GPU free: {free:.0f}GB — still waiting')
        time.sleep(60)
    log(f'GPU free: {gpu_free_gb():.0f}GB — STARTING!')

    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 models
    exp_dir = 'runs/detect/Detection_experiments'
    priority = ['v6_1_s_refined','v12_seed_42','v12_seed_123','v14_seed_789','v14_seed_999','v15_seed_333','v15_yolo11n']
    model_paths = []
    for name in priority:
        path = os.path.join(exp_dir, name, 'weights', 'best.pt')
        if os.path.exists(path): model_paths.append((name, path))
    log(f'Models: {len(model_paths)} — {[n for n,_ in model_paths]}')

    scales = [1280, 1536, 1920]
    flips = [False, True]
    brights = [1.0, 1.2]

    # ── Baseline ──
    log('Running baseline...')
    m0 = YOLO(model_paths[0][1])
    base_preds = []
    for img_file in tqdm(val_files, desc='Baseline'):
        img = Image.open(os.path.join(val_img_dir, img_file))
        b, _ = predict_augmented(m0, img, 1536); base_preds.append(b)
    mAP_base, r75_base, cams_base = evaluate('Baseline', base_preds, val_files, val_img_dir, val_label_dir)
    log(f'Baseline: mAP50-95={mAP_base:.4f} IoU@75={r75_base:.4f}')
    del m0; gc.collect(); torch.cuda.empty_cache()

    # ── Sequential model prediction ──
    log('Running multi-model predictions...')
    model_preds = {}
    for name, path in model_paths:
        log(f'  Model: {name}')
        m = YOLO(path)
        img_preds = []
        for img_file in tqdm(val_files, desc=name, leave=False):
            img = Image.open(os.path.join(val_img_dir, img_file))
            bl, sl = [], []
            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: bl.append(b); sl.append(s)
            img_preds.append((bl, sl))
        model_preds[name] = img_preds
        del m; gc.collect(); torch.cuda.empty_cache()

    # ── Kitchen Sink (standard WBF) ──
    log('Kitchen Sink standard...')
    ks_preds = []
    for idx in range(len(val_files)):
        all_b, all_s = [], []
        for name, _ in model_paths:
            bl, sl = model_preds[name][idx]
            for b, s in zip(bl, sl):
                if len(b) > 0: all_b.append(b); all_s.append(s)
        ks_preds.append(wbf(all_b, all_s)[0])
    mAP_ks, r75_ks, cams_ks = evaluate('Kitchen Sink', ks_preds, val_files, val_img_dir, val_label_dir)
    log(f'Kitchen Sink: mAP50-95={mAP_ks:.4f} IoU@75={r75_ks:.4f} (+{mAP_ks-mAP_base:+.4f})')

    # ── Camera-adaptive WBF ──
    log('Camera-adaptive WBF...')
    # Load camera stats for adaptive thresholds
    if os.path.exists('logs/camera_stats.json'):
        with open('logs/camera_stats.json') as f: cam_stats = json.load(f)
    else:
        cam_stats = {}

    cam_thresholds = {}
    for cam in ['EastLeft','EastRight','WestLeft','WestRight']:
        if cam in cam_stats:
            diff = cam_stats[cam].get('difficulty_score', 0.5)
            cam_thresholds[cam] = 0.50 if diff > 0.6 else 0.55
        else:
            cam_thresholds[cam] = 0.55

    adapt_preds = []
    for idx in range(len(val_files)):
        cam = val_files[idx].split('_2025')[0]
        iou_thr = cam_thresholds.get(cam, 0.55)
        all_b, all_s = [], []
        for name, _ in model_paths:
            bl, sl = model_preds[name][idx]
            for b, s in zip(bl, sl):
                if len(b) > 0: all_b.append(b); all_s.append(s)
        adapt_preds.append(wbf(all_b, all_s, iou_thr=iou_thr)[0])
    mAP_adapt, r75_adapt, cams_adapt = evaluate('Adaptive WBF', adapt_preds, val_files, val_img_dir, val_label_dir)
    log(f'Adaptive WBF: mAP50-95={mAP_adapt:.4f} IoU@75={r75_adapt:.4f} (+{mAP_adapt-mAP_base:+.4f})')

    # ── Final Summary ──
    all_results = {
        'timestamp': datetime.now().isoformat(),
        'baseline': {'mAP50-95': round(mAP_base,4), 'IoU@75': round(r75_base,4), 'per_camera': cams_base},
        'kitchen_sink': {'mAP50-95': round(mAP_ks,4), 'IoU@75': round(r75_ks,4), 'delta': round(mAP_ks-mAP_base,4), 'per_camera': cams_ks},
        'adaptive_wbf': {'mAP50-95': round(mAP_adapt,4), 'IoU@75': round(r75_adapt,4), 'delta': round(mAP_adapt-mAP_base,4), 'per_camera': cams_adapt},
        'n_models': len(model_paths),
        'model_names': [n for n,_ in model_paths],
        'n_sources_per_model': len(scales)*len(flips)*len(brights),
    }
    with open('logs/pipeline_master_results.json', 'w') as f:
        json.dump(all_results, f, indent=2)

    log('='*60)
    log('PIPELINE MASTER COMPLETE')
    log(f'  Baseline:      {mAP_base:.4f}')
    log(f'  Kitchen Sink:  {mAP_ks:.4f}  (+{mAP_ks-mAP_base:+.4f})')
    log(f'  Adaptive WBF:  {mAP_adapt:.4f}  (+{mAP_adapt-mAP_base:+.4f})')
    log(f'  Best:          {max(mAP_ks, mAP_adapt):.4f}')
    log(f'  Results saved to logs/pipeline_master_results.json')
    log('='*60)


if __name__ == '__main__':
    main()