postprocessing/geometry.py

"""Contour extraction, sparse anchor selection, and curve rendering."""

from __future__ import annotations

import math
from typing import Any

import cv2
import numpy as np
from PIL import Image, ImageDraw

from .types import AdaptiveContourConfig, PixelSpacing, SplineContourConfig


def finite_xy(value: Any) -> np.ndarray:
    try:
        points = np.asarray(value, dtype=np.float64)
    except Exception:
        return np.empty((0, 2), dtype=np.float64)
    if points.ndim != 2 or points.shape[1] != 2:
        return np.empty((0, 2), dtype=np.float64)
    return points[np.isfinite(points).all(axis=1)]


def strip_duplicate_endpoint(points: Any, tolerance: float = 1e-6) -> np.ndarray:
    cleaned = finite_xy(points)
    if len(cleaned) > 1 and np.linalg.norm(cleaned[0] - cleaned[-1]) <= tolerance:
        cleaned = cleaned[:-1]
    return cleaned


def largest_external_contour(mask: np.ndarray) -> np.ndarray:
    """Extract the largest external boundary as ``[x, y]`` pixel points."""

    array = np.asarray(mask)
    if array.ndim != 2:
        raise ValueError(f"mask must be 2D, got shape {array.shape}")
    binary = (array > 0).astype(np.uint8)
    if not binary.any():
        raise ValueError("mask has no foreground pixels")
    contours, _ = cv2.findContours(binary, cv2.RETR_EXTERNAL, cv2.CHAIN_APPROX_NONE)
    if not contours:
        raise ValueError("no external contour could be extracted from mask")
    points = max(contours, key=cv2.contourArea).reshape(-1, 2).astype(np.float64)
    points = strip_duplicate_endpoint(points)
    if len(points) < 3:
        raise ValueError("largest mask component has fewer than 3 boundary points")
    return points


def myocardium_ring_boundaries(mask: np.ndarray) -> tuple[np.ndarray, np.ndarray]:
    """Return endocardial (inner) and epicardial (outer) boundaries of a ring."""

    array = np.asarray(mask)
    if array.ndim != 2:
        raise ValueError(f"mask must be 2D, got shape {array.shape}")
    binary = (array > 0).astype(np.uint8)
    if not binary.any():
        raise ValueError("myocardium mask has no foreground pixels")
    contours, hierarchy = cv2.findContours(binary, cv2.RETR_CCOMP, cv2.CHAIN_APPROX_NONE)
    if hierarchy is None or not contours:
        raise ValueError("no myocardium contours could be extracted")
    hierarchy = hierarchy[0]
    outer_indices = [index for index, item in enumerate(hierarchy) if int(item[3]) < 0]
    if not outer_indices:
        raise ValueError("myocardium ring has no outer boundary")
    outer_index = max(outer_indices, key=lambda index: cv2.contourArea(contours[index]))
    child_indices = [
        index for index, item in enumerate(hierarchy) if int(item[3]) == outer_index
    ]
    if not child_indices:
        raise ValueError("myocardium mask must contain an endocardial hole")
    inner_index = max(child_indices, key=lambda index: cv2.contourArea(contours[index]))
    inner = strip_duplicate_endpoint(contours[inner_index].reshape(-1, 2))
    outer = strip_duplicate_endpoint(contours[outer_index].reshape(-1, 2))
    if len(inner) < 3 or len(outer) < 3:
        raise ValueError("myocardium boundaries require at least 3 points")
    return inner.astype(np.float64), outer.astype(np.float64)


def xy_to_physical(points: Any, pixel_spacing: PixelSpacing) -> np.ndarray:
    cleaned = finite_xy(points)
    row_mm, column_mm = pixel_spacing
    return cleaned * np.asarray([column_mm, row_mm], dtype=np.float64)


def physical_to_xy(points: Any, pixel_spacing: PixelSpacing) -> np.ndarray:
    cleaned = finite_xy(points)
    row_mm, column_mm = pixel_spacing
    return cleaned / np.asarray([column_mm, row_mm], dtype=np.float64)


def _closed_indices(start: int, end: int, n_points: int) -> np.ndarray:
    if start < end:
        return np.arange(start, end + 1, dtype=np.int64)
    return np.concatenate(
        (np.arange(start, n_points, dtype=np.int64), np.arange(0, end + 1, dtype=np.int64))
    )


def _point_segment_distances(points: np.ndarray, start: np.ndarray, end: np.ndarray) -> np.ndarray:
    direction = end - start
    denominator = float(np.dot(direction, direction))
    if denominator <= 1e-12:
        return np.linalg.norm(points - start, axis=1)
    fraction = np.clip(((points - start) @ direction) / denominator, 0.0, 1.0)
    projected = start + fraction[:, None] * direction
    return np.linalg.norm(points - projected, axis=1)


def _turn_scores(points: np.ndarray, window: int = 3) -> np.ndarray:
    n_points = len(points)
    scores = np.zeros(n_points, dtype=np.float64)
    if n_points < 5:
        return scores
    window = max(1, min(int(window), max(1, n_points // 4)))
    for index in range(n_points):
        previous = points[index] - points[(index - window) % n_points]
        following = points[(index + window) % n_points] - points[index]
        previous_norm = float(np.linalg.norm(previous))
        following_norm = float(np.linalg.norm(following))
        if previous_norm <= 1e-6 or following_norm <= 1e-6:
            continue
        cosine = float(
            np.clip(np.dot(previous, following) / (previous_norm * following_norm), -1.0, 1.0)
        )
        scores[index] = abs(math.pi - math.acos(cosine))
    return scores


def _far_from_selected(
    points: np.ndarray, index: int, selected: set[int], minimum_spacing: float
) -> bool:
    if not selected or minimum_spacing <= 0:
        return True
    return all(
        float(np.linalg.norm(points[index] - points[selected_index])) >= minimum_spacing
        for selected_index in selected
    )


def _add_turning_points(
    points: np.ndarray, selected: set[int], target: int, minimum_spacing: float
) -> None:
    scores = _turn_scores(points)
    for index in np.argsort(scores)[::-1]:
        index = int(index)
        if index in selected:
            continue
        if _far_from_selected(points, index, selected, minimum_spacing) or len(selected) < 4:
            selected.add(index)
            if len(selected) >= target:
                return


def select_adaptive_control_indices(
    physical_points: Any, config: AdaptiveContourConfig
) -> np.ndarray:
    """Select deterministic anchors using the recovered task04 algorithm."""

    config.validate()
    points = strip_duplicate_endpoint(physical_points)
    n_points = len(points)
    if n_points < 3:
        raise ValueError("contour requires at least 3 finite points")
    if n_points <= config.max_control_points:
        return np.arange(n_points, dtype=np.int64)
    selected = {
        int(np.argmin(points[:, 0])),
        int(np.argmax(points[:, 0])),
        int(np.argmin(points[:, 1])),
        int(np.argmax(points[:, 1])),
    }
    _add_turning_points(points, selected, config.min_control_points, config.min_spacing_mm)
    if len(selected) < config.min_control_points:
        _add_turning_points(points, selected, config.min_control_points, 0.0)
    while len(selected) < config.max_control_points:
        ordered = sorted(selected)
        best_index = None
        best_distance = -1.0
        for position, start in enumerate(ordered):
            end = ordered[(position + 1) % len(ordered)]
            segment = _closed_indices(start, end, n_points)
            if len(segment) <= 2:
                continue
            interior = segment[1:-1]
            distances = _point_segment_distances(points[interior], points[start], points[end])
            for local_position in np.argsort(distances)[::-1]:
                index = int(interior[int(local_position)])
                if _far_from_selected(points, index, selected, config.min_spacing_mm):
                    distance = float(distances[int(local_position)])
                    if distance > best_distance:
                        best_distance = distance
                        best_index = index
                    break
        if best_index is None:
            break
        if len(selected) >= config.min_control_points and best_distance <= config.tolerance_mm:
            break
        selected.add(best_index)
    return np.asarray(sorted(selected), dtype=np.int64)


def render_tension_curve(
    control_points: Any, tension: float, samples_per_segment: int
) -> np.ndarray:
    """Render the CMR-annotator-compatible closed cubic Bezier curve."""

    points = strip_duplicate_endpoint(control_points)
    if len(points) < 3:
        raise ValueError("control contour requires at least 3 points")
    if len(points) < 4:
        samples = []
        for index, start in enumerate(points):
            end = points[(index + 1) % len(points)]
            for step in range(samples_per_segment):
                samples.append(start + step / samples_per_segment * (end - start))
        return np.asarray(samples, dtype=np.float64)
    handle_scale = max(0.0, float(tension)) / 6.0
    samples = []
    for index in range(len(points)):
        p0 = points[(index - 1) % len(points)]
        p1 = points[index]
        p2 = points[(index + 1) % len(points)]
        p3 = points[(index + 2) % len(points)]
        c1 = p1 + (p2 - p0) * handle_scale
        c2 = p2 - (p3 - p1) * handle_scale
        for step in range(samples_per_segment):
            t = step / samples_per_segment
            one_minus = 1.0 - t
            samples.append(
                one_minus**3 * p1
                + 3 * one_minus**2 * t * c1
                + 3 * one_minus * t**2 * c2
                + t**3 * p2
            )
    return np.asarray(samples, dtype=np.float64)


def render_periodic_bspline(
    physical_points: Any, config: SplineContourConfig
) -> np.ndarray:
    config.validate()
    points = strip_duplicate_endpoint(physical_points)
    if len(points) < 4:
        return points.copy()
    try:
        from scipy.interpolate import splprep, splev

        tck, _ = splprep(
            [points[:, 0], points[:, 1]], s=config.smoothing, per=True, k=3
        )
        # Preserve the historical SAX renderer, including its repeated closing
        # sample at parameter 1.0.
        parameter = np.linspace(0.0, 1.0, config.n_points)
        x_values, y_values = splev(parameter, tck)
        return np.column_stack([x_values, y_values]).astype(np.float64)
    except Exception:
        # The historical implementation also preserved the raw contour when a
        # degenerate spline could not be fit.
        return points.copy()


def compute_curvature(points: Any) -> np.ndarray:
    """Compute the January converter's discrete closed-contour curvature."""

    cleaned = finite_xy(points)
    if len(cleaned) < 3:
        raise ValueError("contour requires at least 3 finite points")
    padded = np.vstack([cleaned[-2:], cleaned, cleaned[:2]])
    dx = np.gradient(padded[:, 0])
    dy = np.gradient(padded[:, 1])
    ddx = np.gradient(dx)
    ddy = np.gradient(dy)
    numerator = np.abs(dx * ddy - dy * ddx)
    denominator = (dx**2 + dy**2 + 1e-10) ** 1.5
    return (numerator / denominator)[2:-2]


def curvature_sample(points: Any, n_control_points: int, min_weight: float = 0.1) -> np.ndarray:
    """Select the historical fixed-count curvature-weighted control points."""

    dense = finite_xy(points)
    if n_control_points < 3:
        raise ValueError("n_control_points must be at least 3")
    if n_control_points >= len(dense):
        return dense.copy()
    weights = np.abs(compute_curvature(dense)) + float(min_weight)
    weights /= weights.sum()
    cumulative = np.cumsum(weights)
    cumulative[-1] = 1.0
    positions = np.linspace(
        0.0, 1.0 - 1.0 / n_control_points, n_control_points
    )
    indices = np.unique(
        np.clip(np.searchsorted(cumulative, positions), 0, len(dense) - 1)
    )
    while len(indices) < n_control_points:
        gaps: list[tuple[int, int]] = []
        ordered = sorted(int(index) for index in indices)
        for position, start in enumerate(ordered):
            next_position = (position + 1) % len(ordered)
            end = ordered[next_position]
            if next_position == 0:
                end += len(dense)
            if end - start > 1:
                gaps.append((end - start, ((start + end) // 2) % len(dense)))
        if not gaps:
            break
        gaps.sort(reverse=True)
        indices = np.append(indices, gaps[0][1])
    return dense[np.sort(indices)[:n_control_points]]


def render_control_point_bspline(
    control_points: Any,
    n_points: int,
    smoothing: float = 0.0,
) -> np.ndarray:
    """Reconstruct sparse controls exactly as the January converter did."""

    points = strip_duplicate_endpoint(control_points)
    if len(points) < 3:
        return points.copy()
    closed = np.vstack([points, points[0]])
    try:
        from scipy.interpolate import splprep, splev

        tck, _ = splprep(
            [closed[:, 0], closed[:, 1]], s=float(smoothing), per=True
        )
        parameter = np.linspace(0.0, 1.0, int(n_points))
        x_values, y_values = splev(parameter, tck)
        return np.column_stack([x_values, y_values]).astype(np.float64)
    except Exception:
        return points.copy()


def contour_overlap_iou(
    first: Any, second: Any, resolution: int = 200
) -> float:
    """Rasterized contour IoU used by the historical sparse-point selector."""

    first_points = finite_xy(first)
    second_points = finite_xy(second)
    if len(first_points) < 3 or len(second_points) < 3:
        return 0.0
    all_points = np.vstack([first_points, second_points])
    minimum = all_points.min(axis=0) - 5.0
    maximum = all_points.max(axis=0) + 5.0
    extent = float(np.max(maximum - minimum))
    if not np.isfinite(extent) or extent <= 0:
        return 0.0
    scale = int(resolution) / extent

    def rasterize(points: np.ndarray) -> np.ndarray:
        scaled = ((points - minimum) * scale).astype(np.int32)
        mask = np.zeros((int(resolution), int(resolution)), dtype=np.uint8)
        cv2.fillPoly(mask, [scaled], 1)
        return mask

    first_mask = rasterize(first_points)
    second_mask = rasterize(second_points)
    intersection = int(np.logical_and(first_mask, second_mask).sum())
    union = int(np.logical_or(first_mask, second_mask).sum())
    return float(intersection / union) if union else 0.0


def select_sparse_spline_controls(
    smoothed_physical_points: Any, config: SplineContourConfig
) -> tuple[np.ndarray, np.ndarray]:
    """Select 6/8/10 controls and return their final dense reconstruction."""

    config.validate()
    dense = finite_xy(smoothed_physical_points)
    if len(dense) < 3:
        raise ValueError("contour requires at least 3 finite points")
    best_control: np.ndarray | None = None
    best_reconstruction: np.ndarray | None = None
    best_iou = -1.0
    for count in config.control_point_counts:
        control = curvature_sample(dense, count)
        reconstruction = render_control_point_bspline(
            control,
            n_points=config.n_points,
            smoothing=config.reconstruction_smoothing,
        )
        iou = contour_overlap_iou(dense, reconstruction)
        if iou > best_iou:
            best_control = control
            best_reconstruction = reconstruction
            best_iou = iou
        if iou >= config.control_point_iou_threshold:
            break
    if best_control is None or best_reconstruction is None:  # pragma: no cover
        raise ValueError("no sparse spline candidate could be constructed")
    return best_control, best_reconstruction


def rasterize_contour(points: Any, shape: tuple[int, int]) -> np.ndarray:
    cleaned = strip_duplicate_endpoint(points)
    rows, columns = shape
    image = Image.new("L", (int(columns), int(rows)), 0)
    if len(cleaned) >= 3:
        # Keep the July reference renderer's sub-pixel polygon semantics so
        # archived area and IoU metrics remain numerically reproducible.
        ImageDraw.Draw(image).polygon(
            [(float(x), float(y)) for x, y in cleaned], fill=1
        )
    return np.asarray(image, dtype=np.uint8)