from skimage.measure._regionprops import (
    RegionProperties,
    _cached,
    _infer_number_of_required_args,
)
from typing import Optional, List, Dict, Any, Tuple
import numpy as np
import inspect
import json
import os
from scipy.ndimage import find_objects
from celldetective.log_manager import get_logger

logger = get_logger(__name__)

abs_path = os.sep.join([os.path.split(os.path.dirname(os.path.realpath(__file__)))[0]])

with open(os.sep.join([abs_path, "regionprops", "props.json"])) as f:
    PROPS = json.load(f)

COL_DTYPES = {
    "area": float,
    "area_bbox": float,
    "area_convex": float,
    "area_filled": float,
    "axis_major_length": float,
    "axis_minor_length": float,
    "bbox": int,
    "centroid": float,
    "centroid_local": float,
    "centroid_weighted": float,
    "centroid_weighted_local": float,
    "coords": object,
    "coords_scaled": object,
    "eccentricity": float,
    "equivalent_diameter_area": float,
    "euler_number": int,
    "extent": float,
    "feret_diameter_max": float,
    "image": object,
    "image_convex": object,
    "image_filled": object,
    "image_intensity": object,
    "inertia_tensor": float,
    "inertia_tensor_eigvals": float,
    "intensity_max": float,
    "intensity_mean": float,
    "intensity_min": float,
    "intensity_std": float,
    "label": int,
    "moments": float,
    "moments_central": float,
    "moments_hu": float,
    "moments_normalized": float,
    "moments_weighted": float,
    "moments_weighted_central": float,
    "moments_weighted_hu": float,
    "moments_weighted_normalized": float,
    "num_pixels": int,
    "orientation": float,
    "perimeter": float,
    "perimeter_crofton": float,
    "slice": object,
    "solidity": float,
}

OBJECT_COLUMNS = [col for col, dtype in COL_DTYPES.items() if dtype == object]
PROP_VALS = set(PROPS.values())

_require_intensity_image = (
    "image_intensity",
    "intensity_max",
    "intensity_mean",
    "intensity_median",
    "intensity_min",
    "intensity_std",
    "moments_weighted",
    "moments_weighted_central",
    "centroid_weighted",
    "centroid_weighted_local",
    "moments_weighted_hu",
    "moments_weighted_normalized",
)


class CustomRegionProps(RegionProperties):
    """
    From https://github.com/scikit-image/scikit-image/blob/main/skimage/measure/_regionprops.py with a modification to not mask the intensity image itself before measurements
    """

    def __init__(
        self, channel_names: Optional[List[str]], *args: Any, **kwargs: Any
    ) -> None:
        """
        Initialize CustomRegionProps.

        Parameters
        ----------
        channel_names : list
            List of channel names.
        *args
            Variable length argument list.
        **kwargs
            Arbitrary keyword arguments.
        """

        self.channel_names = channel_names
        if isinstance(self.channel_names, np.ndarray):
            self.channel_names = list(self.channel_names)
        super().__init__(*args, **kwargs)

    def __getattr__(self, attr: str) -> Any:
        """
        Get attribute.

        Parameters
        ----------
        attr : str
            Attribute name.

        Returns
        -------
        object
            Attribute value.
        """

        if self.channel_names is not None and self._multichannel:
            assert (
                len(self.channel_names) == self._intensity_image.shape[-1]
            ), "Mismatch between provided channel names and the number of channels in the image..."

        if attr == "__setstate__":
            return self.__getattribute__(attr)

        if self._intensity_image is None and attr in _require_intensity_image:
            raise AttributeError(
                f"Attribute '{attr}' unavailable when `intensity_image` "
                f"has not been specified."
            )
        if attr in self._extra_properties:
            func = self._extra_properties[attr]
            n_args = _infer_number_of_required_args(func)
            # determine whether func requires intensity image
            if n_args == 2:
                if self._intensity_image is not None:
                    if self._multichannel:
                        arg_dict = dict(inspect.signature(func).parameters)
                        if (
                            self.channel_names is not None
                            and "target_channel" in arg_dict
                        ):
                            multichannel_list = [
                                np.nan for i in range(self.image_intensity.shape[-1])
                            ]
                            len_output = 1
                            default_channel = arg_dict["target_channel"]._default

                            if default_channel in self.channel_names:

                                idx = self.channel_names.index(default_channel)
                                res = func(self.image, self.image_intensity[..., idx])
                                if isinstance(res, tuple):
                                    len_output = len(res)
                                else:
                                    len_output = 1

                                if len_output > 1:
                                    multichannel_list = [
                                        [np.nan] * len_output
                                        for c in range(len(self.channel_names))
                                    ]
                                    multichannel_list[idx] = res
                                else:
                                    multichannel_list = [
                                        np.nan for c in range(len(self.channel_names))
                                    ]
                                    multichannel_list[idx] = res

                            else:
                                print(
                                    f"Warning... Channel required by custom measurement ({default_channel}) could not be found in your data..."
                                )

                            return np.stack(multichannel_list, axis=-1)
                        else:
                            multichannel_list = [
                                func(self.image, self.image_intensity[..., i])
                                for i in range(self.image_intensity.shape[-1])
                            ]
                            return np.stack(multichannel_list, axis=-1)
                    else:
                        return func(self.image, self.image_intensity)
                else:
                    raise AttributeError(
                        f"intensity image required to calculate {attr}"
                    )
            elif n_args == 1:
                return func(self.image)
            else:
                raise AttributeError(
                    f"Custom regionprop function's number of arguments must "
                    f"be 1 or 2, but {attr} takes {n_args} arguments."
                )
        elif attr in PROPS and attr.lower() == attr:
            if (
                self._intensity_image is None
                and PROPS[attr] in _require_intensity_image
            ):
                raise AttributeError(
                    f"Attribute '{attr}' unavailable when `intensity_image` "
                    f"has not been specified."
                )
            # retrieve deprecated property (excluding old CamelCase ones)
            return getattr(self, PROPS[attr])
        else:
            raise AttributeError(f"'{type(self)}' object has no attribute '{attr}'")

    @property
    @_cached
    def image_intensity(self):
        """
        Get intensity image.

        Returns
        -------
        ndarray
            Intensity image.
        """
        if self._intensity_image is None:
            raise AttributeError("No intensity image specified.")
        image = (
            self.image
            if not self._multichannel
            else np.expand_dims(self.image, self._ndim)
        )
        return self._intensity_image[self.slice]


def regionprops(
    label_image: np.ndarray,
    intensity_image: Optional[np.ndarray] = None,
    cache: bool = True,
    channel_names: Optional[List[str]] = None,
    *,
    extra_properties: Optional[List[Any]] = None,
    spacing: Optional[Tuple[float, ...]] = None,
    offset: Optional[np.ndarray] = None,
) -> List[CustomRegionProps]:
    """
    Measure properties of labeled image regions.

    Parameters
    ----------
    label_image : (N, M, P) ndarray
        Labeled input image. Integer array where non-zero values represent distinct regions.
    intensity_image : (N, M, P, C) ndarray, optional
        Intensity image with same size as labeled image. Default is None.
    cache : bool, optional
        Whether to cache calculated properties. Default is True.
    channel_names : list, optional
        Names of channels in intensity image. Default is None.
    extra_properties : list, optional
        Extra properties to measure. Default is None.
    spacing : tuple, optional
        Pixel spacing. Default is None.
    offset : tuple, optional
        Offset of the labels. Default is None.

    Returns
    -------
    list
        List of RegionProperties.

    Notes
    -----
    From https://github.com/scikit-image/scikit-image/blob/main/skimage/measure/_regionprops.py with a modification to use CustomRegionProps
    """

    if label_image.ndim not in (2, 3):
        raise TypeError("Only 2-D and 3-D images supported.")

    if not np.issubdtype(label_image.dtype, np.integer):
        if np.issubdtype(label_image.dtype, bool):
            raise TypeError(
                "Non-integer image types are ambiguous: "
                "use skimage.measure.label to label the connected "
                "components of label_image, "
                "or label_image.astype(np.uint8) to interpret "
                "the True values as a single label."
            )
        else:
            raise TypeError("Non-integer label_image types are ambiguous")

    if offset is None:
        offset_arr = np.zeros((label_image.ndim,), dtype=int)
    else:
        offset_arr = np.asarray(offset)
        if offset_arr.ndim != 1 or offset_arr.size != label_image.ndim:
            raise ValueError(
                "Offset should be an array-like of integers "
                "of shape (label_image.ndim,); "
                f"{offset} was provided."
            )

    regions = []

    objects = find_objects(label_image)
    for i, sl in enumerate(objects):
        if sl is None:
            continue

        label = i + 1

        props = CustomRegionProps(
            channel_names,
            sl,
            label,
            label_image,
            intensity_image,
            cache,
            spacing=spacing,
            extra_properties=extra_properties,
            offset=offset_arr,
        )
        regions.append(props)

    return regions


def _props_to_dict(
    regions: List[CustomRegionProps],
    properties: Tuple[str, ...] = ("label", "bbox"),
    separator: str = "-",
) -> Dict[str, np.ndarray]:
    """
    Convert region properties to dictionary.

    Parameters
    ----------
    regions : list
        List of RegionProperties.
    properties : tuple, optional
        Properties to include. Default is ("label", "bbox").
    separator : str, optional
        Separator for multi-column properties. Default is "-".

    Returns
    -------
    dict
        Dictionary of properties.
    """

    out = {}
    n = len(regions)
    for prop in properties:
        r = regions[0]
        # Copy the original property name so the output will have the
        # user-provided property name in the case of deprecated names.
        orig_prop = prop
        # determine the current property name for any deprecated property.
        prop = PROPS.get(prop, prop)
        rp = getattr(r, prop)
        if prop in COL_DTYPES:
            dtype = COL_DTYPES[prop]
        else:
            func = r._extra_properties[prop]
            # dtype = _infer_regionprop_dtype(
            # 	func,
            # 	intensity=r._intensity_image is not None,
            # 	ndim=r.image.ndim,
            # )
            dtype = np.float64

        # scalars and objects are dedicated one column per prop
        # array properties are raveled into multiple columns
        # for more info, refer to notes 1
        if np.isscalar(rp) or prop in OBJECT_COLUMNS or dtype is np.object_:
            column_buffer = np.empty(n, dtype=dtype)
            for i in range(n):
                column_buffer[i] = regions[i][prop]
            out[orig_prop] = np.copy(column_buffer)
        else:
            # precompute property column names and locations
            modified_props = []
            locs = []
            for ind in np.ndindex(np.shape(rp)):
                modified_props.append(separator.join(map(str, (orig_prop,) + ind)))
                locs.append(ind if len(ind) > 1 else ind[0])

            # fill temporary column data_array
            n_columns = len(locs)
            column_data = np.empty((n, n_columns), dtype=dtype)
            for k in range(n):
                # we coerce to a numpy array to ensure structures like
                # tuple-of-arrays expand correctly into columns
                rp = np.asarray(regions[k][prop])
                for i, loc in enumerate(locs):
                    column_data[k, i] = rp[loc]

            # add the columns to the output dictionary
            for i, modified_prop in enumerate(modified_props):
                out[modified_prop] = column_data[:, i]
    return out


def regionprops_table(
    label_image: np.ndarray,
    intensity_image: Optional[np.ndarray] = None,
    properties: Tuple[str, ...] = ("label", "bbox"),
    *,
    cache: bool = True,
    separator: str = "-",
    extra_properties: Optional[List[Any]] = None,
    spacing: Optional[Tuple[float, ...]] = None,
    channel_names: Optional[List[str]] = None,
) -> Dict[str, np.ndarray]:
    """
    Compute region properties and return as a pandas-compatible table.

    Parameters
    ----------
    label_image : (N, M) ndarray
        Labeled input image.
    intensity_image : (N, M, C) ndarray, optional
        Intensity image. Default is None.
    properties : tuple, optional
        Properties to include. Default is ("label", "bbox").
    cache : bool, optional
        Whether to cache calculated properties. Default is True.
    separator : str, optional
        Separator for multi-column properties. Default is "-".
    extra_properties : list, optional
        Extra properties to measure. Default is None.
    spacing : tuple, optional
        Pixel spacing. Default is None.
    channel_names : list, optional
        Names of channels in intensity image. Default is None.

    Returns
    -------
    dict
        Dictionary of properties.

    Notes
    -----
    From https://github.com/scikit-image/scikit-image/blob/main/skimage/measure/_regionprops.py
    """
    regions = regionprops(
        label_image,
        intensity_image=intensity_image,
        cache=cache,
        extra_properties=extra_properties,
        spacing=spacing,
        channel_names=channel_names,
    )
    if extra_properties is not None:
        properties = list(properties) + [prop.__name__ for prop in extra_properties]
    if len(regions) == 0:
        ndim = label_image.ndim
        label_image = np.zeros((3,) * ndim, dtype=int)
        label_image[(1,) * ndim] = 1
        if intensity_image is not None:
            intensity_image = np.zeros(
                label_image.shape + intensity_image.shape[ndim:],
                dtype=intensity_image.dtype,
            )
        regions = regionprops(
            label_image,
            intensity_image=intensity_image,
            cache=cache,
            extra_properties=extra_properties,
            spacing=spacing,
            channel_names=channel_names,
        )
        out_d = _props_to_dict(regions, properties=properties, separator=separator)
        return {k: v[:0] for k, v in out_d.items()}

    good_props = []
    for prop in properties:
        try:
            nan_test = [np.isnan(getattr(r, prop)) for r in regions]
            if not np.all(nan_test):
                good_props.append(prop)
        except AttributeError:
            logger.warning(f"Could not measure {prop}... Skip...")

    return _props_to_dict(regions, properties=good_props, separator=separator)
