from pathlib import Path
from typing import NamedTuple

import numpy as np
import tifffile


class MicroscopyImage(NamedTuple):
    """현미경 이미지 컨테이너."""
    data: np.ndarray  # shape: (channels, height, width) or (height, width)
    channel_names: list[str]
    pixel_size_um: float  # 픽셀 하나가 물리적으로 몇 마이크로미터인가
    filename: str


def load_microscopy_image(
    path: Path,
    channel_names: list[str] | None = None,
    pixel_size_um: float = 0.325,
) -> MicroscopyImage:
    """TIFF 로딩. OME-TIFF metadata에서 pixel size 자동 추출 가능."""
    data = tifffile.imread(path)
    if data.ndim == 2:
        data = data[np.newaxis, ...]  # (H, W) → (1, H, W)
    if channel_names is None:
        channel_names = [f"ch{i}" for i in range(data.shape[0])]
    return MicroscopyImage(
        data=data,
        channel_names=channel_names,
        pixel_size_um=pixel_size_um,
        filename=path.name,
    )


def normalize_channel(
    channel_data: np.ndarray,
    lower_percentile: float = 1.0,
    upper_percentile: float = 99.5,
) -> np.ndarray:
    """percentile 기반 정규화. min-max보다 outlier에 견고."""
    p_low = np.percentile(channel_data, lower_percentile)
    p_high = np.percentile(channel_data, upper_percentile)
    if p_high - p_low < 1e-6:
        return np.zeros_like(channel_data, dtype=np.float32)
    normalized = np.clip((channel_data - p_low) / (p_high - p_low), 0, 1)
    return normalized.astype(np.float32)

from cellpose import models, io as cp_io
from cellpose.denoise import DenoiseModel


class CellposeSegmenter:
    """Cellpose 3 복원 + 세그멘테이션 통합."""

    def __init__(
        self,
        model_type: str = "cyto3",       # cyto3: 세포체 최신, nuclei: 핵 전용
        restore_type: str = "denoise",   # denoise / deblur / upsample / None
        device: str = "cuda",
    ):
        self.model = models.CellposeModel(
            gpu=(device == "cuda"),
            model_type=model_type,
        )
        self.restore_type = restore_type
        if restore_type:
            self.denoiser = DenoiseModel(
                model_type=f"{restore_type}_{model_type}",
                gpu=(device == "cuda"),
            )
        else:
            self.denoiser = None

    def segment(
        self,
        image: np.ndarray,
        diameter: float | None = None,
        channels: list[int] = [0, 0],
        cellprob_threshold: float = 0.0,
        flow_threshold: float = 0.4,
    ) -> dict:
        """
        image: (H, W) grayscale or (C, H, W) multi-channel.
        diameter: 예상 세포 지름 (px). None이면 자동 추정.
        channels: [세포체 채널, 핵 채널] indices. [0,0] = grayscale.
        Returns: {masks, flows, styles, diams}
        """
        if self.denoiser is not None:
            image = self.denoiser.eval(image, channels=channels)[0]
        masks, flows, styles = self.model.eval(
            image,
            diameter=diameter,
            channels=channels,
            cellprob_threshold=cellprob_threshold,
            flow_threshold=flow_threshold,
        )
        return {
            "masks": masks,        # (H, W) int, 0=배경, 1..N=세포 ID
            "flows": flows,        # gradient flow visualization
            "diameter_used": diameter or self.model.diam_labels,
        }

from segment_anything_2 import SAM2ImagePredictor


class SAM2Refiner:
    """Cellpose 마스크를 SAM 2로 refine."""

    def __init__(self, model_id: str = "facebook/sam2-hiera-base", device: str = "cuda"):
        self.predictor = SAM2ImagePredictor.from_pretrained(model_id)
        self.predictor.model.to(device)

    def refine_mask(
        self,
        image: np.ndarray,
        cellpose_mask: np.ndarray,
        cell_id: int,
    ) -> np.ndarray:
        """특정 세포 ID의 Cellpose 마스크를 SAM 2로 refine."""
        self.predictor.set_image(image)
        # Cellpose 마스크의 bounding box + centroid를 SAM 프롬프트로 사용
        cell_mask = (cellpose_mask == cell_id)
        if not cell_mask.any():
            return cell_mask
        ys, xs = np.where(cell_mask)
        bbox = np.array([xs.min(), ys.min(), xs.max(), ys.max()])
        centroid = np.array([[xs.mean(), ys.mean()]])
        refined_masks, scores, _ = self.predictor.predict(
            point_coords=centroid,
            point_labels=np.array([1]),  # foreground
            box=bbox,
            multimask_output=True,
        )
        # 가장 높은 스코어의 마스크 선택
        best_idx = scores.argmax()
        return refined_masks[best_idx]

    def refine_uncertain_cells(
        self,
        image: np.ndarray,
        cellpose_result: dict,
        uncertainty_threshold: float = 0.5,
    ) -> np.ndarray:
        """Cellpose flow의 확신도가 낮은 세포만 SAM으로 refine."""
        masks = cellpose_result["masks"].copy()
        # flows[2]는 cell probability map
        cell_probs = cellpose_result["flows"][2] if len(cellpose_result["flows"]) > 2 else None
        if cell_probs is None:
            return masks
        for cell_id in np.unique(masks):
            if cell_id == 0:
                continue
            cell_region = (masks == cell_id)
            mean_prob = cell_probs[cell_region].mean()
            if mean_prob < uncertainty_threshold:
                refined = self.refine_mask(image, masks, cell_id)
                masks[cell_region] = 0
                masks[refined] = cell_id
        return masks

from dataclasses import dataclass, asdict
from typing import Iterable

from skimage.measure import regionprops, regionprops_table
import pandas as pd


@dataclass
class CellFeatures:
    """세포 하나의 형태·강도 특성."""
    cell_id: int
    area_px: int
    area_um2: float
    perimeter_px: float
    circularity: float           # 4π·area / perimeter² (원형=1)
    eccentricity: float          # 편심도 (원=0, 선=1)
    solidity: float              # area / convex_hull_area
    centroid_y: float
    centroid_x: float
    mean_intensity_per_channel: dict[str, float]
    max_intensity_per_channel: dict[str, float]


def extract_features(
    mask: np.ndarray,
    intensity_channels: dict[str, np.ndarray],
    pixel_size_um: float,
) -> list[CellFeatures]:
    """
    mask: (H, W) int, 0=배경, 1..N=세포 ID.
    intensity_channels: {channel_name: (H, W) array}.
    """
    features = []
    for prop in regionprops(mask):
        perimeter = prop.perimeter if prop.perimeter > 0 else 1e-6
        circularity = (4 * np.pi * prop.area) / (perimeter ** 2)
        mean_intensity = {
            name: float(ch[prop.coords[:, 0], prop.coords[:, 1]].mean())
            for name, ch in intensity_channels.items()
        }
        max_intensity = {
            name: float(ch[prop.coords[:, 0], prop.coords[:, 1]].max())
            for name, ch in intensity_channels.items()
        }
        features.append(CellFeatures(
            cell_id=int(prop.label),
            area_px=int(prop.area),
            area_um2=float(prop.area * (pixel_size_um ** 2)),
            perimeter_px=float(prop.perimeter),
            circularity=float(circularity),
            eccentricity=float(prop.eccentricity),
            solidity=float(prop.solidity),
            centroid_y=float(prop.centroid[0]),
            centroid_x=float(prop.centroid[1]),
            mean_intensity_per_channel=mean_intensity,
            max_intensity_per_channel=max_intensity,
        ))
    return features


def features_to_dataframe(features: Iterable[CellFeatures]) -> pd.DataFrame:
    """다운스트림 분석용 표 형태."""
    rows = []
    for f in features:
        row = asdict(f)
        for ch_name, val in row.pop("mean_intensity_per_channel").items():
            row[f"mean_{ch_name}"] = val
        for ch_name, val in row.pop("max_intensity_per_channel").items():
            row[f"max_{ch_name}"] = val
        rows.append(row)
    return pd.DataFrame(rows)

def visualize_with_napari(
    image: np.ndarray,
    masks: np.ndarray,
    channel_names: list[str],
) -> None:
    """napari로 이미지 + 마스크 오버레이 시각화.
    
    실행 시 GUI 창이 열림. 원격 서버에서는 X11 forward 또는 
    napari 웹 뷰어(napari-remote) 사용.
    """
    import napari
    
    viewer = napari.Viewer()
    for i, ch_name in enumerate(channel_names):
        viewer.add_image(
            image[i] if image.ndim == 3 else image,
            name=ch_name,
            colormap=["green", "red", "blue"][i % 3],
            blending="additive",
        )
    viewer.add_labels(masks, name="cell masks", opacity=0.5)
    napari.run()

def full_pipeline(
    image_path: Path,
    channel_names: list[str],
    pixel_size_um: float,
    output_csv: Path,
    use_sam_refinement: bool = False,
    device: str = "cuda",
) -> pd.DataFrame:
    """이미지 1장 → 세포 특성 DataFrame."""
    img = load_microscopy_image(image_path, channel_names, pixel_size_um)
    
    # 세포체 채널 정규화 (예: FITC)
    cell_channel_idx = channel_names.index("FITC") if "FITC" in channel_names else 0
    normalized = normalize_channel(img.data[cell_channel_idx])
    
    segmenter = CellposeSegmenter(model_type="cyto3", restore_type="denoise", device=device)
    result = segmenter.segment(normalized)
    masks = result["masks"]
    
    if use_sam_refinement:
        refiner = SAM2Refiner(device=device)
        masks = refiner.refine_uncertain_cells(normalized, result)
    
    intensity_channels = {
        name: normalize_channel(img.data[i])
        for i, name in enumerate(channel_names)
    }
    features = extract_features(masks, intensity_channels, pixel_size_um)
    df = features_to_dataframe(features)
    df["source_image"] = img.filename
    df.to_csv(output_csv, index=False)
    return df