"""Gradio app: detect cells in a fluorescence image and return red-channel
grayscale images with cell + nucleus outlines drawn in yellow.

One output image is produced per detected cell, matching the documentation
style: grayscale background + two concentric yellow outlines, nothing else.
"""
from __future__ import annotations

import os

import cv2
import gradio as gr
import numpy as np
from PIL import Image

from quantification import analyze_image

DEFAULT_N_CELLS = 5
DEFAULT_DILATION_RADIUS = 12
OUTLINE_COLOR_BGR_AS_RGB = (255, 255, 0)  # yellow in RGB
OUTLINE_THICKNESS = 2

EXAMPLES_DIR = os.path.join(os.path.dirname(__file__), "examples")
DEFAULT_EXAMPLE = os.path.join(EXAMPLES_DIR, "Picture1.jpg")


def _ensure_rgb(arr: np.ndarray) -> np.ndarray:
    if arr.ndim == 2:
        arr = np.stack([arr, arr, arr], axis=-1)
    if arr.shape[2] == 4:
        arr = arr[..., :3]
    if arr.dtype != np.uint8:
        arr = np.clip(arr, 0, 255).astype(np.uint8)
    return arr


def _draw_cell_outline(
    gray_rgb: np.ndarray,
    cell_mask: np.ndarray,
    nucleus_mask: np.ndarray,
) -> np.ndarray:
    """Draw the outer (cell) and inner (nucleus) outlines on a copy of `gray_rgb`."""
    out = gray_rgb.copy()
    for mask in (cell_mask, nucleus_mask):
        contours, _ = cv2.findContours(
            mask.astype(np.uint8), cv2.RETR_EXTERNAL, cv2.CHAIN_APPROX_SIMPLE
        )
        cv2.drawContours(
            out, contours, -1, OUTLINE_COLOR_BGR_AS_RGB, OUTLINE_THICKNESS
        )
    return out


def process_image(image_path: str | None, n_cells: int, dilation_radius: int):
    """Return a list of one annotated image per detected cell."""
    if image_path is None:
        return []

    image_pil = Image.open(image_path).convert("RGB")
    image_rgb = _ensure_rgb(np.array(image_pil))

    # Background for outputs: the red channel rendered as a grayscale RGB.
    red = image_rgb[..., 0]
    gray_rgb = np.stack([red, red, red], axis=-1)

    cells = analyze_image(
        image_rgb,
        n_cells=int(n_cells),
        dilation_radius=int(dilation_radius),
    )

    return [
        _draw_cell_outline(gray_rgb, c.cell_mask, c.nucleus_mask) for c in cells
    ]


def build_demo() -> gr.Blocks:
    description = (
        "Upload a fluorescence image (RGB: blue = nuclei, red = cytoplasm). "
        "The app detects representative cells and returns the red channel as "
        "grayscale with the cell + nucleus boundaries drawn in yellow — one "
        "output image per cell."
    )

    with gr.Blocks(title="Cell Boundary Detection") as demo:
        gr.Markdown("# Cell Boundary Detection")
        gr.Markdown(description)

        with gr.Row():
            with gr.Column(scale=1):
                image_in = gr.Image(
                    label="Input image",
                    type="filepath",
                    value=DEFAULT_EXAMPLE if os.path.exists(DEFAULT_EXAMPLE) else None,
                )
                n_cells_slider = gr.Slider(
                    minimum=1,
                    maximum=10,
                    value=DEFAULT_N_CELLS,
                    step=1,
                    label="Number of cells",
                )
                dilation_slider = gr.Slider(
                    minimum=4,
                    maximum=30,
                    value=DEFAULT_DILATION_RADIUS,
                    step=1,
                    label="Cytoplasm ring thickness (pixels)",
                )
                run_btn = gr.Button("Detect cells", variant="primary")

            with gr.Column(scale=2):
                gallery = gr.Gallery(
                    label="Detected cells (one per image)",
                    columns=2,
                    height=620,
                    show_label=True,
                    object_fit="contain",
                )

        run_btn.click(
            fn=process_image,
            inputs=[image_in, n_cells_slider, dilation_slider],
            outputs=[gallery],
        )

        # Example images (other defaults from prior dataset).
        example_files = []
        if os.path.isdir(EXAMPLES_DIR):
            example_files = sorted(
                os.path.join(EXAMPLES_DIR, f)
                for f in os.listdir(EXAMPLES_DIR)
                if f.lower().endswith((".jpg", ".jpeg", ".png", ".tif", ".tiff"))
            )
        if example_files:
            gr.Examples(
                examples=[[p, DEFAULT_N_CELLS, DEFAULT_DILATION_RADIUS]
                          for p in example_files],
                inputs=[image_in, n_cells_slider, dilation_slider],
                outputs=[gallery],
                fn=process_image,
                cache_examples=False,
                label="Example images",
            )

        # Preload outputs for the default image on app start.
        if os.path.exists(DEFAULT_EXAMPLE):
            demo.load(
                fn=process_image,
                inputs=[image_in, n_cells_slider, dilation_slider],
                outputs=[gallery],
            )

    return demo


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