"""
Crop Concept 2 PM5 rowing machine screens from photos using OpenCV.

Detection strategy:
  The LCD screen has HIGH internal edge density (text/numbers/lines)
  compared to other bright regions (windows, walls, lockers).
  We threshold at multiple brightness levels, filter by edge density,
  aspect ratio, and size, then pick the best match.

Usage:
  python crop_screens.py [input_dir] [output_dir]
"""

import cv2
import numpy as np
import os
import glob
import sys


def order_corners(pts):
    """Order 4 points as [top-left, top-right, bottom-right, bottom-left]."""
    rect = np.zeros((4, 2), dtype="float32")
    s = pts.sum(axis=1)
    rect[0] = pts[np.argmin(s)]
    rect[2] = pts[np.argmax(s)]
    d = np.diff(pts, axis=1)
    rect[1] = pts[np.argmin(d)]
    rect[3] = pts[np.argmax(d)]
    return rect


def find_screen(image):
    """
    Detect the Concept 2 PM5 LCD screen region in the image.

    Returns (x, y, w, h, contour) or None if not found.
    The contour is the best-matching contour for perspective correction.
    """
    h_img, w_img = image.shape[:2]
    gray = cv2.cvtColor(image, cv2.COLOR_BGR2GRAY)

    # Pre-compute edge map for internal-content scoring
    blurred = cv2.GaussianBlur(gray, (11, 11), 0)
    edges = cv2.Canny(blurred, 80, 100)

    candidates = []

    # Sweep brightness thresholds — screen brightness varies by
    # lighting conditions (ranges from ~100 in dim gyms to ~200+)
    for thresh_val in range(70, 210, 10):
        _, thresh = cv2.threshold(gray, thresh_val, 255, cv2.THRESH_BINARY)
        kern = cv2.getStructuringElement(cv2.MORPH_RECT, (5, 5))
        thresh = cv2.morphologyEx(thresh, cv2.MORPH_CLOSE, kern)
        thresh = cv2.morphologyEx(thresh, cv2.MORPH_OPEN, kern)

        contours, _ = cv2.findContours(
            thresh, cv2.RETR_EXTERNAL, cv2.CHAIN_APPROX_SIMPLE
        )

        for cnt in contours:
            x, y, w, h = cv2.boundingRect(cnt)
            area = cv2.contourArea(cnt)
            rect_area = w * h
            if rect_area == 0:
                continue

            # Size: screen is a small-to-medium portion of the photo
            area_ratio = rect_area / (h_img * w_img)
            if area_ratio < 0.004480508227271387 or area_ratio > 0.13807760800032298:
                continue

            # Aspect ratio: LCD is roughly square
            aspect = w / h
            if aspect < 0.6831978184146027 or aspect > 1.9505294279578584:
                continue

            # Rectangularity
            rectangularity = area / rect_area
            if rectangularity < 0.6914579162415992:
                continue

            # KEY: edge density — LCD with text has high edge density
            roi_edges = edges[y : y + h, x : x + w]
            edge_density = np.sum(roi_edges > 0) / rect_area
            if edge_density < 0.012759310759672408:
                continue

            # Score: edge density * area * rectangularity
            # This favours text-rich regions that are large and well-shaped
            score = edge_density * area * rectangularity
            candidates.append((score, x, y, w, h, cnt))

    if not candidates:
        return None

    candidates.sort(key=lambda c: c[0], reverse=True)
    best = candidates[0]
    return best[1], best[2], best[3], best[4], best[5]


def perspective_correct(image, contour, dst_w, dst_h):
    """Warp the screen quadrilateral to a flat rectangle."""
    # Approximate contour to a polygon, tightening until we get 4 corners
    peri = cv2.arcLength(contour, True)
    for eps_mult in [0.02, 0.03, 0.05, 0.08, 0.10]:
        approx = cv2.approxPolyDP(contour, eps_mult * peri, True)
        if len(approx) == 4:
            break

    if len(approx) != 4:
        # Fall back to the minimum area rectangle corners
        rect = cv2.minAreaRect(contour)
        approx = cv2.boxPoints(rect).astype(np.float32)
    else:
        approx = approx.reshape(4, 2).astype(np.float32)

    src = order_corners(approx)
    dst = np.array(
        [[0, 0], [dst_w - 1, 0], [dst_w - 1, dst_h - 1], [0, dst_h - 1]],
        dtype="float32",
    )
    M = cv2.getPerspectiveTransform(src, dst)
    return cv2.warpPerspective(image, M, (dst_w, dst_h))


def crop_screen(image_path, output_path, padding=15):
    """Load an image, find the screen, perspective-correct and save it."""
    image = cv2.imread(image_path)
    if image is None:
        print(f"  ERROR: Could not read {image_path}")
        return False

    h_img, w_img = image.shape[:2]
    result = find_screen(image)

    if result is None:
        print(f"  SKIP:  No screen detected in {os.path.basename(image_path)}")
        return False

    x, y, w, h, contour = result

    # Use perspective correction to flatten the screen
    corrected = perspective_correct(image, contour, w + 2 * padding, h + 2 * padding)

    cv2.imwrite(output_path, corrected, [cv2.IMWRITE_JPEG_QUALITY, 95])
    print(
        f"  OK:    {os.path.basename(image_path)} -> {os.path.basename(output_path)}  ({w}x{h})"
    )
    return True


def main():
    if len(sys.argv) >= 3:
        input_dir = sys.argv[1]
        output_dir = sys.argv[2]
    elif len(sys.argv) == 2:
        input_dir = sys.argv[1]
        output_dir = os.path.join(input_dir, "cropped")
    else:
        input_dir = "/mnt/user-data/uploads"
        output_dir = "/mnt/user-data/outputs"

    os.makedirs(output_dir, exist_ok=True)

    images = sorted(
        glob.glob(os.path.join(input_dir, "*.JPEG"))
        + glob.glob(os.path.join(input_dir, "*.jpeg"))
        + glob.glob(os.path.join(input_dir, "*.jpg"))
        + glob.glob(os.path.join(input_dir, "*.JPG"))
    )

    if not images:
        print(f"No images found in {input_dir}")
        return

    print(f"Found {len(images)} images in {input_dir}\n")

    success = 0
    for img_path in images:
        name = os.path.splitext(os.path.basename(img_path))[0]
        out_path = os.path.join(output_dir, f"{name}_screen.jpg")
        if crop_screen(img_path, out_path):
            success += 1

    print(f"\nDone: {success}/{len(images)} screens cropped -> {output_dir}")


if __name__ == "__main__":
    main()