"""
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 find_screen(image):
    """
    Detect the Concept 2 PM5 LCD screen region in the image.

    Returns (x, y, w, h) bounding box or None if not found.
    """
    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, (5, 5), 0)
    edges = cv2.Canny(blurred, 50, 150)

    candidates = []

    # Sweep brightness thresholds — screen brightness varies by
    # lighting conditions (ranges from ~100 in dim gyms to ~200+)
    for thresh_val in range(120, 200, 10):
        _, thresh = cv2.threshold(gray, thresh_val, 255, cv2.THRESH_BINARY)
        kern = cv2.getStructuringElement(cv2.MORPH_RECT, (11, 11))
        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.005 or area_ratio > 0.12:
                continue

            # Aspect ratio: LCD is roughly square (0.5 to 1.6)
            aspect = w / h
            if aspect < 0.5 or aspect > 1.6:
                continue

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

            # KEY: edge density — LCD with text > 0.03, plain surfaces < 0.01
            roi_edges = edges[y : y + h, x : x + w]
            edge_density = np.sum(roi_edges > 0) / rect_area
            if edge_density < 0.03:
                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))

    if not candidates:
        return None

    candidates.sort(key=lambda c: c[0], reverse=True)
    return candidates[0][1:]


def crop_screen(image_path, output_path, padding=15):
    """Load an image, find the screen, crop 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 = result

    # Add padding, clamped to image bounds
    x1 = max(0, x - padding)
    y1 = max(0, y - padding)
    x2 = min(w_img, x + w + padding)
    y2 = min(h_img, y + h + padding)

    cropped = image[y1:y2, x1:x2]
    cv2.imwrite(output_path, cropped, [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()