"""Snake Active Contour Bbox Refinement
========================================
Classic CV approach: run active contour on image gradient inside each YOLO box,
fit the snake to the goat's actual boundary, then fit a tighter box.

Zero training cost. CPU-based.
"""
import sys, os
import numpy as np
from PIL import Image
from tqdm import tqdm
from scipy import ndimage

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

from ultralytics import YOLO


def snake_refine(image, box, n_iter=30, alpha=0.01, beta=0.1, gamma=0.001):
    """Refine a single bounding box using active contour.

    Args:
        image: PIL Image (full image)
        box: [x1, y1, x2, y2] in pixel coords
        n_iter: snake iterations
        alpha: continuity (smoothness along contour)
        beta: curvature (bending penalty)
        gamma: step size

    Returns:
        refined_box: [x1, y1, x2, y2] or None if failed
    """
    x1, y1, x2, y2 = box
    w, h = x2 - x1, y2 - y1

    if w < 8 or h < 8:
        return None

    # Expand ROI for context
    pad = 0.3
    rx1 = max(0, int(x1 - w * pad))
    ry1 = max(0, int(y1 - h * pad))
    rx2 = min(image.size[0], int(x2 + w * pad))
    ry2 = min(image.size[1], int(y2 + h * pad))

    roi = image.crop((rx1, ry1, rx2, ry2))
    roi_gray = np.array(roi.convert('L'), dtype=float)

    # Edge map (negative gradient magnitude attracts snake)
    gy, gx = np.gradient(roi_gray)
    edge = np.sqrt(gx**2 + gy**2)
    edge = edge / (edge.max() + 1e-8)

    # Initialize snake as ellipse inside the original box (in ROI coords)
    roi_h, roi_w = roi_gray.shape
    cx_roi = (x1 + x2) / 2 - rx1
    cy_roi = (y1 + y2) / 2 - ry1
    rw = w * 0.9  # Slightly smaller than YOLO box
    rh = h * 0.9

    # Create initial contour (ellipse with 40 points)
    n_pts = 40
    theta = np.linspace(0, 2*np.pi, n_pts)
    snake_x = cx_roi + rw/2 * np.cos(theta)
    snake_y = cy_roi + rh/2 * np.sin(theta)

    # Clip to ROI
    snake_x = np.clip(snake_x, 1, roi_w - 2)
    snake_y = np.clip(snake_y, 1, roi_h - 2)

    # Precompute external energy: negative edge + distance to edge
    external = -edge
    # Add small inward balloon force
    balloon = -0.05

    # Snake evolution
    for _ in range(n_iter):
        # Get snake as integer coords for interpolation
        sx = np.clip(snake_x.astype(int), 0, roi_w - 1)
        sy = np.clip(snake_y.astype(int), 0, roi_h - 1)

        # External force: gradient of edge map at snake points
        # Use simple gradient descent on external energy
        for i in range(n_pts):
            x, y = int(snake_x[i]), int(snake_y[i])
            x = max(1, min(roi_w - 2, x))
            y = max(1, min(roi_h - 2, y))

            # Local gradient of external energy
            dx = external[y, min(x+1, roi_w-1)] - external[y, max(x-1, 0)]
            dy = external[min(y+1, roi_h-1), x] - external[max(y-1, 0), x]

            # Internal force: tension + rigidity
            prev_i = (i - 1) % n_pts
            next_i = (i + 1) % n_pts

            # Tension: pull towards neighbors (alpha)
            tension = alpha * (snake_x[prev_i] + snake_x[next_i] - 2*snake_x[i])

            # Update
            snake_x[i] += gamma * (-dx * 50 + tension) + balloon * (snake_x[i] - cx_roi) * 0.01
            snake_y[i] += gamma * (-dy * 50 + tension) + balloon * (snake_y[i] - cy_roi) * 0.01

        # Re-parameterize: resample evenly
        # Simple: just clip and continue
        snake_x = np.clip(snake_x, 1, roi_w - 2)
        snake_y = np.clip(snake_y, 1, roi_h - 2)

    # Fit bounding box to refined contour
    new_x1 = snake_x.min() + rx1
    new_y1 = snake_y.min() + ry1
    new_x2 = snake_x.max() + rx1
    new_y2 = snake_y.max() + ry1

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


def compute_iou(b1, b2):
    x1 = max(b1[0], b2[0]); y1 = max(b1[1], b2[1])
    x2 = min(b1[2], b2[2]); y2 = 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():
    yolo = 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')])

    # Test on first 5 images only (snake is slow)
    test_files = val_files[:5]

    print('Snake refinement test on 5 images:')
    for img_file in test_files:
        img_path = os.path.join(val_img_dir, img_file)
        img = Image.open(img_path)

        # YOLO predict
        results = yolo.predict(img, imgsz=1536, conf=0.25, verbose=False)

        if not results or len(results[0].boxes) == 0:
            continue

        boxes = results[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]
                    ])

        # Refine each box
        for i, box in enumerate(boxes[:10]):  # Max 10 per image
            refined = snake_refine(img, box, n_iter=20)
            if refined is not None:
                # Measure improvement
                best_gt_iou = 0
                for gt in gt_boxes:
                    iou = compute_iou(box, gt)
                    if iou > best_gt_iou:
                        best_gt_iou = iou

                best_ref_iou = 0
                for gt in gt_boxes:
                    iou = compute_iou(refined, gt)
                    if iou > best_ref_iou:
                        best_ref_iou = iou

                if best_ref_iou > best_gt_iou + 0.01:
                    delta = best_ref_iou - best_gt_iou
                    print(f'  {img_file} box[{i}]: IoU {best_gt_iou:.3f} -> {best_ref_iou:.3f} (+{delta:.3f}) ✓')
                elif best_ref_iou < best_gt_iou - 0.01:
                    delta = best_ref_iou - best_gt_iou
                    print(f'  {img_file} box[{i}]: IoU {best_gt_iou:.3f} -> {best_ref_iou:.3f} ({delta:.3f}) ✗')

    print('\nSnake test complete. CPU refinement works, no GPU needed.')


if __name__ == '__main__':
    main()