from itertools import product
from typing import Callable, Union

import numpy as np
import pandas as pd
import xarray as xr

from starfish.core.imagestack.imagestack import ImageStack
from starfish.core.types import SpotAttributes, SpotFindingResults
from starfish.types import Axes, Features, Number


def measure_intensities_at_spot_locations_in_image(
        image: Union[np.ndarray, xr.DataArray],
        spots: SpotAttributes,
        measurement_function: Callable[[Union[np.ndarray, xr.DataArray]], Number],
        radius_is_gyration: bool = False,
) -> pd.Series:
    """measure the intensity of each spot in spots in the corresponding image

    Parameters
    ----------
    image : Union[np.ndarray, xr.DataArray],
        3-d volume in which to measure intensities
    spots : pd.DataFrame
        SpotAttributes table containing coordinates and radii of spots
    measurement_function : Callable[[Sequence], Number])
        Function to apply over the spot volumes to identify the intensity (e.g. max, mean, ...)
    radius_is_gyration : bool
        if True, indicates that the radius corresponds to radius of gyration, which is a
        function of spot intensity, but typically is a smaller unit than the sigma generated by
        blob_log. In this case, the spot's bounding box is rounded up instead of down when
        measuring intensity. (default False)

    Returns
    -------
    pd.Series :
        Intensities for each spot in SpotAttributes

    """

    def fn(row: pd.Series) -> Number:
        d = image[row['z_min']:row['z_max'], row['y_min']:row['y_max'], row['x_min']:row['x_max']]
        return measurement_function(d)

    if radius_is_gyration:
        radius = np.ceil(spots.data[Features.SPOT_RADIUS]).astype(int) + 1  # round up
    else:
        radius = spots.data[Features.SPOT_RADIUS].astype(int)  # truncate down to nearest int
    for v, max_size in zip(['z', 'y', 'x'], image.shape):
        # numpy does exclusive max indexing, so need to subtract 1 from min to get centered box
        spots.data[f'{v}_min'] = np.clip(spots.data[v] - (radius - 1), 0, None)
        spots.data[f'{v}_max'] = np.clip(spots.data[v] + radius, None, max_size)
    return spots.data[['z_min', 'z_max', 'y_min', 'y_max', 'x_min', 'x_max']].astype(int).apply(
        fn,
        axis=1
    )


def measure_intensities_at_spot_locations_across_imagestack(
        data_image: ImageStack,
        reference_spots: SpotAttributes,
        measurement_function: Callable[[Union[np.ndarray, xr.DataArray]], Number],
        radius_is_gyration: bool = False) -> SpotFindingResults:
    """given spots found from a reference image, find those spots across a data_image

    Parameters
    ----------
    data_image : ImageStack
        ImageStack containing multiple volumes for which spots' intensities must be calculated
    reference_spots : SpotAttributes
        Spots found in a reference image
    measurement_function : Callable[[Sequence], Number])
        Function to apply over the spot volumes to identify the intensity (e.g. max, mean, ...)
    radius_is_gyration : bool
        if True, indicates that the radius corresponds to radius of gyration, which is
        a function of spot intensity, but typically is a smaller unit than the sigma generated
        by blob_log. In this case, the spot's bounding box is rounded up instead of down when
        measuring intensity. (default False)

    Returns
    -------
    SpotFindingResults :
        A Dict of tile indices and their corresponding measured SpotAttributes

    """

    ch_labels = data_image.axis_labels(Axes.CH)
    round_labels = data_image.axis_labels(Axes.ROUND)

    spot_results = SpotFindingResults(imagestack_coords=data_image.xarray.coords,
                                      log=data_image.log)
    # measure spots in each tile
    indices = product(ch_labels, round_labels)
    for c, r in indices:
        tile_indices = {Axes.ROUND: r, Axes.CH: c}
        if reference_spots.data.empty:
            # if no spots found don't measure
            spot_results[tile_indices] = reference_spots
        else:
            image, _ = data_image.get_slice({Axes.CH: c, Axes.ROUND: r})
            blob_intensities: pd.Series = measure_intensities_at_spot_locations_in_image(
                image,
                reference_spots,
                measurement_function,
                radius_is_gyration=radius_is_gyration
            )
            # copy reference spot positions and attributes
            tile_spots = SpotAttributes(reference_spots.data.copy())
            # fill in intensities
            tile_spots.data[Features.INTENSITY] = blob_intensities
            spot_results[tile_indices] = tile_spots
    return spot_results
