import numpy as np
from skimage import exposure

from .params import AnalysisParameters


class CellposeSegmenter:
    def __init__(self, params: AnalysisParameters):
        self.params = params
        self.model = self._load_model()

    def _load_model(self):
        from cellpose import models
        return models.CellposeModel(
            pretrained_model=self.params.pretrained_model,
            gpu=self.params.gpu,
        )

    def segment(self, image: np.ndarray) -> np.ndarray:
        if image.ndim == 2:
            img = np.stack([image, image, image], axis=0)
        elif image.ndim == 3:
            if image.shape[-1] in (3, 4):
                img = np.transpose(image[:, :, :3], (2, 0, 1))
            elif image.shape[0] in (3, 4):
                img = image[:3]
            else:
                raise ValueError(f"Unexpected image shape: {image.shape}")
        else:
            raise ValueError(f"Unexpected image ndim: {image.ndim}")

        # Contrast stretch (Fiji-style auto B/C)
        for c in range(img.shape[0]):
            p2, p98 = np.percentile(img[c], (2, 98))
            img[c] = exposure.rescale_intensity(img[c], in_range=(p2, p98))

        # Bubbles are dark on light background — invert for Cellpose
        img_inv = 255 - img if img.max() > 1 else 1 - img

        masks, _, _ = self.model.eval(
            img_inv,
            diameter=self.params.diameter,
            flow_threshold=self.params.flow_threshold,
            cellprob_threshold=self.params.cellprob_threshold,
        )
        return masks