# Copyright https://en.wikipedia.org/wiki/List_of_baryons
# MIT License
#
# Copyright (c) 2018 eng-tools
#
# Permission is hereby granted, free of charge, to any person obtaining a copy
# of this software and associated documentation files (the "Software"), to deal
# in the Software without restriction, including without limitation the rights
# to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
# copies of the Software, and to permit persons to whom the Software is
# furnished to do so, subject to the following conditions:
#
# The above copyright notice and this permission notice shall be included in all
# copies or substantial portions of the Software.
#
# THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
# IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
# FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
# AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
# LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
# OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
# SOFTWARE.
# ==============================================================================

"""Grid score calculations.
"""
from __future__ import absolute_import, division, print_function

import copy

import matplotlib.pyplot as plt
import numpy as np
import skimage
from scipy import ndimage, stats
from skimage.segmentation import watershed

from neuralplayground.plotting.plot_utils import make_plot_rate_map


class GridScorer(object):
    """Class for scoring ratemaps given trajectories."""

    def __init__(self, nbins):
        """Scoring ratemaps given trajectories.

        Args:
          nbins: Number of bins per dimension in the ratemap.
          coords_range: Environment coordinates range.
          mask_parameters: parameters for the masks that analyze the angular
            autocorrelation of the 2D autocorrelation.
          min_max: Correction.
        """
        self._nbins = nbins

    def autoCorr2D(self, A, nodwell, tol=1e-10):
        """
        Performs a spatial autocorrelation on the array A
        Parameters
        ----------
        A : array_like
            Either 2 or 3D. In the former it is simply the binned up ratemap
            where the two dimensions correspond to x and y.
            If 3D then the first two dimensions are x
            and y and the third (last dimension) is 'stack' of ratemaps
        nodwell : array_like
            A boolean array corresponding the bins in the ratemap that
            weren't visited. See Notes below.
        tol : float, optional
            Values below this are set to zero to deal with v small values
            thrown up by the fft. Default 1e-10
        Returns
        -------
        sac : array_like
            The spatial autocorrelation in the relevant dimensionality
        Notes
        -----
        The nodwell input can usually be generated by:
        >>> nodwell = ~np.isfinite(A)
        """

        assert np.ndim(A) == 2
        m, n = np.shape(A)
        o = 1
        x = np.reshape(A, (m, n, o))
        nodwell = np.reshape(nodwell, (m, n, o))
        x[nodwell] = 0
        # [Step 1] Obtain FFTs of x, the sum of squares and bins visited
        Fx = np.fft.fft(np.fft.fft(x, 2 * m - 1, axis=0), 2 * n - 1, axis=1)
        FsumOfSquares_x = np.fft.fft(np.fft.fft(np.power(x, 2), 2 * m - 1, axis=0), 2 * n - 1, axis=1)
        Fn = np.fft.fft(
            np.fft.fft(np.invert(nodwell).astype(int), 2 * m - 1, axis=0),
            2 * n - 1,
            axis=1,
        )
        # [Step 2] Multiply the relevant transforms and invert to obtain the
        # equivalent convolutions
        rawCorr = np.fft.fftshift(
            np.real(np.fft.ifft(np.fft.ifft(Fx * np.conj(Fx), axis=1), axis=0)),
            axes=(0, 1),
        )
        sums_x = np.fft.fftshift(
            np.real(np.fft.ifft(np.fft.ifft(np.conj(Fx) * Fn, axis=1), axis=0)),
            axes=(0, 1),
        )
        sumOfSquares_x = np.fft.fftshift(
            np.real(np.fft.ifft(np.fft.ifft(Fn * np.conj(FsumOfSquares_x), axis=1), axis=0)),
            axes=(0, 1),
        )
        N = np.fft.fftshift(
            np.real(np.fft.ifft(np.fft.ifft(Fn * np.conj(Fn), axis=1), axis=0)),
            axes=(0, 1),
        )
        # [Step 3] Account for rounding errors.
        rawCorr[np.abs(rawCorr) < tol] = 0
        sums_x[np.abs(sums_x) < tol] = 0
        sumOfSquares_x[np.abs(sumOfSquares_x) < tol] = 0
        N = np.round(N)
        N[N <= 1] = np.nan
        # [Step 4] Compute correlation matrix
        mapStd = np.sqrt((sumOfSquares_x * N) - sums_x**2)
        mapCovar = (rawCorr * N) - sums_x * sums_x[::-1, :, :][:, ::-1, :][:, :, :]
        A = np.squeeze(mapCovar / mapStd / mapStd[::-1, :, :][:, ::-1, :][:, :, :])

        return np.asarray(A)

    def grid_field_props(self, A, maxima="centroid", allProps=True, **kwargs):
        """
        Extracts various measures from a spatial autocorrelogram
        Parameters
        ----------
        A : array_like
            The spatial autocorrelogram (SAC)
        maxima : str, optional
            The method used to detect the peaks in the SAC.
            Legal values are 'single' and 'centroid'. Default 'centroid'
        allProps : bool, optional
            Whether to return a dictionary that contains the attempt to fit an
            ellipse around the edges of the central size peaks. See below
            Default True
        Returns
        -------
        props : dict
            A dictionary containing measures of the SAC. Keys include:
            * gridness scorescore
            * scale
            * orientation
            * coordinates of the peaks (nominally 6) closest to SAC centre
            * a binary mask around the extent of the 6 central fields
            * values of the rotation procedure used to calculate gridness
            * ellipse axes and angle (if allProps is True and the it worked)

        Notes
        -----
        The output from this method can be used as input to the show() method
        of this class.
        When it is the plot produced will display a lot more informative.

        See Also
        --------
        ephysiopy.common.binning.autoCorr2D()
        """

        A_tmp = copy.copy(A)
        A_tmp[~np.isfinite(A.astype(float))] = -1
        A_tmp[A_tmp <= 0] = -1
        A_sz = np.array(np.shape(A.astype(float)))
        # [STAGE 1] find peaks & identify 7 closest to centre
        if "min_distance" in kwargs:
            min_distance = kwargs.pop("min_distance")
        else:
            min_distance = np.ceil(np.min(A_sz / 2) / 8.0).astype(int)

        peak_idx, field_labels = self._get_field_labels(A_tmp, neighbours=7, **kwargs)

        # a fcn for the labeled_comprehension function that returns
        # linear indices in A where the values in A for each label are
        # greater than half the max in that labeled region

        def fn(val, pos):
            return pos[val > (np.max(val) / 2)]

        nLbls = np.max(field_labels)
        indices = ndimage.labeled_comprehension(A_tmp, field_labels, np.arange(0, nLbls), fn, np.ndarray, 0, True)
        # turn linear indices into coordinates
        coords = [np.unravel_index(i, np.shape(A)) for i in indices]
        half_peak_labels = np.zeros_like(A)
        for peak_id, coord in enumerate(coords):
            xc, yc = coord
            half_peak_labels[xc, yc] = peak_id

        # Get some statistics about the labeled regions
        # fieldPerim = bwperim(half_peak_labels)
        lbl_range = np.arange(0, nLbls)
        # meanRInLabel = ndimage.mean(A, half_peak_labels, lbl_range)
        # nPixelsInLabel = np.bincount(np.ravel(half_peak_labels.astype(int)))
        # sumRInLabel = ndimage.sum_labels(A, half_peak_labels, lbl_range)
        # maxRInLabel = ndimage.maximum(A, half_peak_labels, lbl_range)
        peak_coords = ndimage.maximum_position(A, half_peak_labels, lbl_range)

        # Get some distance and morphology measures
        centre = np.floor(np.array(np.shape(A)) / 2)
        centred_peak_coords = peak_coords - centre
        peak_dist_to_centre = np.hypot(centred_peak_coords.T[0], centred_peak_coords.T[1])
        closest_peak_idx = np.argsort(peak_dist_to_centre)
        central_peak_label = closest_peak_idx[0]
        closest_peak_idx = closest_peak_idx[1 : np.min((7, len(closest_peak_idx) - 1))]
        # closest_peak_idx should now the indices of the labeled 6 peaks
        # surrounding the central peak at the image centre
        scale = np.median(peak_dist_to_centre[closest_peak_idx])
        orientation = np.nan
        orientation = self.grid_orientation(centred_peak_coords, closest_peak_idx)

        central_pt = peak_coords[central_peak_label]
        x = np.linspace(-central_pt[0], central_pt[0], A_sz[0])
        y = np.linspace(-central_pt[1], central_pt[1], A_sz[1])
        xv, yv = np.meshgrid(x, y, indexing="ij")
        dist_to_centre = np.hypot(xv, yv)
        # get the max distance of the half-peak width labeled fields
        # from the centre of the image
        max_dist_from_centre = 0
        for peak_id, _coords in enumerate(coords):
            if peak_id in closest_peak_idx:
                xc, yc = _coords
                if np.any(xc) and np.any(yc):
                    xc = xc - np.floor(A_sz[0] / 2)
                    yc = yc - np.floor(A_sz[1] / 2)
                    d = np.max(np.hypot(xc, yc))
                    if d > max_dist_from_centre:
                        max_dist_from_centre = d

        # Set the outer bits and the central region of the SAC to nans
        # getting ready for the correlation procedure
        dist_to_centre[np.abs(dist_to_centre) > max_dist_from_centre] = 0
        dist_to_centre[half_peak_labels == central_peak_label] = 0
        dist_to_centre[dist_to_centre != 0] = 1
        dist_to_centre = dist_to_centre.astype(bool)
        sac_middle = copy.copy(A)
        sac_middle[~dist_to_centre] = np.nan

        if "step" in kwargs.keys():
            step = kwargs.pop("step")
        else:
            step = 30

        try:
            gridscore, rotationCorrVals, rotationArr = self.gridness(sac_middle, step)
        except Exception:
            gridscore, rotationCorrVals, rotationArr = np.nan, np.nan, np.nan

        im_centre = central_pt

        if allProps:
            # attempt to fit an ellipse around the outer edges of the nearest
            # peaks to the centre of the SAC. First find the outer edges for
            # the closest peaks using a ndimages labeled_comprehension
            try:

                def fn2(val, pos):
                    xc, yc = np.unravel_index(pos, A_sz)
                    xc = xc - np.floor(A_sz[0] / 2)
                    yc = yc - np.floor(A_sz[1] / 2)
                    idx = np.argmax(np.hypot(xc, yc))
                    return xc[idx], yc[idx]

                ellipse_coords = ndimage.labeled_comprehension(A, half_peak_labels, closest_peak_idx, fn2, tuple, 0, True)

                ellipse_fit_coords = np.array([(x, y) for x, y in ellipse_coords])
                from skimage.measure import EllipseModel

                E = EllipseModel()
                E.estimate(ellipse_fit_coords)
                im_centre = E.params[0:2]
                ellipse_axes = E.params[2:4]
                ellipse_angle = E.params[-1]
                ellipseXY = E.predict_xy(np.linspace(0, 2 * np.pi, 50), E.params)

                # get the min containing circle given the eliipse minor axis
                from skimage.measure import CircleModel

                _params = list(im_centre)
                _params.append(np.min(ellipse_axes))
                circleXY = CircleModel().predict_xy(np.linspace(0, 2 * np.pi, 50), params=_params)
            except (TypeError, ValueError):  # non-iterable x and y
                ellipse_axes = None
                ellipse_angle = (None, None)
                ellipseXY = None
                circleXY = None

        # collect all the following keywords into a dict for output
        closest_peak_coords = np.array(peak_coords)[closest_peak_idx]
        dictKeys = (
            "gridscore",
            "scale",
            "orientation",
            "closest_peak_coords",
            "dist_to_centre",
            "ellipse_axes",
            "ellipse_angle",
            "ellipseXY",
            "circleXY",
            "im_centre",
            "rotationArr",
            "rotationCorrVals",
        )
        outDict = dict.fromkeys(dictKeys, np.nan)
        for thiskey in outDict.keys():
            outDict[thiskey] = locals()[thiskey]
            # neat trick: locals is a dict holding all locally scoped variables
        return outDict

    def polar(self, x, y, deg=False):
        """
        Converts from rectangular coordinates to polar ones
        Parameters
        ----------
        x, y : array_like, list_like
            The x and y coordinates
        deg : int
            radian if deg=0; degree if deg=1
        Returns
        -------
        p : array_like
            The polar version of x and y
        """
        if deg:
            return np.hypot(x, y), 180.0 * np.arctan2(y, x) / np.pi
        else:
            return np.hypot(x, y), np.arctan2(y, x)

    def get_scores(self, rate_map):
        """Get summary of scrores for grid cells.

        Parameters:
        -----------
        rate_map : np.ndarray
            2D array of firing rate map.

        Returns:
        --------
        sac : np.ndarray
            2D array of spatial autocorrelation.
        stGrd : dict
            Dictionary of grid cell properties.

        """
        sac = self.autoCorr2D(rate_map, ~np.isfinite(rate_map), tol=1e-11)
        stGrd = self.grid_field_props(sac)

        return (sac, stGrd)

    def _get_field_labels(self, A: np.ndarray, **kwargs) -> tuple:
        """
        Returns a labeled version of A after finding the peaks
        in A and finding the watershed basins from the markers
        found from those peaks. Used in field_props() and
        grid_field_props()
        Parameters
        -----------------
        A : np.ndarray
        Valid kwargs:
        min_distance : float
            The distance in bins between fields to separate the regions
            of the image
        clear_border : boolskimage
            Input to skimage.feature.peak_local_max. The number of
            pixels to ignore at the edge of the image
        """
        clear_border = True
        if "clear_border" in kwargs:
            clear_border = kwargs.pop("clear_border")

        min_distance = 1
        if "min_distance" in kwargs:
            min_distance = kwargs.pop("min_distance")

        A[~np.isfinite(A)] = -1
        A[A < 0] = -1

        peak_coords = skimage.feature.peak_local_max(A, min_distance=min_distance, exclude_border=clear_border)
        peaksMask = np.zeros_like(A, dtype=bool)
        peaksMask[tuple(peak_coords.T)] = True
        peaksLabel, nLbls = ndimage.label(peaksMask)
        ws = watershed(image=-A, markers=peaksLabel)
        return peak_coords, ws

    def grid_orientation(self, peakCoords, closestPeakIdx):
        """
        Calculates the orientation angle of a grid field.
        The orientation angle is the angle of the first peak working
        counter-clockwise from 3 o'clock
        Parameters
        ----------
        peakCoords : array_like
            The peak coordinates as pairs of xy
        closestPeakIdx : array_like
            A 1D array of the indices in peakCoords of the peaks closest
            to the centre of the SAC
        Returns
        -------
        peak_orientation : float
            The first value in an array of the angles of the peaks in the SAC
            working counter-clockwise from a line extending from the
            middle of the SAC to 3 o'clock.
        """
        if len(peakCoords) < 3 or closestPeakIdx.size == 0:
            return np.nan
        else:
            # Assume that the first entry in peakCoords is
            # the central peak of the SAC
            peaks = peakCoords[closestPeakIdx]
            peaks = peaks - peakCoords[closestPeakIdx[0]]
            theta = self.polar(peaks[:, 1], -peaks[:, 0], deg=1)[1]
            return np.sort(theta.compress(theta >= 0))[0]

    def plot_sac(
        self, sac, mask_params=None, ax=None, title=None, score="", *args, **kwargs
    ):  # pylint: disable=keyword-arg-before-vararg
        """Plot spatial autocorrelogram.
        Parameters

        ----------
        sac : np.ndarray
            2D array of spatial autocorrelation.
        mask_params : tuple, optional
            (inner_radius, outer_radius) of the ring mask.
        ax : matplotlib.axes.Axes, optional
            The axes to plot on.
        title : str, optional
            The title of the plot.
        score : str, optional
            The score of the plot.
        *args, **kwargs : optional
            Additional arguments to pass to ax.imshow().

        Returns

        -------
        ax : matplotlib.axes.Axes
            The axes the plot was drawn on.

        """
        if ax is None:
            ax = plt.gca()
        # Plot the sac
        ax = make_plot_rate_map(sac, ax, "sac plot with gridscore:" + score, "width", "depth", "sac rate")
        # ax.imshow(sac, interpolation="none", *args, **kwargs)
        # ax.pcolormesh(useful_sac, *args, **kwargs)
        # Plot a ring for the adequate mask
        if mask_params is not None:
            center = self._nbins - 1
            ax.add_artist(
                plt.Circle(
                    (center, center),
                    mask_params[0] * self._nbins,
                    # lw=bump_size,
                    fill=False,
                    edgecolor="k",
                )
            )
            ax.add_artist(
                plt.Circle(
                    (center, center),
                    mask_params[1] * self._nbins,
                    # lw=bump_size,
                    fill=False,
                    edgecolor="k",
                )
            )
        ax.axis("off")
        if title is not None:
            ax.set_title(title)

    def gridness(self, image, step):
        """
        Calculates the gridness score in a grid cell SAC.

        Briefly, the data in `image` is rotated in `step` amounts and
        each rotated array is correlated with the original.
        The maximum of the values at 30, 90 and 150 degrees
        is the subtracted from the minimum of the values at 60, 120
        and 180 degrees to give the grid score.

        Parameters
        ----------
        image : array_like
            The spatial autocorrelogram
        step : int, optional
            The amount to rotate the SAC in each step of the rotational
            correlation procedure

        Returns
        -------
        gridmeasures : 3-tuple
            The gridscore, the correlation values at each `step` and
            the rotational array

        Notes
        -----
        The correlation performed is a Pearsons R. Some rescaling of the
        values in `image` is performed following rotation.

        See Also
        --------
        skimage.transform.rotate : for how the rotation of `image` is done
        skimage.exposure.rescale_intensity : for the resscaling following
        rotation

        """
        # TODO: add options in here for whether the full range of correlations
        # are wanted or whether a reduced set is wanted (i.e. at the 30-tuples)
        from collections import OrderedDict

        rotationalCorrVals = OrderedDict.fromkeys(np.arange(0, 181, step), np.nan)
        rotationArr = np.zeros(len(rotationalCorrVals)) * np.nan
        # autoCorrMiddle needs to be rescaled or the image rotation falls down
        # as values are cropped to lie between 0 and 1.0
        in_range = (np.nanmin(image), np.nanmax(image))
        out_range = (0, 1)
        import skimage

        autoCorrMiddleRescaled = skimage.exposure.rescale_intensity(image, in_range, out_range)
        origNanIdx = np.isnan(autoCorrMiddleRescaled.ravel())
        for idx, angle in enumerate(rotationalCorrVals.keys()):
            rotatedA = skimage.transform.rotate(autoCorrMiddleRescaled, angle=angle, cval=np.nan, order=3)
            # ignore nans
            rotatedNanIdx = np.isnan(rotatedA.ravel())
            allNans = np.logical_or(origNanIdx, rotatedNanIdx)
            # get the correlation between the original and rotated images
            rotationalCorrVals[angle] = stats.pearsonr(autoCorrMiddleRescaled.ravel()[~allNans], rotatedA.ravel()[~allNans])[0]
            rotationArr[idx] = rotationalCorrVals[angle]
        gridscore = np.nanmin((rotationalCorrVals[60], rotationalCorrVals[120])) - np.nanmax(
            (rotationalCorrVals[150], rotationalCorrVals[30], rotationalCorrVals[90])
        )
        return gridscore, rotationalCorrVals, rotationArr

    def plot_grid_score(self, r_out_im, plot=True, ax=None):
        """
        Get the grid score of the network

        Parameters
        ----------
        r_out_im : np.ndarray
            2D array of spatial autocorrelation.
        plot : bool, optional
            Whether to plot the spatial autocorrelation.
        ax : matplotlib.axes.Axes, optional
            The axes to plot on.

        Returns
        -------
        grid_score : float
            Grid score of the network
        """
        score = self.get_scores(np.asarray(r_out_im))
        if plot:
            if ax is None:
                fig, ax = plt.subplots()
            gridscore = str(np.around(score[1]["gridscore"], decimals=4, out=None))
            self.plot_sac(sac=score[0], ax=ax, score=gridscore)
        else:
            ax.text(0, 0.7, "Grid_score: " + str(np.around(score[1]["gridscore"], decimals=4, out=None)), fontsize=10)
            ax.set_axis_off()
        return score
