from __future__ import annotations

from dataclasses import dataclass
from typing import Callable

import cv2
import numpy as np
from skimage.filters import frangi, threshold_otsu
from skimage.morphology import remove_small_objects, skeletonize


METHOD_DESCRIPTIONS = {
    "hessian_ridge": {
        "label": "Hessian / Frangi 细线增强",
        "description": "多尺度 Hessian 管状结构响应，适合低对比细导丝候选提取。",
        "uses_temporal": False,
    },
    "edge_morphology": {
        "label": "边缘 + 形态学",
        "description": "CLAHE、黑帽增强、Canny 边缘与线性形态学连接。",
        "uses_temporal": False,
    },
    "temporal_difference": {
        "label": "视频时序差分",
        "description": "利用相邻帧运动候选抑制静态骨骼和背景结构。",
        "uses_temporal": True,
    },
    "fusion": {
        "label": "融合模式",
        "description": "融合 Hessian、边缘形态学和时序差分，作为默认稳健输出。",
        "uses_temporal": True,
    },
    "compare": {
        "label": "多方法对比",
        "description": "对同一帧同时运行多种方法，便于调参和方案比较。",
        "uses_temporal": True,
    },
}


@dataclass(frozen=True)
class SegmentationOutput:
    method: str
    mask: np.ndarray
    overlay: np.ndarray
    metrics: dict[str, float | int]


def normalize01(image: np.ndarray) -> np.ndarray:
    image = image.astype(np.float32)
    low = float(np.percentile(image, 1))
    high = float(np.percentile(image, 99))
    if high <= low:
        return np.zeros_like(image, dtype=np.float32)
    return np.clip((image - low) / (high - low), 0.0, 1.0)


def to_gray(frame: np.ndarray) -> np.ndarray:
    if frame.ndim == 2:
        return frame
    return cv2.cvtColor(frame, cv2.COLOR_BGR2GRAY)


def clahe_gray(frame: np.ndarray) -> np.ndarray:
    gray = to_gray(frame)
    clahe = cv2.createCLAHE(clipLimit=2.2, tileGridSize=(8, 8))
    return clahe.apply(gray)


def _adaptive_cutoff(response: np.ndarray, sensitivity: float) -> float:
    response = normalize01(response)
    nonzero = response[response > 0]
    if nonzero.size < 16:
        return 1.0
    sensitivity = float(np.clip(sensitivity, 0.05, 0.95))
    percentile = 99.2 - sensitivity * 22.0
    percentile_cut = float(np.percentile(nonzero, percentile))
    try:
        otsu_cut = float(threshold_otsu(nonzero))
    except ValueError:
        otsu_cut = percentile_cut
    return max(min(percentile_cut, 0.98), otsu_cut * 0.72)


def clean_mask(mask: np.ndarray, min_area: int = 12) -> np.ndarray:
    binary = mask.astype(bool)
    binary = remove_small_objects(binary, max_size=max(1, int(min_area)))
    cleaned = binary.astype(np.uint8) * 255
    kernel = cv2.getStructuringElement(cv2.MORPH_ELLIPSE, (3, 3))
    cleaned = cv2.morphologyEx(cleaned, cv2.MORPH_CLOSE, kernel, iterations=1)
    return cleaned


def hessian_ridge_mask(frame: np.ndarray, sensitivity: float = 0.56) -> np.ndarray:
    enhanced = clahe_gray(frame)
    inverted = 255 - enhanced
    normalized = normalize01(inverted)
    response = frangi(
        normalized,
        sigmas=(0.7, 1.1, 1.7, 2.3),
        alpha=0.55,
        beta=0.55,
        gamma=12,
        black_ridges=False,
    )
    response = normalize01(response)
    cutoff = _adaptive_cutoff(response, sensitivity)
    mask = response >= cutoff
    return clean_mask(mask, min_area=10)


def edge_morphology_mask(frame: np.ndarray, sensitivity: float = 0.56) -> np.ndarray:
    enhanced = clahe_gray(frame)
    blur = cv2.GaussianBlur(enhanced, (3, 3), 0)
    kernel = cv2.getStructuringElement(cv2.MORPH_RECT, (15, 15))
    blackhat = cv2.morphologyEx(blur, cv2.MORPH_BLACKHAT, kernel)
    dark_line = normalize01(blackhat)
    cutoff = _adaptive_cutoff(dark_line, min(0.95, sensitivity + 0.1))
    candidate = (dark_line >= cutoff).astype(np.uint8) * 255
    low = int(20 + (1.0 - sensitivity) * 65)
    high = int(70 + (1.0 - sensitivity) * 120)
    edges = cv2.Canny(blur, low, high)
    candidate = cv2.dilate(candidate, cv2.getStructuringElement(cv2.MORPH_ELLIPSE, (3, 3)), iterations=1)
    edges = cv2.bitwise_and(edges, candidate)
    line_h = cv2.getStructuringElement(cv2.MORPH_RECT, (5, 1))
    line_v = cv2.getStructuringElement(cv2.MORPH_RECT, (1, 5))
    connected = cv2.morphologyEx(edges, cv2.MORPH_CLOSE, line_h, iterations=1)
    connected = cv2.morphologyEx(connected, cv2.MORPH_CLOSE, line_v, iterations=1)
    connected = cv2.bitwise_or(connected, cv2.erode(candidate, cv2.getStructuringElement(cv2.MORPH_ELLIPSE, (2, 2))))
    return clean_mask(connected, min_area=8)


def temporal_difference_mask(
    frame: np.ndarray,
    previous_frame: np.ndarray | None,
    sensitivity: float = 0.56,
) -> np.ndarray:
    ridge = hessian_ridge_mask(frame, sensitivity=sensitivity)
    if previous_frame is None:
        return ridge
    current = cv2.GaussianBlur(clahe_gray(frame), (5, 5), 0)
    previous = cv2.GaussianBlur(clahe_gray(previous_frame), (5, 5), 0)
    diff = cv2.absdiff(current, previous)
    diff = normalize01(diff)
    cutoff = _adaptive_cutoff(diff, min(0.92, sensitivity + 0.12))
    moving = (diff >= cutoff).astype(np.uint8) * 255
    moving = cv2.dilate(moving, cv2.getStructuringElement(cv2.MORPH_ELLIPSE, (3, 3)), iterations=1)
    blended = cv2.bitwise_or(cv2.bitwise_and(ridge, moving), cv2.bitwise_and(ridge, cv2.dilate(moving, None)))
    if int(np.count_nonzero(blended)) < 8:
        blended = cv2.bitwise_or(ridge, moving)
    return clean_mask(blended, min_area=8)


def fusion_mask(
    frame: np.ndarray,
    previous_frame: np.ndarray | None = None,
    sensitivity: float = 0.56,
) -> np.ndarray:
    ridge = hessian_ridge_mask(frame, sensitivity=sensitivity)
    edge = edge_morphology_mask(frame, sensitivity=sensitivity)
    temporal = temporal_difference_mask(frame, previous_frame, sensitivity=sensitivity)
    votes = (
        (ridge > 0).astype(np.uint8)
        + (edge > 0).astype(np.uint8)
        + (temporal > 0).astype(np.uint8)
    )
    fused = votes >= 2
    if int(np.count_nonzero(fused)) < 8:
        fused = votes >= 1
    return clean_mask(fused, min_area=10)


def overlay_mask(frame: np.ndarray, mask: np.ndarray, color: tuple[int, int, int] = (0, 220, 255)) -> np.ndarray:
    if frame.ndim == 2:
        base = cv2.cvtColor(frame, cv2.COLOR_GRAY2BGR)
    else:
        base = frame.copy()
    color_layer = np.zeros_like(base)
    color_layer[mask > 0] = color
    return cv2.addWeighted(base, 0.78, color_layer, 0.72, 0)


def mask_metrics(mask: np.ndarray) -> dict[str, float | int]:
    binary = mask > 0
    coverage = float(np.count_nonzero(binary) / binary.size) if binary.size else 0.0
    skeleton = skeletonize(binary)
    component_count, _ = cv2.connectedComponents(binary.astype(np.uint8))
    return {
        "coverage": round(coverage, 6),
        "mask_pixels": int(np.count_nonzero(binary)),
        "skeleton_length": int(np.count_nonzero(skeleton)),
        "components": max(0, int(component_count) - 1),
    }


def segment_frame(
    frame: np.ndarray,
    method: str = "fusion",
    previous_frame: np.ndarray | None = None,
    sensitivity: float = 0.56,
) -> SegmentationOutput:
    method_map: dict[str, Callable[..., np.ndarray]] = {
        "hessian_ridge": hessian_ridge_mask,
        "edge_morphology": edge_morphology_mask,
        "temporal_difference": temporal_difference_mask,
        "fusion": fusion_mask,
    }
    if method not in method_map:
        raise ValueError(f"Unknown segmentation method: {method}")
    if method in {"temporal_difference", "fusion"}:
        mask = method_map[method](frame, previous_frame, sensitivity)
    else:
        mask = method_map[method](frame, sensitivity)
    return SegmentationOutput(
        method=method,
        mask=mask,
        overlay=overlay_mask(frame, mask),
        metrics=mask_metrics(mask),
    )


def compare_frame(
    frame: np.ndarray,
    previous_frame: np.ndarray | None = None,
    sensitivity: float = 0.56,
) -> list[SegmentationOutput]:
    return [
        segment_frame(frame, "hessian_ridge", previous_frame, sensitivity),
        segment_frame(frame, "edge_morphology", previous_frame, sensitivity),
        segment_frame(frame, "temporal_difference", previous_frame, sensitivity),
        segment_frame(frame, "fusion", previous_frame, sensitivity),
    ]
