from itertools import product
from typing import Callable

import numpy as np
import pandas as pd

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


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

    Parameters
    ----------
    image : np.ndarray,
        3-d volume in which to measure intensities
    spots : pd.DataFrame
        SpotAttributes table containing coordinates and radii of spots
    measurement_function : Callable[[np.ndarray], 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: PerImageSliceSpotResults,
        measurement_function: Callable[[np.ndarray], 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 : PerImageSliceSpotResults
        Spots found in a reference image
    measurement_function : Callable[[np.ndarray], 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.spot_attrs.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.spot_attrs,
                measurement_function,
                radius_is_gyration=radius_is_gyration
            )
            # copy reference spot positions and attributes
            tile_spots = SpotAttributes(reference_spots.spot_attrs.data.copy())
            # fill in intensities
            tile_spots.data[Features.INTENSITY] = blob_intensities
            spot_results[tile_indices] = PerImageSliceSpotResults(
                spot_attrs=tile_spots, extras=None)
    return spot_results
