import numpy as np
from functools import lru_cache


def _resolve_model_path(model_type: str) -> str:
    if model_type == "finetune":
        return "./models/ft_dic_cyto3_20251107_111510"
    if model_type == "cyto3":
        return "cyto3"
    raise ValueError("model_type must be 'cyto3' or 'finetune'")


@lru_cache(maxsize=None)
def _get_cellpose_model(model_type: str):
    """
    Lazily load and cache the Cellpose model so we pay the load cost once.
    The import is inside the function to avoid heavy import time at plugin import.
    """
    from cellpose import models  # local import to speed plugin import

    model_path = _resolve_model_path(model_type)
    return models.CellposeModel(pretrained_model=str(model_path), gpu=True)

def _unpack_model_eval_ret(ret):
    """
    Normalize return of model.eval() into:
      masks (list of 2D int arrays), flows, styles, diams
    Handles returns with 2/3/4 elements or ndarray.
    """
    masks = flows = styles = diams = None

    # If ret is not a sequence, try to interpret as ndarray of masks
    if not isinstance(ret, (list, tuple)):
        arr = np.asarray(ret)
        if arr.ndim == 2:
            masks = [arr.astype(np.int32)]
        elif arr.ndim == 3:
            masks = [arr[i].astype(np.int32) for i in range(arr.shape[0])]
        else:
            masks = []
        return masks, None, None, None

    # Sequence case
    n = len(ret)
    if n == 4:
        masks, flows, styles, diams = ret
    elif n == 3:
        masks, flows, styles = ret
        diams = None
    elif n == 2:
        masks, flows = ret
        styles = None
        diams = None
    else:
        masks = []

    # Normalize masks into list[np.ndarray(dtype=int32)]
    if masks is None:
        masks = []
    else:
        if isinstance(masks, np.ndarray):
            if masks.ndim == 3:
                masks = [np.asarray(masks[i]).astype(np.int32) for i in range(masks.shape[0])]
            elif masks.ndim == 2:
                masks = [np.asarray(masks).astype(np.int32)]
            else:
                masks = [np.asarray(m).astype(np.int32) for m in masks]
        elif isinstance(masks, (list, tuple)):
            masks = [np.asarray(m).astype(np.int32) for m in masks]
        else:
            masks = [np.asarray(masks).astype(np.int32)]

    return masks, flows, styles, diams

def segment_dic_image(dic_image, model_type="cyto3"):
    model = _get_cellpose_model(model_type)
    ret = model.eval(dic_image, diameter=150, channels=[0,0])
    masks, flows, styles, diams = _unpack_model_eval_ret(ret)
    return masks
