# Allen Institute Software License - This software license is the 2-clause BSD
# license plus a third clause that prohibits redistribution for commercial
# purposes without further permission.
#
# Copyright 2017. Allen Institute. All rights reserved.
#
# Redistribution and use in source and binary forms, with or without
# modification, are permitted provided that the following conditions are met:
#
# 1. Redistributions of source code must retain the above copyright notice,
# this list of conditions and the following disclaimer.
#
# 2. Redistributions in binary form must reproduce the above copyright notice,
# this list of conditions and the following disclaimer in the documentation
# and/or other materials provided with the distribution.
#
# 3. Redistributions for commercial purposes are not permitted without the
# Allen Institute's written permission.
# For purposes of this license, commercial purposes is the incorporation of the
# Allen Institute's software into anything for which you will charge fees or
# other compensation. Contact terms@alleninstitute.org for commercial licensing
# opportunities.
#
# THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS"
# AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE
# IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE
# ARE DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR CONTRIBUTORS BE
# LIABLE FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR
# CONSEQUENTIAL DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF
# SUBSTITUTE GOODS OR SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS
# INTERRUPTION) HOWEVER CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN
# CONTRACT, STRICT LIABILITY, OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE)
# ARISING IN ANY WAY OUT OF THE USE OF THIS SOFTWARE, EVEN IF ADVISED OF THE
# POSSIBILITY OF SUCH DAMAGE.
#
import numpy as np
import scipy.interpolate as si
import scipy.ndimage.filters as filt
import scipy.stats as stats


ON_LUMINANCE = 255
OFF_LUMINANCE = 0


def chi_square_binary(events, LSN_template):
    # note: can only be applied to binary events for trial responses
    #
    # *****INPUT*****
    # events: 2D numpy bool with shape (num_trials,num_cells) for presence
    #   or absence of response on a given trial
    # LSN_template: 3D numpy int8 with shape (num_trials,num_y_pixels,num_x_pixels)
    #   for luminance at each pixel location
    #
    # *****OUTPUT*****
    # chi_square_grid_NLL: 3D numpy float with shape (num_cells,num_y_pixels,num_x_pixels)
    #   that gives the p value for the hypothesis that a receptive field is contained
    #   within a 7x7 pixel mask centered on a given pixel location as measured
    #   by a chi-square test for the responses among the pixels (both on and off)
    #   that fall within the mask.

    num_trials = np.shape(events)[0]
    num_cells = np.shape(events)[1]
    num_y = np.shape(LSN_template)[1]
    num_x = np.shape(LSN_template)[2]

    # for each pixel location, get a mask that is centered on that location
    #    disc_masks has shape (num_y,num_x,num_y,num_x)
    disc_masks = get_disc_masks(LSN_template)

    # determine which trials each pixel is active (i.e not gray),
    #   broken up by ON and OFF pixels.
    #   trial_matrix has shape (num_y,num_x,2,num_trials)
    trial_matrix = build_trial_matrix(LSN_template, num_trials)

    # get the total number of trials each pixel is active (i.e. not gray)
    #   trials_per_pixel has shape (num_y,num_x,2)
    trials_per_pixel = np.sum(trial_matrix, axis=3)

    # get the sum of the number of events across all trials that each pixel is active
    #   events_per_pixel has shape (num_cells,num_y,num_x,2)
    events_per_pixel = get_events_per_pixel(events, trial_matrix)

    # smooth stimulus-triggered average spatially with a gaussian
    for n in range(num_cells):
        for on_off in range(2):
            events_per_pixel[n, :, :, on_off] = smooth_STA(events_per_pixel[n, :, :, on_off])

    # calculate the p_value for each exclusion region
    chi_square_grid = np.zeros((num_cells, num_y, num_x))
    for y in range(num_y):
        for x in range(num_x):
            exclusion_mask = np.ones((num_y, num_x, 2)) * disc_masks[y, x, :, :].reshape(num_y, num_x, 1)
            p_vals, __ = chi_square_within_mask(exclusion_mask, events_per_pixel, trials_per_pixel)
            chi_square_grid[:, y, x] = p_vals

    return chi_square_grid


def get_peak_significance(chi_square_grid_NLL,
                          LSN_template,
                          alpha=0.05):
    # ****INPUT*****
    # chi_square_grid_NLL: result of chi_square_binary(events,LSN_template)
    # LSN_template: 3D numpy int8 with shape (num_trials,num_y_pixels,num_x_pixels)
    #   for luminance at each pixel location
    #
    # *****OUTPUT*****
    # significant_cells: 1D numpy bool with shape (num_cells,) that indicates
    #   whether or not each cell has a location with a significant chance of a
    #   true receptive field
    # best_exclusion_region_list: list of 2D numpy bool with len (num_cells). For each cell,
    #   the array is a mask for the best pixel, or all false

    num_cells = np.shape(chi_square_grid_NLL)[0]
    num_y = np.shape(chi_square_grid_NLL)[1]
    num_x = np.shape(chi_square_grid_NLL)[2]

    chi_square_grid = NLL_to_pvalue(chi_square_grid_NLL)

    # get the average size of all masks in units of number of pixels
    disc_masks = get_disc_masks(LSN_template)
    pixels_per_mask_per_pixel = np.sum(disc_masks, axis=(2, 3)).astype(float)

    # find the smallest p-value and determine if it's significant
    significant_cells = np.zeros((num_cells)).astype(bool)
    best_p = np.zeros((num_cells))
    p_value_correction_factor_per_pixel = (1.0 * num_y * num_x / pixels_per_mask_per_pixel)

    best_exclusion_region_list = []
    corrected_p_value_array_list = []
    for n in range(num_cells):

        # Sidak correction:
        p_value_corrected_per_pixel = 1-np.power((1-chi_square_grid[n, :,:]), p_value_correction_factor_per_pixel)
        corrected_p_value_array_list.append(p_value_corrected_per_pixel)

        y, x = np.unravel_index(p_value_corrected_per_pixel.argmin(), (num_y, num_x))

        # if more than one p-value that maxes out, use the median location
        if np.sum(p_value_corrected_per_pixel == 0.0) > 1:

            y, x = np.unravel_index(np.argwhere(p_value_corrected_per_pixel.flatten() == 0.0)[:, 0], (num_y, num_x))
            center_y, center_x = locate_median(y, x)

        best_p[n] = p_value_corrected_per_pixel[y,x]
        if best_p[n] < alpha:
            significant_cells[n] = True
            best_exclusion_region_list.append(disc_masks[y, x, :,:].astype(np.bool))
        else:
            best_exclusion_region_list.append(np.zeros((disc_masks.shape[0], disc_masks.shape[1]), dtype=np.bool))

    return significant_cells, best_p, corrected_p_value_array_list, best_exclusion_region_list


def locate_median(y, x):

    med_x = np.median(x)
    med_y = np.median(y)
    center_x = x[0]
    center_y = y[0]

    for i in range(len(x)):
        dx = x[i] - med_x
        dy = y[i] - med_y
        dc_x = center_x - med_x
        dc_y = center_y - med_y

        if np.sqrt(dx ** 2 + dy ** 2) < np.sqrt(dc_x ** 2 + dc_y ** 2):
            center_x = x[i]
            center_y = y[i]

    return center_y, center_x


def pvalue_to_NLL(p_values, max_NLL=10.0):
    return np.where(p_values == 0.0, max_NLL, -np.log10(p_values))


def NLL_to_pvalue(NLLs, log_base=10.0):
    return (log_base ** (-NLLs))


def get_events_per_pixel(responses_np, trial_matrix):
    '''Obtain a matrix linking cellular responses to pixel activity.

    Parameters
    ----------
    responses_np : np.ndarray
        Dimensions are (nTrials, nCells). Boolean values indicate presence/absence
        of a response on a given trial.
    trial_matrix : np.ndarray
        Dimensions are (nYPixels, nXPixels, {on, off}, nTrials). Boolean values 
        indicate that a pixel was on/off on a particular trial.

    Returns
    -------
    events_per_pixel : np.ndarray
        Dimensions are (nCells, nYPixels, nXPixels, {on, off}). Values for each 
        cell, pixel, and on/off state are the sum of events for that cell across 
        all trials where the pixel was in the on/off state.

    '''

    num_cells = np.shape(responses_np)[1]
    num_y = np.shape(trial_matrix)[0]
    num_x = np.shape(trial_matrix)[1]

    events_per_pixel = np.zeros((num_cells, num_y, num_x, 2))
    for y in range(num_y):
        for x in range(num_x):
            for on_off in range(2):
                frames = np.argwhere(trial_matrix[y, x, on_off, :])[:, 0]
                events_per_pixel[:, y, x, on_off] = np.sum(responses_np[frames, :], axis=0)

    return events_per_pixel


def smooth_STA(STA, gauss_std=0.75, total_degrees=64):
    '''Smooth an image by convolution with a gaussian kernel

    Parameters
    ----------
    STA : np.ndarray
        Input image
    gauss_std : numeric, optional
        Standard deviation of the gaussian kernel. Will be applied to the 
        upsampled image, so units are visual degrees. Default is 0.75
    total_degrees : int, optional
        Size in visual degrees of the input image along its zeroth (row) axis.
        Used to set the scale factor for up/downsampling.

    Returns
    -------
    STA_smoothed : np.ndarray
        Smoothed image
    '''

    deg_per_pnt = total_degrees // STA.shape[0]
    STA_interpolated = interpolate_RF(STA, deg_per_pnt)
    STA_interpolated_smoothed = filt.gaussian_filter(STA_interpolated, gauss_std)
    STA_smoothed = deinterpolate_RF(STA_interpolated_smoothed, STA.shape[1], STA.shape[0], deg_per_pnt)

    return STA_smoothed


def interpolate_RF(rf_map, deg_per_pnt):
    '''Upsample an image
      
    Parameters
    ----------
    rf_map : np.ndarray
        Input image
    deg_per_pnt : numeric
        scale factor

    Returns
    -------
    interpolated : np.ndarray
        Upsampled image
    '''

    x_pnts = np.shape(rf_map)[1]
    y_pnts = np.shape(rf_map)[0]

    x_coor = np.arange(-(x_pnts - 1) * deg_per_pnt / 2, (x_pnts + 1) * deg_per_pnt / 2, deg_per_pnt)
    y_coor = np.arange(-(y_pnts - 1) * deg_per_pnt / 2, (y_pnts + 1) * deg_per_pnt / 2, deg_per_pnt)

    x_interpolated = np.arange(-(x_pnts - 1) * deg_per_pnt / 2, deg_per_pnt / 2 + (x_pnts / 2 - 1) * deg_per_pnt + 1, 1)
    y_interpolated = np.arange(-(y_pnts - 1) * deg_per_pnt / 2, deg_per_pnt / 2 + (y_pnts / 2 - 1) * deg_per_pnt + 1, 1)

    interpolated = si.interp2d(x_coor, y_coor, rf_map)
    interpolated = interpolated(x_interpolated, y_interpolated)

    return interpolated


def deinterpolate_RF(rf_map, x_pnts, y_pnts, deg_per_pnt):
    '''Downsample an image

    Parameters
    ----------
    rf_map : np.ndarray
        Input image
    x_pnts : np.ndarray
        Count of sample points along the first (column) axis
    y_pnts : np.ndarray
        Count of sample points along the zeroth (row) axis
    deg_per_pnt : numeric
        scale factor
        
    Returns
    -------
    sampled_yx : np.ndarray
        Downsampled image
    '''

    # x_pnts = 28
    # y_pnts = 16

    x_interpolated = np.arange(-(x_pnts - 1) * deg_per_pnt / 2, deg_per_pnt / 2 + (x_pnts / 2 - 1) * deg_per_pnt + 1, 1)
    y_interpolated = np.arange(-(y_pnts - 1) * deg_per_pnt / 2, deg_per_pnt / 2 + (y_pnts / 2 - 1) * deg_per_pnt + 1, 1)

    x_deinterpolate = np.arange(0, len(x_interpolated), deg_per_pnt)
    y_deinterpolate = np.arange(0, len(y_interpolated), deg_per_pnt)

    sampled_y = rf_map[y_deinterpolate, :]
    sampled_yx = sampled_y[:, x_deinterpolate]

    return sampled_yx


def chi_square_within_mask(exclusion_mask, events_per_pixel, trials_per_pixel):
    '''Determine if cells respond preferentially to on/off pixels in a mask using
    a chi2 test.
  
    Parameters
    ----------
    exclusion_mask : np.ndarray
        Dimensions are (nYPixels, nXPixels, {on, off}). Integer indicator for INCLUSION (!) 
        of a pixel within the testing region.
    events_per_pixel : np.ndarray
        Dimensions are (nCells, nYPixels, nXPixels, {on, off}). Integer values 
        are response counts by cell to on/off luminance at each pixel.
    trials_per_pixel : np.ndarray
        Dimensions are (nYPixels, nXPixels, {on, off}). Integer values are 
        counts of trials where a pixel is on/off.

    Returns
    -------
    p_vals : np.ndarray 
        One-dimensional, of length nCells. Float values are p-values 
        for the hypothesis that a given cell has a receptive field within the 
        exclusion mask.
    chi : np.ndarray
        Dimensions are (nCells, nYPixels, nXPixels, {on, off}). Values (float) 
        are squared residual event counts divided by expected event counts.
    '''

    num_y = np.shape(exclusion_mask)[0]
    num_x = np.shape(exclusion_mask)[1]

    # d.f. is number of pixels in mask minus one
    degrees_of_freedom = int(np.sum(exclusion_mask)) - 1

    #   observed_by_pixel has shape (num_cells,num_y,num_x,2)
    expected_by_pixel = get_expected_events_by_pixel(exclusion_mask, events_per_pixel, trials_per_pixel)
    observed_by_pixel = (events_per_pixel * exclusion_mask.reshape(1, num_y, num_x, 2)).astype(float)

    # calculate test statistic given observed and expected
    residual_by_pixel = observed_by_pixel - expected_by_pixel
    chi = (residual_by_pixel ** 2) / expected_by_pixel
    chi_sum = np.nansum(chi, axis=(1, 2, 3))

    # get p-value given test statistic and degrees of freedom
    p_vals = 1.0 - stats.chi2.cdf(chi_sum, degrees_of_freedom)

    return p_vals, chi


def get_expected_events_by_pixel(exclusion_mask, events_per_pixel, trials_per_pixel):
    '''Calculate expected number of events per pixel

    Parameters
    ----------
    exclusion_mask : np.ndarray
        Dimensions are (nYPixels, nXPixels, {on, off}). Integer indicator for INCLUSION (!) 
        of a pixel within the testing region.
    events_per_pixel : np.ndarray
        Dimensions are (nCells, nYPixels, nXPixels, {on, off}). Integer values 
        are response counts by cell to on/off luminance at each pixel.
    trials_per_pixel : np.ndarray
        Dimensions are (nYPixels, nXPixels, {on, off}). Integer values are 
        counts of trials where a pixel is on/off.

    Returns
    -------
    np.ndarray :
        Dimensions (nCells, nYPixels, nXPixels, {on, off}). Float values are 
        pixelwise counts of events expected if events are evenly distributed
        in mask across trials.
    '''

    num_y = np.shape(exclusion_mask)[0]
    num_x = np.shape(exclusion_mask)[1]
    num_cells = np.shape(events_per_pixel)[0]

    exclusion_mask = exclusion_mask.reshape(1, num_y, num_x, 2)
    trials_per_pixel = trials_per_pixel.reshape(1, num_y, num_x, 2)

    masked_trials = exclusion_mask * trials_per_pixel
    masked_events = exclusion_mask * events_per_pixel

    total_trials = np.sum(masked_trials).astype(float)
    total_events_by_cell = np.sum(masked_events, axis=(1, 2, 3)).astype(float)
    
    expected_by_cell_per_trial = total_events_by_cell / total_trials
    return masked_trials * expected_by_cell_per_trial.reshape(num_cells, 1, 1, 1)


def build_trial_matrix(LSN_template, 
                       num_trials, 
                       on_off_luminance=(ON_LUMINANCE, OFF_LUMINANCE)):
    '''Construct indicator arrays for on/off pixels across trials.

    Parameters
    ----------
    LSN_template : np.ndarray
        Dimensions are (nTrials, nYPixels, nXPixels). Luminance values per pixel 
        and trial. The size of the first dimension may be larger than the num_trials 
        argument (in which case only the first num_trials slices will be used) 
        but may not be smaller.
    num_trials : int
        The number of trials (left-justified) to build indicators for.
    on_off_luminance : array-like, optional
        The zeroth element is the luminance value of a pixel when on, the first when off.
        Defaults are [255, 0].

    Returns
    -------
    trial_mat : np.ndarray
        Dimensions are (nYPixels, nXPixels, {on, off}, nTrials). Boolean values 
        indicate that a pixel was on/off on a particular trial.
    '''

    _, num_y, num_x = np.shape(LSN_template)
    trial_mat = np.zeros( (num_y, num_x, 2, num_trials), dtype=bool )

    for y in range(num_y):
        for x in range(num_x):
            for oo, on_off in enumerate(on_off_luminance):

                frame = np.argwhere( LSN_template[:num_trials, y, x] == on_off )[:, 0]
                trial_mat[y, x, oo, frame] = True

    return trial_mat


def get_disc_masks(LSN_template, radius=3, on_luminance=ON_LUMINANCE, off_luminance=OFF_LUMINANCE):
    '''Obtain an indicator mask surrounding each pixel. The mask is a square, excluding pixels which 
    are coactive on any trial with the main pixel.

    Parameters
    ----------
    LSN_template : np.ndarray
        Dimensions are (nTrials, nYPixels, nXPixels). Luminance values per pixel 
        and trial.
    radius : int
        The base mask will be a box whose sides are 2 * radius + 1 in length.
    on_luminance : int, optional
        The value of the luminance for on trials. Default is 255
    off_luminance : int, optional
        The value of the luminance for off trials. Default is 0
  

    Returns
    -------
    masks : np.ndarray
        Dimensions are (nYPixels, nXPixels, nYPixels, nXPixels). The first 2 
        dimensions describe the pixel from which the mask was computed. The last 
        2 serve as the dimensions of the mask images themselves. Masks are binary 
        arrays of type float, with 1 indicating inside, 0 outside.
    '''

    num_y = np.shape(LSN_template)[1]
    num_x = np.shape(LSN_template)[2]

    # convert template to true on trials a pixel is not gray and false when gray
    LSN_binary = np.where(LSN_template == off_luminance, 1, LSN_template)
    LSN_binary = np.where(LSN_binary == on_luminance, 1, LSN_binary)
    LSN_binary = np.where(LSN_binary == 1, 1.0, 0.0)

    # get number of trials each pixel is not gray
    on_trials = LSN_binary.sum(axis=0).astype(float)  # shape is (num_y,num_x)

    masks = np.zeros((num_y, num_x, num_y, num_x))
    for y in range(num_y):
        for x in range(num_x):
            trials_not_gray = np.argwhere( LSN_binary[:, y, x] > 0 )[:, 0]
            raw_mask = np.divide( LSN_binary[trials_not_gray, :, :].sum(axis=0), on_trials ) 

            center_y, center_x = np.unravel_index( raw_mask.argmax(), (num_y, num_x) )

            # include center pixel in mask
            raw_mask[center_y, center_x] = 0.0

            x_max = center_x + radius + 1
            if x_max > num_x:
                x_max = num_x
            x_min = center_x - radius
            if x_min < 0:
                x_min = 0
            y_max = center_y + radius + 1
            if y_max > num_y:
                y_max = num_y
            y_min = center_y - radius
            if y_min < 0:
                y_min = 0

            # don't include far away pixels that just happen
            # to not have any trials in common with center pixel
            clean_mask = np.ones(np.shape(raw_mask))
            clean_mask[y_min:y_max, x_min:x_max] = raw_mask[y_min:y_max, x_min:x_max]

            masks[y, x, :, :] = clean_mask

    masks = np.where(masks > 0, 0.0, 1.0)

    return masks

