"""
Improved larva counting application for Hugging Face Spaces.

This version exposes additional parameters through Gradio sliders to allow
the user to tune the preprocessing and contour‑filtering steps.  It also
modularises some of the hard‑coded constants in the original implementation.

Parameters exposed via the UI:

* **Umbral**: Threshold value for binary segmentation.  Setting 0 triggers
  Otsu's automatic threshold.
* **Área mínima / máxima**: Rejects contours outside this size range.
* **Forma mínima / máxima**: Controls the acceptable ellipse axis ratio of
  detected contours; values between 0 and 1.  Useful for eliminating
  elongated or highly circular artefacts.
* **Solidez mínima**: Rejects contours with low solidity (area divided by
  convex hull area), which helps to discard irregular flour particles.
* **Kernel morfológico**: Size of the structuring element used during the
  morphological opening step; larger kernels remove more noise but can
  merge nearby larvae.
* **Iteraciones morfológicas**: Number of times the morphological opening is
  applied.
* **CLAHE clipLimit** and **tileGridSize**: Adjust the contrast limited
  adaptive histogram equalisation used to emphasise bright dots in the
  image.
* **Gauss blur (fondo)**: Kernel size for the Gaussian blur that estimates
  the background.  Larger kernels remove broader illumination gradients.
* **Median blur**: Kernel size for the median filter used to smooth the
  preprocessed image.
* **Border recorte**: Number of pixels trimmed from each edge of the
  resized image.  Adjust if the frame contains noise or if larvae are
  close to the border.

The counting logic remains similar to the original: after thresholding and
morphological filtering, contours are filtered by area, shape and
solidity.  For contours larger than the maximum single‑larva area, the
estimated number of larvae is computed by dividing by the median area of
small contours.
"""

import gradio as gr
import cv2
import numpy as np
import statistics

# ----- Constants -----
# Target image size; images are resized to this resolution for processing.
IMG_W = 2047
IMG_H = 1148

# Default parameters used when sliders are at their initial positions.
DEFAULT_BORDER = 6
DEFAULT_CLIP = 2.5
DEFAULT_TILE = 8
DEFAULT_BG_BLUR = 25
DEFAULT_MEDIAN_BLUR = 3
DEFAULT_SHAPE_MIN = 0.55
DEFAULT_SHAPE_MAX = 0.95
DEFAULT_MIN_SOLIDITY = 0.7

# Global state for accumulated count.  In a production system you could use
# gr.State or another mechanism to avoid globals.
global_count = 0
median_single_area = None


def ellipse_ratio(cnt: np.ndarray) -> float | None:
    """Return the minor/major axis ratio of the best‑fit ellipse for a contour.

    If the contour has fewer than 5 points (required by cv2.fitEllipse),
    returns ``None``.

    Args:
        cnt: Contour as returned by ``cv2.findContours``.

    Returns:
        float between 0 and 1, or ``None`` if fitting fails.
    """
    if len(cnt) < 5:
        return None
    try:
        (_, _), (MA, ma), _ = cv2.fitEllipse(cnt)
    except cv2.error:
        return None
    # ratio of minor axis to major axis
    return min(MA, ma) / max(MA, ma)


def contour_solidity(cnt: np.ndarray) -> float:
    """Compute the solidity of a contour (area divided by convex hull area)."""
    area = cv2.contourArea(cnt)
    if area <= 0:
        return 0.0
    hull = cv2.convexHull(cnt)
    hull_area = cv2.contourArea(hull)
    if hull_area == 0:
        return 0.0
    return float(area) / float(hull_area)


def preprocess(
    image_bgr: np.ndarray,
    clip_limit: float = DEFAULT_CLIP,
    tile_grid_size: int = DEFAULT_TILE,
    bg_blur: int = DEFAULT_BG_BLUR,
    median_blur: int = DEFAULT_MEDIAN_BLUR,
    border: int = DEFAULT_BORDER,
) -> tuple[np.ndarray, np.ndarray]:
    """Resize and enhance the input image.

    The function performs the following steps:

    1. Resize the image to ``IMG_W`` × ``IMG_H`` using bilinear interpolation.
    2. Crop ``border`` pixels from each side.
    3. Convert to grayscale.
    4. Apply CLAHE to emphasise bright points.
    5. Subtract a blurred background to remove gradients.
    6. Normalise to full 0‑255 range.
    7. Apply median filtering to reduce noise.

    Args:
        image_bgr: Original image in BGR colour space.
        clip_limit: CLAHE clip limit; higher values increase local contrast.
        tile_grid_size: Size of the grid for CLAHE (in pixels).  The same
            value is used for both dimensions.
        bg_blur: Kernel size (odd integer) for the Gaussian blur used to
            estimate the background.
        median_blur: Kernel size (odd integer) for the median filter.
        border: Number of pixels to trim from each edge after resizing.

    Returns:
        A tuple ``(sub, img)`` where ``sub`` is the processed grayscale
        image and ``img`` is the resized colour image (for overlaying
        detections).
    """
    # Resize to a consistent working resolution
    img = cv2.resize(image_bgr, (IMG_W, IMG_H), interpolation=cv2.INTER_LINEAR)

    # Crop the border region
    if border > 0:
        roi = img[border : IMG_H - border, border : IMG_W - border]
    else:
        roi = img.copy()

    # Convert to grayscale
    gray = cv2.cvtColor(roi, cv2.COLOR_BGR2GRAY)

    # Apply CLAHE (Contrast Limited Adaptive Histogram Equalization)
    tile_size = max(1, int(tile_grid_size))
    clahe = cv2.createCLAHE(clipLimit=float(clip_limit), tileGridSize=(tile_size, tile_size))
    gray = clahe.apply(gray)

    # Estimate background via Gaussian blur
    # Ensure the blur kernel is odd and at least 3
    bg_blur = int(bg_blur) if int(bg_blur) % 2 == 1 else int(bg_blur) + 1
    bg = cv2.GaussianBlur(gray, (bg_blur, bg_blur), 0)

    # Subtract background and normalise
    sub = cv2.subtract(gray, bg)
    sub = cv2.normalize(sub, None, 0, 255, cv2.NORM_MINMAX)

    # Median blur to reduce noise from flour granules
    m_size = int(median_blur) if int(median_blur) % 2 == 1 else int(median_blur) + 1
    sub = cv2.medianBlur(sub, m_size)

    return sub, img


def detect_larvas(
    image_bgr: np.ndarray,
    thresh_value: int = 10,
    min_area: int = 6,
    max_area_single: int = 40,
    shape_min: float = DEFAULT_SHAPE_MIN,
    shape_max: float = DEFAULT_SHAPE_MAX,
    min_solidity: float = DEFAULT_MIN_SOLIDITY,
    morph_kernel: int = 3,
    morph_iter: int = 1,
    clip_limit: float = DEFAULT_CLIP,
    tile_grid_size: int = DEFAULT_TILE,
    bg_blur: int = DEFAULT_BG_BLUR,
    median_blur: int = DEFAULT_MEDIAN_BLUR,
    border: int = DEFAULT_BORDER,
) -> tuple[np.ndarray, int]:
    """Detect and count larvae in the input image.

    Applies preprocessing, thresholding, morphological filtering and contour
    analysis.  Contours are filtered by area, ellipse axis ratio and
    solidity.  Large contours are divided by the median area of single
    larvae to estimate the number of larvae they contain.

    Args:
        image_bgr: Original image in BGR colour space.
        thresh_value: Threshold for binarisation; 0 triggers Otsu's method.
        min_area: Minimum contour area to accept (in pixels²).
        max_area_single: Maximum area considered as one larva (in pixels²).
        shape_min, shape_max: Acceptable range of ellipse axis ratio.
        min_solidity: Minimum solidity to accept a contour.
        morph_kernel: Size of the morphological kernel (odd integer).
        morph_iter: Number of morphological opening iterations.
        clip_limit, tile_grid_size, bg_blur, median_blur, border: Parameters
            passed to ``preprocess``.

    Returns:
        A tuple ``(output_image, total)`` where ``output_image`` is the
        colour image with contours drawn and ``total`` is the estimated
        number of larvae.
    """
    global median_single_area

    gray_proc, base_img = preprocess(
        image_bgr,
        clip_limit=clip_limit,
        tile_grid_size=tile_grid_size,
        bg_blur=bg_blur,
        median_blur=median_blur,
        border=border,
    )

    # Thresholding
    if thresh_value == 0:
        # Otsu's threshold
        _, th = cv2.threshold(gray_proc, 0, 255, cv2.THRESH_BINARY + cv2.THRESH_OTSU)
    else:
        _, th = cv2.threshold(gray_proc, int(thresh_value), 255, cv2.THRESH_BINARY)

    # Morphological opening to remove small noise
    k_size = int(morph_kernel)
    # Ensure kernel size is odd and >= 1
    if k_size < 1:
        k_size = 1
    if k_size % 2 == 0:
        k_size += 1
    kernel = np.ones((k_size, k_size), np.uint8)
    iters = max(1, int(morph_iter))
    th = cv2.morphologyEx(th, cv2.MORPH_OPEN, kernel, iterations=iters)

    # Find external contours
    contours, _ = cv2.findContours(th, cv2.RETR_EXTERNAL, cv2.CHAIN_APPROX_SIMPLE)

    good = []
    areas_all = []
    areas_single = []

    for c in contours:
        area = cv2.contourArea(c)
        # Discard extremely large regions (limit at 5000 px² as in original)
        if area < min_area or area > 5000:
            continue
        # Shape ratio filter
        ratio = ellipse_ratio(c)
        if ratio is None or not (shape_min <= ratio <= shape_max):
            continue
        # Solidity filter
        sol = contour_solidity(c)
        if sol < min_solidity:
            continue
        good.append(c)
        areas_all.append(area)
        if area <= max_area_single:
            areas_single.append(area)

    # Estimate median area of single larvae
    if areas_single:
        median_single_area = statistics.median_low(areas_single)
    elif areas_all:
        median_single_area = statistics.median_low(areas_all)
    else:
        # No detections
        out = base_img.copy()
        cv2.putText(out, "LARVAS: 0", (40, 80), cv2.FONT_HERSHEY_SIMPLEX, 2, (0, 0, 255), 3)
        return out, 0

    # Count larvae
    total = 0
    for c in good:
        a = cv2.contourArea(c)
        if a <= max_area_single:
            total += 1
        else:
            # Estimate number of larvae in a cluster
            est = int(round(a / median_single_area))
            total += max(1, est)

    # Draw contours on the original‑sized image
    out = base_img.copy()
    for c in good:
        # Shift contour coordinates by the border offset
        if border > 0:
            c_shifted = c + np.array([[border, border]])
        else:
            c_shifted = c
        cv2.drawContours(out, [c_shifted], -1, (0, 255, 0), 1)
    cv2.putText(out, f"LARVAS: {total}", (40, 80), cv2.FONT_HERSHEY_SIMPLEX, 2, (0, 0, 255), 3)

    return out, total


def process(
    image: np.ndarray,
    thresh: int,
    min_a: int,
    max_a: int,
    shape_min: float,
    shape_max: float,
    sol_min: float,
    morph_kernel: int,
    morph_iter: int,
    clip_limit: float,
    tile_grid: int,
    bg_blur: int,
    med_blur: int,
    border: int,
) -> tuple[np.ndarray | None, str, str]:
    """Gradio wrapper for larva detection.

    Accumulates the total count across multiple calls via the global
    ``global_count`` variable.
    """
    global global_count
    if image is None:
        return None, "No subiste imagen", f"Conteo total: {global_count}"

    img_bgr = cv2.cvtColor(np.array(image), cv2.COLOR_RGB2BGR)
    out_img_bgr, n = detect_larvas(
        img_bgr,
        thresh_value=int(thresh),
        min_area=int(min_a),
        max_area_single=int(max_a),
        shape_min=float(shape_min),
        shape_max=float(shape_max),
        min_solidity=float(sol_min),
        morph_kernel=int(morph_kernel),
        morph_iter=int(morph_iter),
        clip_limit=float(clip_limit),
        tile_grid_size=int(tile_grid),
        bg_blur=int(bg_blur),
        median_blur=int(med_blur),
        border=int(border),
    )
    global_count += n
    out_img_rgb = cv2.cvtColor(out_img_bgr, cv2.COLOR_BGR2RGB)
    return out_img_rgb, f"Larvas en la imagen: {n}", f"Conteo total: {global_count}"


def reset_count() -> str:
    """Reset the accumulated count and return a message."""
    global global_count
    global_count = 0
    return f"Conteo total: {global_count}"


# ----- Gradio interface -----
with gr.Blocks() as demo:
    gr.Markdown("## Contador de larvas – versión mejorada")
    with gr.Row():
        # Input column
        with gr.Column(scale=1):
            inp = gr.Image(label="Subí la foto")
            thresh = gr.Slider(0, 255, value=10, step=1, label="Umbral (0=Otsu auto)")
            min_area = gr.Slider(0, 300, value=6, step=1, label="Min área px²")
            max_area_single = gr.Slider(0, 5000, value=40, step=1, label="Máx área 1 larva px²")
            shape_min_s = gr.Slider(0.0, 1.0, value=DEFAULT_SHAPE_MIN, step=0.05, label="Forma mínima")
            shape_max_s = gr.Slider(0.0, 1.0, value=DEFAULT_SHAPE_MAX, step=0.05, label="Forma máxima")
            solidity_min_s = gr.Slider(0.0, 1.0, value=DEFAULT_MIN_SOLIDITY, step=0.05, label="Solidez mínima")
            morph_kernel_s = gr.Slider(3, 11, value=3, step=2, label="Kernel morfológico")
            morph_iter_s = gr.Slider(1, 3, value=1, step=1, label="Iteraciones morfológicas")
            cliplimit_s = gr.Slider(1.0, 5.0, value=DEFAULT_CLIP, step=0.5, label="CLAHE clipLimit")
            tilegrid_s = gr.Slider(4, 16, value=DEFAULT_TILE, step=2, label="CLAHE tileGridSize")
            bg_blur_s = gr.Slider(15, 55, value=DEFAULT_BG_BLUR, step=2, label="Gauss blur (fondo)")
            median_blur_s = gr.Slider(3, 11, value=DEFAULT_MEDIAN_BLUR, step=2, label="Median blur")
            border_s = gr.Slider(0, 20, value=DEFAULT_BORDER, step=1, label="Border recorte")
            btn = gr.Button("Procesar")
            btn_reset = gr.Button("Reset contador")
        # Output column
        with gr.Column(scale=1):
            out_img = gr.Image(label="Resultado")
            out_txt = gr.Textbox(label="Resultado individual")
            out_total = gr.Textbox(label="Resultado acumulado")
    # Bind button clicks
    btn.click(
        process,
        inputs=[
            inp,
            thresh,
            min_area,
            max_area_single,
            shape_min_s,
            shape_max_s,
            solidity_min_s,
            morph_kernel_s,
            morph_iter_s,
            cliplimit_s,
            tilegrid_s,
            bg_blur_s,
            median_blur_s,
            border_s,
        ],
        outputs=[out_img, out_txt, out_total],
    )
    btn_reset.click(reset_count, [], [out_total])

if __name__ == "__main__":
    demo.launch()