import argparse
import re
from pathlib import Path

import cv2
import numpy as np
import pandas as pd
import tifffile as tif
from scipy.ndimage import binary_fill_holes
from skimage.color import hdx_from_rgb, separate_stains
from skimage.segmentation import watershed
from skimage.util import img_as_float32, img_as_ubyte

from cleaning import INPAINT_RADIUS, _as_rgb_uint8, create_debris_mask
from ihc_qc import save_qc_image


IMAGE_EXTENSIONS = {".tif", ".tiff"}
DEFAULT_OUTPUT_PATH = Path("out/quantification/ihc_quantification.csv")
MAX_BOUNDARY_CONTACT_RATIO = 0.3
WATERSHED_MARKER_CORE_FRACTION = 0.55
WATERSHED_MINIMUM_RELATIVE_REGION_AREA = 0.5
CONCENTRATION_PATTERN = re.compile(
    r"^\d+(?:[.,]\d+)?(?:pM|nM|uM|µM|μM|mM)$",
    re.IGNORECASE,
)
MAGNIFICATION_PATTERN = re.compile(r"^\d+(?:[.,]\d+)?x$", re.IGNORECASE)
MARKER_PATTERN = re.compile(r"^M\d+(?:[.,]\d+)?$", re.IGNORECASE)


def parse_filename(image_path: str | Path) -> dict:
    parts = Path(image_path).stem.split()
    if len(parts) < 5:
        raise ValueError(
            f"Could not read metadata from filename: {Path(image_path).name}"
        )

    magnification_index = next(
        (
            index
            for index in range(len(parts) - 1, 1, -1)
            if MAGNIFICATION_PATTERN.fullmatch(parts[index])
        ),
        None,
    )
    if magnification_index is None or magnification_index < 3:
        raise ValueError(
            f"Could not read magnification from filename: {Path(image_path).name}"
        )

    trailing_tokens = parts[magnification_index + 1 :]
    number_indices = [
        index for index, token in enumerate(trailing_tokens) if token.isdigit()
    ]
    if len(number_indices) != 1:
        raise ValueError(
            f"Could not read a unique image number from filename: "
            f"{Path(image_path).name}"
        )

    number_index = number_indices[0]
    image_number = int(trailing_tokens[number_index])
    image_note = " ".join(
        token for index, token in enumerate(trailing_tokens) if index != number_index
    )

    metadata_end = magnification_index
    if MARKER_PATTERN.fullmatch(parts[magnification_index - 1]):
        metadata_end -= 1
    metadata_tokens = parts[2:metadata_end]
    treatment_tokens = [
        token
        for token in metadata_tokens
        if (token == "+" or not token.startswith("+"))
        and not CONCENTRATION_PATTERN.fullmatch(token)
    ]
    if not treatment_tokens:
        raise ValueError(
            f"Could not read treatment from filename: {Path(image_path).name}"
        )
    treatment = re.sub(r"\s*\+\s*", "+", " ".join(treatment_tokens))

    return {
        "image_name": Path(image_path).name,
        "image_path": str(Path(image_path).resolve()),
        "treatment": treatment,
        "image_note": image_note,
        "image_number": image_number,
    }


def read_original_rgb(image_path: str) -> np.ndarray:
    image_path = Path(image_path)
    if image_path.suffix.lower() in {".tif", ".tiff"}:
        return tif.imread(image_path)


def extract_dab(rgb: np.ndarray) -> np.ndarray:
    stains = separate_stains(rgb, hdx_from_rgb)
    dab = stains[..., 1].astype(np.float32)

    return dab


def split_touching_spheroids(
    tissue_mask: np.ndarray,
    min_area: int,
    marker_core_fraction: float = WATERSHED_MARKER_CORE_FRACTION,
    minimum_relative_region_area: float = WATERSHED_MINIMUM_RELATIVE_REGION_AREA,
) -> np.ndarray:
    if not 0 < marker_core_fraction <= 1:
        raise ValueError(
            "marker_core_fraction must be greater than 0 and at most 1."
        )
    if not 0 <= minimum_relative_region_area <= 1:
        raise ValueError(
            "minimum_relative_region_area must be between 0 and 1."
        )

    count, components, stats, _ = cv2.connectedComponentsWithStats(
        tissue_mask.astype(np.uint8),
        connectivity=8,
    )
    kept_components = [
        component
        for component in range(1, count)
        if stats[component, cv2.CC_STAT_AREA] >= min_area
    ]
    filtered_mask = np.isin(components, kept_components)

    distance = cv2.distanceTransform(
        filtered_mask.astype(np.uint8),
        cv2.DIST_L2,
        5,
    )
    marker_mask = np.zeros(filtered_mask.shape, dtype=np.uint8)

    for component in kept_components:
        component_mask = components == component
        maximum_distance = distance[component_mask].max()
        marker_mask[
            component_mask & (distance >= marker_core_fraction * maximum_distance)
        ] = 1

    _, markers = cv2.connectedComponents(marker_mask, connectivity=8)
    candidate_labels = watershed(
        -distance,
        markers,
        mask=filtered_mask,
    ).astype(np.int32)

    labels = np.zeros(candidate_labels.shape, dtype=np.int32)
    next_label = 1

    for component in kept_components:
        component_mask = components == component
        regions = [
            region
            for region in np.unique(candidate_labels[component_mask])
            if region != 0
        ]
        region_areas = [
            int((candidate_labels[component_mask] == region).sum())
            for region in regions
        ]

        split_is_balanced = len(regions) > 1 and min(
            region_areas
        ) >= minimum_relative_region_area * max(region_areas)
        if not split_is_balanced:
            labels[component_mask] = next_label
            next_label += 1
            continue

        for region in regions:
            labels[component_mask & (candidate_labels == region)] = next_label
            next_label += 1

    return labels


def segment_spheroids(
    preview: np.ndarray,
    marker_core_fraction: float = WATERSHED_MARKER_CORE_FRACTION,
    minimum_relative_region_area: float = WATERSHED_MINIMUM_RELATIVE_REGION_AREA,
) -> np.ndarray:
    gray = cv2.cvtColor(preview, cv2.COLOR_RGB2GRAY)

    # thresholding
    C = 15
    block_size = round(min(gray.shape) * 0.14)
    block_size = block_size - 1 if block_size % 2 == 0 else block_size
    tissue_thresh = cv2.adaptiveThreshold(
        gray, 255, cv2.ADAPTIVE_THRESH_GAUSSIAN_C, cv2.THRESH_BINARY_INV, block_size, C
    )

    # filtering
    opening_kernel = 5  # largest noise I want removed
    opening_kernel = cv2.getStructuringElement(
        cv2.MORPH_ELLIPSE,
        (opening_kernel, opening_kernel),
    )
    tissue_thresh = cv2.morphologyEx(
        tissue_thresh,
        cv2.MORPH_OPEN,
        opening_kernel,
    )

    closing_kernel = 51  # smallest element I want to connect
    closing_kernel = cv2.getStructuringElement(
        cv2.MORPH_ELLIPSE,
        (closing_kernel, closing_kernel),
    )
    tissue_thresh = cv2.morphologyEx(
        tissue_thresh,
        cv2.MORPH_CLOSE,
        closing_kernel,
    )

    tissue_filled = binary_fill_holes(tissue_thresh > 0).astype(np.uint8) * 255

    min_area_fraction = (
        0.0025  # component must occupy at least 0.25% of the complete image
    )
    min_area = round(preview.shape[0] * preview.shape[1] * min_area_fraction)
    split_labels = split_touching_spheroids(
        tissue_filled > 0,
        min_area,
        marker_core_fraction=marker_core_fraction,
        minimum_relative_region_area=minimum_relative_region_area,
    )

    components = []
    for component in np.unique(split_labels):
        if component == 0:
            continue

        y, x = np.nonzero(split_labels == component)
        if len(x) >= min_area:
            components.append((component, y.mean(), x.mean()))

    components.sort(key=lambda component: (component[1], component[2]))

    labels = np.zeros(split_labels.shape, dtype=np.int32)
    for spheroid_id, (component, _, _) in enumerate(components, start=1):
        labels[split_labels == component] = spheroid_id

    return labels


def measure_spheroids(
    labels: np.ndarray,
    debris_mask: np.ndarray,
    dab: np.ndarray,
    positive_threshold: None | float = None,
    max_boundary_contact_ratio: float = MAX_BOUNDARY_CONTACT_RATIO,
) -> list[dict]:
    if max_boundary_contact_ratio < 0:
        raise ValueError("max_boundary_contact_ratio must be nonnegative.")

    measurements = []
    debris = debris_mask > 0
    image_boundary = np.zeros(labels.shape, dtype=bool)
    image_boundary[[0, -1], :] = True
    image_boundary[:, [0, -1]] = True

    for spheroid_id in np.unique(labels):
        if spheroid_id == 0:
            continue

        spheroid = labels == spheroid_id
        valid = spheroid & ~debris

        spheroid_area = int(spheroid.sum())
        valid_area = int(valid.sum())
        debris_area = int((spheroid & debris).sum())
        boundary_contact_px = int((spheroid & image_boundary).sum())
        equivalent_diameter = float(2 * np.sqrt(spheroid_area / np.pi))
        boundary_contact_ratio = boundary_contact_px / equivalent_diameter
        touches_image_boundary = boundary_contact_px > 0
        exceeds_boundary_tolerance = bool(
            boundary_contact_ratio > max_boundary_contact_ratio
        )

        if valid_area == 0:
            continue

        dab_values = dab[valid]

        result = {
            "spheroid_id": int(spheroid_id),
            "touches_image_boundary": touches_image_boundary,
            "exceeds_boundary_tolerance": exceeds_boundary_tolerance,
            "boundary_contact_ratio": boundary_contact_ratio,
            "spheroid_area_px": spheroid_area,
            "valid_area_px": valid_area,
            "debris_area_px": debris_area,
            "debris_fraction": debris_area / spheroid_area,
            "mean_dab": dab_values.mean(),
            "median_dab": np.median(dab_values),
            "p75_dab": np.percentile(dab_values, 75),
            "p90_dab": np.percentile(dab_values, 90),
            "p95_dab": np.percentile(dab_values, 95),
            "p99_dab": np.percentile(dab_values, 99),
        }

        if positive_threshold is not None:
            positive = valid & (dab >= positive_threshold)
            positive_values = dab[positive]

            result["positive_fraction"] = positive.sum() / valid.sum()
            result["positive_area_px"] = positive.sum()
            result["positive_mean"] = (
                positive_values.mean() if positive_values.size else np.nan
            )

        measurements.append(result)

    return measurements


def quantify_image(
    image_path: str | Path,
    qc_path: str | Path | None = None,
    positive_threshold: float | None = None,
    max_boundary_contact_ratio: float = MAX_BOUNDARY_CONTACT_RATIO,
    marker_core_fraction: float = WATERSHED_MARKER_CORE_FRACTION,
    minimum_relative_region_area: float = WATERSHED_MINIMUM_RELATIVE_REGION_AREA,
) -> list[dict]:
    metadata = parse_filename(image_path)
    original = read_original_rgb(image_path)
    processing_preview = _as_rgb_uint8(original)
    qc_preview = img_as_ubyte(original)

    debris_mask = create_debris_mask(qc_preview)
    cleaned_preview = cv2.inpaint(
        processing_preview,
        debris_mask,
        INPAINT_RADIUS,
        cv2.INPAINT_TELEA,
    )
    spheroid_labels = segment_spheroids(
        cleaned_preview,
        marker_core_fraction=marker_core_fraction,
        minimum_relative_region_area=minimum_relative_region_area,
    )
    if not spheroid_labels.max():
        raise ValueError(
            "No spheroids were detected. Review the QC segmentation settings."
        )

    float_img = img_as_float32(original)
    dab = extract_dab(float_img)
    measurements = measure_spheroids(
        spheroid_labels,
        debris_mask,
        dab,
        positive_threshold=positive_threshold,
        max_boundary_contact_ratio=max_boundary_contact_ratio,
    )

    if qc_path is not None:
        boundary_spheroid_ids = {
            measurement["spheroid_id"]
            for measurement in measurements
            if measurement["exceeds_boundary_tolerance"]
        }
        save_qc_image(
            qc_path,
            qc_preview,
            spheroid_labels,
            debris_mask,
            dab,
            boundary_spheroid_ids=boundary_spheroid_ids,
            positive_threshold=positive_threshold,
        )

    return [{**metadata, **measurement} for measurement in measurements]


def find_images(data_path: str | Path) -> list[Path]:
    data_path = Path(data_path)
    if data_path.is_file():
        if data_path.suffix.lower() not in IMAGE_EXTENSIONS:
            raise ValueError(f"Input is not a TIFF image: {data_path}")
        image_paths = [data_path]
    elif data_path.is_dir():
        image_paths = sorted(
            path
            for path in data_path.rglob("*")
            if path.is_file() and path.suffix.lower() in IMAGE_EXTENSIONS
        )
    else:
        raise FileNotFoundError(f"Data path does not exist: {data_path}")

    if not image_paths:
        raise ValueError(f"No TIFF images found in: {data_path}")

    return image_paths


def quantify_data(
    data_path: str | Path,
    qc_dir: str | Path | None = None,
    positive_threshold: float | None = None,
    max_boundary_contact_ratio: float = MAX_BOUNDARY_CONTACT_RATIO,
    marker_core_fraction: float = WATERSHED_MARKER_CORE_FRACTION,
    minimum_relative_region_area: float = WATERSHED_MINIMUM_RELATIVE_REGION_AREA,
) -> pd.DataFrame:
    measurements = []
    qc_dir = Path(qc_dir) if qc_dir is not None else None

    for image_path in find_images(data_path):
        qc_path = qc_dir / f"{image_path.stem}_qc.png" if qc_dir is not None else None
        image_measurements = quantify_image(
            image_path,
            qc_path,
            positive_threshold=positive_threshold,
            max_boundary_contact_ratio=max_boundary_contact_ratio,
            marker_core_fraction=marker_core_fraction,
            minimum_relative_region_area=minimum_relative_region_area,
        )
        measurements.extend(image_measurements)
        message = f"{image_path}: {len(image_measurements)} spheroids"
        if qc_path is not None:
            message += f"; QC: {qc_path}"
        print(message)

    if not measurements:
        raise ValueError("No valid spheroid measurements were produced.")

    return pd.DataFrame(measurements)


def parse_args() -> argparse.Namespace:
    parser = argparse.ArgumentParser(
        description="Quantify spheroid DAB intensity and save one row per spheroid."
    )
    parser.add_argument(
        "data_path",
        type=Path,
        help="TIFF image or directory containing TIFF images.",
    )
    parser.add_argument(
        "--output",
        type=Path,
        default=DEFAULT_OUTPUT_PATH,
        help=f"Output CSV path (default: {DEFAULT_OUTPUT_PATH}).",
    )
    parser.add_argument(
        "--positive-threshold",
        type=float,
        default=None,
        help="Optional DAB-positive threshold used for measurements and QC.",
    )
    parser.add_argument(
        "--max-boundary-contact-ratio",
        type=float,
        default=MAX_BOUNDARY_CONTACT_RATIO,
        help=(
            "Maximum tolerated boundary contact ratio before a spheroid is flagged "
            f"(default: {MAX_BOUNDARY_CONTACT_RATIO})."
        ),
    )
    parser.add_argument(
        "--watershed-marker-core-fraction",
        type=float,
        default=WATERSHED_MARKER_CORE_FRACTION,
        help=(
            "Distance-transform fraction used to create watershed markers "
            f"(default: {WATERSHED_MARKER_CORE_FRACTION})."
        ),
    )
    parser.add_argument(
        "--watershed-minimum-relative-region-area",
        type=float,
        default=WATERSHED_MINIMUM_RELATIVE_REGION_AREA,
        help=(
            "Smallest accepted watershed region relative to the largest region "
            f"(default: {WATERSHED_MINIMUM_RELATIVE_REGION_AREA})."
        ),
    )
    return parser.parse_args()


def main() -> None:
    args = parse_args()
    qc_dir = args.output.parent / "qc"
    results = quantify_data(
        args.data_path,
        qc_dir,
        positive_threshold=args.positive_threshold,
        max_boundary_contact_ratio=args.max_boundary_contact_ratio,
        marker_core_fraction=args.watershed_marker_core_fraction,
        minimum_relative_region_area=args.watershed_minimum_relative_region_area,
    )
    args.output.parent.mkdir(parents=True, exist_ok=True)
    results.to_csv(args.output, index=False)
    print(f"Saved {len(results)} spheroids to {args.output}")


if __name__ == "__main__":
    main()