"""Cell fluorescence quantification.

Implements the protocol from the documentation:
    Cytoplasm Area   = Cell Area  - Nucleus Area
    Cytoplasm IntDen = Cell IntDen - Nucleus IntDen
    Mean Cytoplasm   = Cytoplasm IntDen / Cytoplasm Area

The nucleus is segmented from the blue (DAPI) channel.
The cell region is defined by dilating each nucleus by a fixed radius,
with a Voronoi-style constraint so cells do not overlap.
Fluorescence intensity is measured in the red channel.
"""
from __future__ import annotations

import warnings
from dataclasses import dataclass

import cv2
import numpy as np
from scipy import ndimage as ndi
from skimage import filters, measure, morphology, segmentation
from skimage.feature import peak_local_max

warnings.filterwarnings("ignore", category=FutureWarning)
warnings.filterwarnings("ignore", category=UserWarning)


@dataclass
class CellMeasurement:
    """Holds the measurements for a single detected cell."""

    cell_id: int
    centroid: tuple[float, float]
    cell_area: int
    nucleus_area: int
    cytoplasm_area: int
    cell_intden: float
    nucleus_intden: float
    cytoplasm_intden: float
    mean_cytoplasm: float
    cell_mask: np.ndarray
    nucleus_mask: np.ndarray


def _segment_nuclei(
    blue_channel: np.ndarray,
    min_area: int = 800,
    max_area: int = 30000,
) -> np.ndarray:
    """Segment nuclei from the blue (DAPI) channel.

    Returns an int32 label image with one positive integer per nucleus.
    """
    if blue_channel.dtype != np.uint8:
        blue_channel = blue_channel.astype(np.uint8)

    blurred = cv2.GaussianBlur(blue_channel, (7, 7), 1.8)

    # Otsu threshold on smoothed channel.
    try:
        thresh = filters.threshold_otsu(blurred)
    except ValueError:
        return np.zeros_like(blue_channel, dtype=np.int32)
    nuclei_mask = blurred > thresh

    # Clean small holes / objects and smooth boundaries.
    nuclei_mask = morphology.remove_small_holes(nuclei_mask, area_threshold=500)
    nuclei_mask = morphology.remove_small_objects(nuclei_mask, min_size=min_area)
    nuclei_mask = morphology.opening(nuclei_mask, morphology.disk(2))

    if not nuclei_mask.any():
        return np.zeros_like(blue_channel, dtype=np.int32)

    # Watershed splitting on the distance transform of the binary mask.
    distance = ndi.distance_transform_edt(nuclei_mask)
    smoothed_dist = cv2.GaussianBlur(distance.astype(np.float32), (5, 5), 1.0)
    coords = peak_local_max(
        smoothed_dist,
        min_distance=25,
        labels=nuclei_mask,
        threshold_abs=5,
    )
    if len(coords) == 0:
        return np.zeros_like(blue_channel, dtype=np.int32)

    marker_mask = np.zeros(distance.shape, dtype=bool)
    marker_mask[tuple(coords.T)] = True
    markers, _ = ndi.label(marker_mask)
    labels = segmentation.watershed(-smoothed_dist, markers, mask=nuclei_mask)

    # Drop nuclei outside the size band, re-label sequentially.
    props = measure.regionprops(labels)
    new_labels = np.zeros_like(labels, dtype=np.int32)
    next_id = 1
    for p in props:
        if min_area <= p.area <= max_area:
            new_labels[labels == p.label] = next_id
            next_id += 1
    return new_labels


def _expand_to_cells(
    nuclei_labels: np.ndarray,
    dilation_radius: int = 12,
) -> np.ndarray:
    """Expand each nucleus outward by `dilation_radius` pixels.

    Cells do not overlap: each background pixel is assigned to its nearest
    nucleus, then we keep only pixels within `dilation_radius` of any nucleus.
    """
    if nuclei_labels.max() == 0:
        return np.zeros_like(nuclei_labels, dtype=np.int32)

    background = nuclei_labels == 0
    # ndi.distance_transform_edt with return_indices gives nearest foreground pixel.
    dist, nearest_inds = ndi.distance_transform_edt(background, return_indices=True)
    expanded = nuclei_labels.copy()
    expanded[background] = nuclei_labels[tuple(nearest_inds[:, background])]

    within_radius = dist <= dilation_radius
    cell_labels = np.where(within_radius, expanded, 0).astype(np.int32)
    return cell_labels


def _measure_cells(
    nuclei_labels: np.ndarray,
    cell_labels: np.ndarray,
    red_channel: np.ndarray,
) -> list[CellMeasurement]:
    """Compute per-cell statistics from the red channel."""
    r = red_channel.astype(np.float32)

    nucleus_props = {p.label: p for p in measure.regionprops(nuclei_labels)}
    results: list[CellMeasurement] = []

    unique_labels = np.unique(cell_labels)
    for label in unique_labels:
        if label == 0 or label not in nucleus_props:
            continue
        cell_mask = cell_labels == label
        nuc_mask = nuclei_labels == label

        cell_area = int(cell_mask.sum())
        nuc_area = int(nuc_mask.sum())
        if cell_area <= nuc_area:
            continue
        cyto_area = cell_area - nuc_area

        cell_intden = float(r[cell_mask].sum())
        nuc_intden = float(r[nuc_mask].sum())
        cyto_intden = cell_intden - nuc_intden
        mean_cyto = cyto_intden / cyto_area if cyto_area > 0 else 0.0

        results.append(
            CellMeasurement(
                cell_id=int(label),
                centroid=nucleus_props[label].centroid,
                cell_area=cell_area,
                nucleus_area=nuc_area,
                cytoplasm_area=cyto_area,
                cell_intden=cell_intden,
                nucleus_intden=nuc_intden,
                cytoplasm_intden=cyto_intden,
                mean_cytoplasm=mean_cyto,
                cell_mask=cell_mask,
                nucleus_mask=nuc_mask,
            )
        )
    return results


def _select_representative_cells(
    measurements: list[CellMeasurement],
    nuclei_labels: np.ndarray,
    n_cells: int = 5,
    min_centroid_distance: float = 120.0,
) -> list[CellMeasurement]:
    """Select up to n_cells well-formed, spread-out cells."""
    if not measurements:
        return []

    H, W = nuclei_labels.shape
    border_margin = 5

    def touches_border(m: CellMeasurement) -> bool:
        ys, xs = np.where(m.nucleus_mask)
        return bool(
            ys.min() <= border_margin
            or xs.min() <= border_margin
            or ys.max() >= H - border_margin - 1
            or xs.max() >= W - border_margin - 1
        )

    def solidity(m: CellMeasurement) -> float:
        props = measure.regionprops(m.nucleus_mask.astype(np.int32))
        if not props:
            return 0.0
        return float(props[0].solidity)

    # 1) Filter out border-touching and badly shaped nuclei when we have enough.
    candidates = [m for m in measurements if not touches_border(m) and solidity(m) > 0.85]
    if len(candidates) < n_cells:
        candidates = [m for m in measurements if not touches_border(m)]
    if len(candidates) < n_cells:
        candidates = list(measurements)

    # 2) Bias toward cells with above-median mean cytoplasm intensity
    #    (avoids picking dim background-like regions).
    if len(candidates) > n_cells:
        means = np.array([m.mean_cytoplasm for m in candidates])
        threshold = float(np.median(means)) * 0.85
        above = [m for m in candidates if m.mean_cytoplasm >= threshold]
        if len(above) >= n_cells:
            candidates = above

    # 3) Prefer larger, better-formed nuclei.
    candidates.sort(key=lambda m: -m.nucleus_area)

    # 4) Pick spread-out cells.
    selected: list[CellMeasurement] = []
    used: list[tuple[float, float]] = []
    for m in candidates:
        cy, cx = m.centroid
        if all(np.hypot(cy - uy, cx - ux) > min_centroid_distance for uy, ux in used):
            selected.append(m)
            used.append((cy, cx))
        if len(selected) >= n_cells:
            break

    # 5) Fill remaining slots if needed (drop the spacing constraint).
    if len(selected) < n_cells:
        for m in candidates:
            if m not in selected:
                selected.append(m)
            if len(selected) >= n_cells:
                break

    return selected[:n_cells]


def analyze_image(
    image_rgb: np.ndarray,
    n_cells: int = 5,
    dilation_radius: int = 12,
) -> list[CellMeasurement]:
    """Full pipeline: segment, expand, measure, and select cells.

    Parameters
    ----------
    image_rgb : (H, W, 3) uint8 array
        The fluorescence image (blue = nuclei, red = cytoplasm marker).
    n_cells : int
        How many cells to report (default 5, matching the protocol).
    dilation_radius : int
        Pixel radius for the cytoplasm ring around each nucleus.
    """
    if image_rgb.ndim != 3 or image_rgb.shape[2] < 3:
        raise ValueError("Input must be an RGB image with at least 3 channels.")

    red = image_rgb[..., 0]
    blue = image_rgb[..., 2]

    nuclei_labels = _segment_nuclei(blue)
    if nuclei_labels.max() == 0:
        return []
    cell_labels = _expand_to_cells(nuclei_labels, dilation_radius=dilation_radius)
    all_cells = _measure_cells(nuclei_labels, cell_labels, red)
    selected = _select_representative_cells(all_cells, nuclei_labels, n_cells=n_cells)

    # Re-number the selected cells 1..N for display.
    renumbered: list[CellMeasurement] = []
    for new_id, m in enumerate(selected, start=1):
        renumbered.append(
            CellMeasurement(
                cell_id=new_id,
                centroid=m.centroid,
                cell_area=m.cell_area,
                nucleus_area=m.nucleus_area,
                cytoplasm_area=m.cytoplasm_area,
                cell_intden=m.cell_intden,
                nucleus_intden=m.nucleus_intden,
                cytoplasm_intden=m.cytoplasm_intden,
                mean_cytoplasm=m.mean_cytoplasm,
                cell_mask=m.cell_mask,
                nucleus_mask=m.nucleus_mask,
            )
        )
    return renumbered


def draw_overlay(
    image_rgb: np.ndarray,
    cells: list[CellMeasurement],
    outline_color: tuple[int, int, int] = (255, 255, 0),
    thickness: int = 2,
) -> np.ndarray:
    """Draw cell + nucleus outlines and labels onto a copy of the image."""
    out = image_rgb.copy()
    for m in cells:
        for mask in (m.cell_mask, m.nucleus_mask):
            contours, _ = cv2.findContours(
                mask.astype(np.uint8), cv2.RETR_EXTERNAL, cv2.CHAIN_APPROX_SIMPLE
            )
            cv2.drawContours(out, contours, -1, outline_color, thickness)

        cy, cx = m.centroid
        label = f"Cell {m.cell_id}"
        (tw, th), _ = cv2.getTextSize(label, cv2.FONT_HERSHEY_SIMPLEX, 0.7, 2)
        x0 = int(cx - tw // 2)
        y0 = int(cy + th // 2)
        cv2.rectangle(
            out,
            (x0 - 4, y0 - th - 4),
            (x0 + tw + 4, y0 + 4),
            (255, 255, 255),
            -1,
        )
        cv2.putText(
            out,
            label,
            (x0, y0),
            cv2.FONT_HERSHEY_SIMPLEX,
            0.7,
            (0, 0, 0),
            2,
        )
    return out


def cells_to_records(cells: list[CellMeasurement]) -> list[dict]:
    """Convert a list of CellMeasurement objects to plain dict records."""
    rows = []
    for m in cells:
        rows.append(
            {
                "Cell": f"Cell {m.cell_id}",
                "Total cell area": m.cell_area,
                "Nucleus area": m.nucleus_area,
                "Cytoplasm area": m.cytoplasm_area,
                "Total Cell IntDen": round(m.cell_intden, 2),
                "Nucleus IntDen": round(m.nucleus_intden, 2),
                "Cytoplasm IntDen": round(m.cytoplasm_intden, 2),
                "Mean Cytoplasm Fluorescence": round(m.mean_cytoplasm, 4),
            }
        )
    return rows