"""Edge-Guided Bbox Refinement
================================
For each box edge, scan outward and find the image gradient peak.
More constrained than Snake, less likely to drift to wrong edges.
CPU-only, 0 training.
"""
import sys, os
import numpy as np
import cv2
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 refine_edge(gray, x1, y1, x2, y2, edge, scan_range=10):
    """Refine one edge of a bbox using gradient search.

    edge: 'left', 'right', 'top', 'bottom'
    Returns: adjusted coordinate
    """
    h, w = gray.shape
    best_pos = None
    best_grad = 0

    if edge == 'left':
        x_center = int(x1)
        for dx in range(-scan_range, scan_range + 1):
            new_x = int(x1 + dx)
            if new_x < 1 or new_x > w - 2:
                continue
            # Gradient along vertical strip at new_x
            strip = gray[max(0, int(y1)):min(h, int(y2)), new_x]
            if len(strip) < 3:
                continue
            grad = np.abs(np.diff(strip)).mean()
            if grad > best_grad:
                best_grad = grad
                best_pos = new_x
        return best_pos if best_pos is not None else x1

    elif edge == 'right':
        x_center = int(x2)
        for dx in range(-scan_range, scan_range + 1):
            new_x = int(x2 + dx)
            if new_x < 1 or new_x > w - 2:
                continue
            strip = gray[max(0, int(y1)):min(h, int(y2)), new_x]
            if len(strip) < 3:
                continue
            grad = np.abs(np.diff(strip)).mean()
            if grad > best_grad:
                best_grad = grad
                best_pos = new_x
        return best_pos if best_pos is not None else x2

    elif edge == 'top':
        y_center = int(y1)
        for dy in range(-scan_range, scan_range + 1):
            new_y = int(y1 + dy)
            if new_y < 1 or new_y > h - 2:
                continue
            strip = gray[new_y, max(0, int(x1)):min(w, int(x2))]
            if len(strip) < 3:
                continue
            grad = np.abs(np.diff(strip)).mean()
            if grad > best_grad:
                best_grad = grad
                best_pos = new_y
        return best_pos if best_pos is not None else y1

    elif edge == 'bottom':
        y_center = int(y2)
        for dy in range(-scan_range, scan_range + 1):
            new_y = int(y2 + dy)
            if new_y < 1 or new_y > h - 2:
                continue
            strip = gray[new_y, max(0, int(x1)):min(w, int(x2))]
            if len(strip) < 3:
                continue
            grad = np.abs(np.diff(strip)).mean()
            if grad > best_grad:
                best_grad = grad
                best_pos = new_y
        return best_pos if best_pos is not None else y2

    return x1 if edge in ('left', 'right') else y1


def refine_box(img_gray, box, scan_range=8):
    """Refine all 4 edges of a box using gradient search."""
    x1, y1, x2, y2 = box
    w, h = x2 - x1, y2 - y1
    if w < 10 or h < 10:
        return box

    new_x1 = refine_edge(img_gray, x1, y1, x2, y2, 'left', scan_range)
    new_x2 = refine_edge(img_gray, x1, y1, x2, y2, 'right', scan_range)
    new_y1 = refine_edge(img_gray, x1, y1, x2, y2, 'top', scan_range)
    new_y2 = refine_edge(img_gray, x1, y1, x2, y2, 'bottom', scan_range)

    # Constrain: don't let the box shrink too much
    min_w, min_h = w * 0.5, h * 0.5
    max_w, max_h = w * 1.5, h * 1.5

    if new_x2 - new_x1 < min_w:
        mid = (new_x1 + new_x2) / 2
        new_x1 = mid - min_w / 2
        new_x2 = mid + min_w / 2
    if new_y2 - new_y1 < min_h:
        mid = (new_y1 + new_y2) / 2
        new_y1 = mid - min_h / 2
        new_y2 = mid + min_h / 2
    if new_x2 - new_x1 > max_w:
        mid = (new_x1 + new_x2) / 2
        new_x1 = mid - max_w / 2
        new_x2 = mid + max_w / 2
    if new_y2 - new_y1 > max_h:
        mid = (new_y1 + new_y2) / 2
        new_y1 = mid - max_h / 2
        new_y2 = mid + max_h / 2

    return np.array([new_x1, new_y1, new_x2, new_y2])


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 main():
    model = YOLO('runs/detect/Detection_experiments/v6_1_s_refined/weights/best.pt')
    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')])

    # Quick test on 5 images
    test_files = val_files[:5]
    improved, degraded, total = 0, 0, 0

    print('Edge-Guided Refinement test:')
    for img_file in test_files:
        img_path = os.path.join(val_img_dir, img_file)
        img = Image.open(img_path)
        gray = np.array(img.convert('L'), dtype=float)

        r = model.predict(img, imgsz=1536, conf=0.25, verbose=False)
        if not r or len(r[0].boxes) == 0:
            continue

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

        # GT
        lf = img_file.replace('.jpg', '.txt')
        gt_boxes = []
        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]])

        for box in boxes[:10]:
            refined = refine_box(gray, box)
            total += 1

            best_orig = max(compute_iou(box, gt) for gt in gt_boxes) if gt_boxes else 0
            best_ref = max(compute_iou(refined, gt) for gt in gt_boxes) if gt_boxes else 0

            if best_ref > best_orig + 0.005:
                improved += 1
            elif best_ref < best_orig - 0.005:
                degraded += 1

    print(f'  Improved: {improved}/{total} ({improved/total*100:.0f}%)')
    print(f'  Degraded: {degraded}/{total} ({degraded/total*100:.0f}%)')
    print(f'  Unchanged: {total-improved-degraded}/{total}')

    if improved > degraded:
        print('  Edge refinement LOOKS PROMISING!')
    else:
        print('  Edge refinement not clearly beneficial')


if __name__ == '__main__':
    main()