from collections.abc import Callable
from typing import Literal, Union

import numpy as np
from skimage import filters, measure, morphology


def _extract_largest_object(binary_image):
    """Keep only the largest object in a binary image.

    Parameters
    ----------
    binary_image : np.ndarray
        A binary image.

    Returns
    -------
    np.ndarray
        A binary image containing only the largest object.
    """
    labeled_image = measure.label(binary_image)
    regions = measure.regionprops(labeled_image)
    largest_region = max(regions, key=lambda region: region.area)
    return labeled_image == largest_region.label


def _threshold_image(
    image: np.ndarray,
    method: Literal["triangle", "otsu", "isodata"] = "triangle",
) -> Union[np.ndarray, None]:
    """Threshold an image using the specified method to get a binary mask.

    Parameters
    ----------
    image : np.ndarray
        Image to threshold.
    method : str
        Thresholding method to use. One of 'triangle', 'otsu', and 'isodata'
        (corresponding to methods from the skimage.filters module).
        Defaults to 'triangle'.

    Returns
    -------
    np.ndarray
        A binary mask.
    """

    method_to_func: dict[str, Callable] = {
        "triangle": filters.threshold_triangle,
        "otsu": filters.threshold_otsu,
        "isodata": filters.threshold_isodata,
    }
    if method in method_to_func.keys():
        thresholded = method_to_func[method](image)
        return image > thresholded
    else:
        raise ValueError(f"Unknown thresholding method {method}")


def create_mask(
    image: np.ndarray,
    gauss_sigma: float = 3,
    threshold_method: Literal["triangle", "otsu", "isodata"] = "triangle",
    closing_size: int = 5,
    erode_size: int = 0,
) -> np.ndarray:
    """Threshold image and create a mask for the largest object.

    The mask is generated by applying a Gaussian filter to the image,
    thresholding the smoothed image, keeping only the largest object, and
    applying a binary closing operation to the mask, optionally followed by
    binary erosion.

    Parameters
    ----------
    image : np.ndarray
        A 3D image to generate the mask from.
    gauss_sigma : float
        Standard deviation for Gaussian kernel (in pixels) to smooth image
        before thresholding. Set to 0 to skip smoothing.
    threshold_method : str
        Thresholding method to use. One of 'triangle', 'otsu', and 'isodata'
        (corresponding to methods from the skimage.filters module).
        Defaults to 'triangle'.
    closing_size : int
        Size of the binary closing footprint (in pixels) to apply to the mask.
        Set to 0 to skip closing. This is useful to remove small holes in the
        mask. Larger values will remove larger holes.
    erode_size : int
        Size of the binary erosion footprint (in pixels) to apply to the mask.
        Set to 0 to skip erosion. This is useful to remove small structures
        at the mask boundary. Larger values will remove larger structures.

    Returns
    -------
    mask : np.ndarray
        A binary mask of the largest object in the image.
    """

    # Check input
    if image.ndim != 3:
        raise ValueError("Image must be 3D")

    if gauss_sigma > 0:
        data_smoothed = filters.gaussian(image, sigma=gauss_sigma)
    else:
        data_smoothed = image

    binary = _threshold_image(data_smoothed, method=threshold_method)
    mask = _extract_largest_object(binary)

    if closing_size > 0:
        mask = morphology.binary_closing(
            mask, footprint=np.ones((closing_size,) * image.ndim)
        )

    if erode_size > 0:
        mask = morphology.binary_erosion(
            mask, footprint=np.ones((erode_size,) * image.ndim)
        )

    return mask
