# utils.py
# ECG Condition Classifier - Image Preprocessing

from PIL import Image, ImageFilter, ImageEnhance, ImageOps
import numpy as np

TARGET_SIZE = (224, 224)
MIN_SIZE = 50


def load_and_preprocess(image):
    """
    Preprocesses an ECG image for ViT inference.

    Key steps:
    1. Convert to RGB
    2. Auto-crop whitespace/borders
    3. Convert to grayscale then back to RGB (removes color noise from
       pink/beige ECG paper, watermarks, etc.)
    4. Enhance contrast to make waveform lines pop
    5. Resize to 224x224 for ViT
    """
    if image is None:
        raise ValueError("No image provided.")

    # Step 1: Ensure RGB
    img = image.convert("RGB")

    # Step 2: Auto-crop borders/whitespace
    img = _autocrop(img)

    # Step 3: Grayscale normalization
    # Converts pink/beige/colored ECG paper to clean black-on-white.
    # This removes color distribution shift between training data and
    # stock photos, dramatically improving model confidence.
    gray = img.convert("L")

    # Step 4: Enhance contrast — make ECG traces sharper
    enhancer = ImageEnhance.Contrast(gray)
    gray = enhancer.enhance(2.5)

    # Step 5: Sharpen edges (waveform lines)
    gray = gray.filter(ImageFilter.SHARPEN)

    # Step 6: Back to RGB (model expects 3-channel input)
    img = gray.convert("RGB")

    # Step 7: Resize to ViT input size
    img = img.resize(TARGET_SIZE, Image.LANCZOS)

    return img


def _autocrop(img, threshold=240, padding=10):
    """
    Crops away white/light borders around the ECG tracing area.
    Helps the model focus on the waveform rather than margins.
    """
    try:
        gray = img.convert("L")
        arr = np.array(gray)

        # Find rows/cols that contain dark pixels (the ECG trace)
        mask = arr < threshold
        rows = np.any(mask, axis=1)
        cols = np.any(mask, axis=0)

        if not rows.any() or not cols.any():
            return img  # Nothing to crop

        rmin, rmax = np.where(rows)[0][[0, -1]]
        cmin, cmax = np.where(cols)[0][[0, -1]]

        # Add small padding back
        h, w = arr.shape
        rmin = max(0, rmin - padding)
        rmax = min(h, rmax + padding)
        cmin = max(0, cmin - padding)
        cmax = min(w, cmax + padding)

        return img.crop((cmin, rmin, cmax, rmax))
    except Exception:
        return img  # Fall back to original if crop fails


def validate_image(image):
    """
    Validates that the uploaded image is usable.
    Returns (is_valid: bool, message: str)
    """
    if image is None:
        return False, "No image uploaded. Please upload an ECG image."

    try:
        w, h = image.size
        if w < MIN_SIZE or h < MIN_SIZE:
            return False, f"Image too small ({w}x{h}). Please upload a higher resolution ECG."
    except Exception as e:
        return False, f"Invalid image: {str(e)}"

    return True, "Valid"


def get_image_info(image):
    """Returns a summary string of image properties (for debugging)."""
    if image is None:
        return "No image"
    return f"Size: {image.size}, Mode: {image.mode}"