# 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 matplotlib.pyplot as plt
from matplotlib import ticker
import numpy as np
import matplotlib.gridspec as gridspec
import matplotlib.pyplot as plt
import matplotlib.patches as mpatches

DEFAULT_CMAP = 'magma'

def plot_ellipses(gaussian_fit_dict, ax=None, show=True, close=True, save_file_name=None, color='b'):
    '''Example Usage:
    oeid, cell_index, stimulus = 512176430, 12, 'locally_sparse_noise'
    brain_observatory_cache = BrainObservatoryCache()
    data_set = brain_observatory_cache.get_ophys_experiment_data(oeid)
    lsn = LocallySparseNoise(data_set, stimulus)
    result = compute_receptive_field_with_postprocessing(data_set, cell_index, stimulus, alpha=.05, number_of_shuffles=5000)
    plot_ellipses(result['off']['gaussian_fit'], color='r')
    '''

    if ax is None:
        fig, ax = plt.subplots(1)
        ax.set_xlim(0, 130)
        ax.set_ylim(0, 74)
        plt.axis('off')

    on_comp = len(gaussian_fit_dict['attrs']['center_x'])
    for i in range(on_comp):
        xy = (gaussian_fit_dict['attrs']['center_x'][i], gaussian_fit_dict['attrs']['center_y'][i])
        width = 3 * np.abs(gaussian_fit_dict['attrs']['width_x'][i])
        height = 3 * np.abs(gaussian_fit_dict['attrs']['width_y'][i])
        angle = gaussian_fit_dict['attrs']['rotation'][i]
        if np.logical_not(any(np.isnan(xy))):
            ellipse = mpatches.Ellipse(xy, width=width, height=height, angle=angle, lw=2, edgecolor=color,
                                       facecolor='none')
            ax.add_artist(ellipse)

    if not save_file_name is None:
        fig.savefig(save_file_name)

    if show == True:
        plt.show()

    if close:
        plt.close(fig)

    return ax

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

def plot_chi_square_summary(rf_data, ax=None, cax=None, cmap=DEFAULT_CMAP):
    if ax is None:
        ax = plt.gca()

    chi_squared_grid = rf_data['chi_squared_analysis']['pvalues']['data']
    chi_square_grid_NLL = pvalue_to_NLL(chi_squared_grid)
    clim = (0, max(2,chi_square_grid_NLL.max()))
    img = ax.imshow(chi_square_grid_NLL, interpolation='none', origin='lower', clim=clim, cmap=cmap)

    if cax is None:
        cb = ax.figure.colorbar(img, ax=ax, ticks=clim)
    else:
        cb = ax.figure.colorbar(img, cax=cax, ticks=clim)

    tick_locator = ticker.MaxNLocator(nbins=5)
    cb.locator = tick_locator
    cb.update_ticks()
    ax.axes.get_xaxis().set_visible(False)
    ax.axes.get_yaxis().set_visible(False)
    ax.set_title('Significant: %s (min_p=%s)' % (rf_data['chi_squared_analysis']['attrs']['significant'], 
                                                 rf_data['chi_squared_analysis']['attrs']['min_p']) )    

def plot_msr_summary(lsn, cell_index, ax_on, ax_off, ax_cbar=None, cmap=None):
    min_clim = lsn.mean_response[:, :, cell_index,:].min()
    max_clim = lsn.mean_response[:, :, cell_index,:].max()
    plot_fields(lsn.mean_response[:, :, cell_index, 0], 
                lsn.mean_response[:, :, cell_index, 1], 
                ax_on, ax_off, clim=(min_clim, max_clim), cmap=cmap, cbar_axes=ax_cbar)

def plot_fields(on_data, off_data, on_axes, off_axes, cbar_axes=None, clim=None, cmap=DEFAULT_CMAP):
    if cbar_axes is None:
        on_axes.figure.subplots_adjust(right=0.9)
        cbar_axes = on_axes.figure.add_axes([0.93, 0.37, 0.02, .28])
    
    if clim is None:
        clim_max = max(np.nanmax(on_data), np.nanmax(off_data))
        clim = (0,clim_max)
    on_axes.imshow(on_data, clim=clim, cmap=cmap, interpolation='none', origin='lower')
    on_axes.set_title("on")
    img = off_axes.imshow(off_data, clim=clim, cmap=cmap, interpolation='none', origin='lower')
    off_axes.set_title("off")
    cb = cbar_axes.figure.colorbar(img, cax=cbar_axes, ticks=clim)
    tick_locator = ticker.MaxNLocator(nbins=5)
    cb.locator = tick_locator
    cb.update_ticks()
    for frame in [on_axes, off_axes]:
        frame.axes.get_xaxis().set_visible(False)
        frame.axes.get_yaxis().set_visible(False)

def plot_rts_summary(rf_data, ax_on, ax_off, ax_cbar=None, cmap=DEFAULT_CMAP):
    rts_on = rf_data['on']['rts']['data']
    rts_off = rf_data['off']['rts']['data']
    plot_fields(rts_on, rts_off, ax_on, ax_off, cbar_axes=ax_cbar, cmap=cmap)

def plot_rts_blur_summary(rf_data, ax_on, ax_off, ax_cbar=None, cmap=DEFAULT_CMAP):
    rts_on_blur = rf_data['on']['rts_convolution']['data']
    rts_off_blur = rf_data['off']['rts_convolution']['data']
    plot_fields(rts_on_blur, rts_off_blur, ax_on, ax_off, cbar_axes=ax_cbar, cmap=cmap)

def plot_p_values(rf_data, ax_on, ax_off, ax_cbar=None, cmap=DEFAULT_CMAP):
    pvalues_on = rf_data['on']['pvalues']['data']
    pvalues_off = rf_data['off']['pvalues']['data']
    clim_max = max(pvalues_on.max(), pvalues_off.max())
    plot_fields(pvalues_on, pvalues_off, ax_on, ax_off, cbar_axes=ax_cbar, clim=(0, clim_max/2), cmap=cmap)

def plot_mask(rf_data, ax_on, ax_off, ax_cbar=None, cmap=DEFAULT_CMAP):
    pvalues_on = rf_data['on']['pvalues']['data']
    pvalues_off = rf_data['off']['pvalues']['data']

    rf_on = pvalues_on.copy()
    rf_off = pvalues_off.copy()

    rf_on[np.logical_not(rf_data['on']['fdr_mask']['data'].sum(axis=0))] = np.nan
    rf_off[np.logical_not(rf_data['off']['fdr_mask']['data'].sum(axis=0))] = np.nan

    plot_fields(rf_on, rf_off, ax_on, ax_off, cbar_axes=ax_cbar, cmap=cmap)

def plot_gaussian_fit(rf_data, ax_on, ax_off, ax_cbar=None, cmap=DEFAULT_CMAP):

    gf_on_exists = 'gaussian_fit' in rf_data['on']
    gf_off_exists = 'gaussian_fit' in rf_data['off']

    if not gf_on_exists and not gf_off_exists:
        return

    img_data_on = rf_data['on']['gaussian_fit']['data'].sum(axis=0) if gf_on_exists else None
    img_data_off = rf_data['off']['gaussian_fit']['data'].sum(axis=0) if gf_off_exists else None

    if gf_on_exists and gf_off_exists:
        plot_fields(img_data_on, img_data_off, ax_on, ax_off, cbar_axes=ax_cbar, cmap=cmap)
    else:
        if gf_on_exists:
            img_data_off = np.zeros(img_data_on.shape)
        else:
            img_data_on = np.zeros(img_data_off.shape)

        plot_fields(img_data_on, img_data_off, ax_on, ax_off, cbar_axes=ax_cbar, cmap=cmap)

def plot_receptive_field_data(rf, lsn, show=True, save_file_name=None, close=True, cmap=DEFAULT_CMAP):
    cell_index = rf['attrs']['cell_index']

    # Prepare plotting figure:n
    number_of_major_rows = 7 if lsn else 6
    pwidth = 1.7
    pheight = 1.0
    fig = plt.figure(figsize=(pwidth*2.3, pheight*number_of_major_rows))
    gsp = gridspec.GridSpec(number_of_major_rows, 3, width_ratios=[1,1,.1], right=0.9)
    ax_list = []

    # Plot chi-square summary:
    row = 0
    curr_axes = fig.add_subplot(gsp[row,:2])
    cbar_axes = fig.add_subplot(gsp[row,-1])
    ax_list += [curr_axes]
    plot_chi_square_summary(rf, ax=curr_axes, cax=cbar_axes, cmap=cmap)

    # MSR plot:
    if not lsn is None:
        row += 1
        curr_on_axes = fig.add_subplot(gsp[row, 0])
        curr_off_axes = fig.add_subplot(gsp[row, 1])
        cbar_axes = fig.add_subplot(gsp[row, 2])
        ax_list += [curr_on_axes, curr_off_axes]
        plot_msr_summary(lsn, cell_index, curr_on_axes, curr_off_axes, cbar_axes, cmap=cmap)

    # RTS no blur:
    row += 1
    curr_on_axes = fig.add_subplot(gsp[row, 0])
    curr_off_axes = fig.add_subplot(gsp[row, 1])
    cbar_axes = fig.add_subplot(gsp[row,2])
    ax_list += [curr_on_axes, curr_off_axes]
    plot_rts_summary(rf, curr_on_axes, curr_off_axes, cbar_axes, cmap=cmap)

    # RTS no blur:
    row += 1
    curr_on_axes = fig.add_subplot(gsp[row, 0])
    curr_off_axes = fig.add_subplot(gsp[row, 1])
    cbar_axes = fig.add_subplot(gsp[row,2])
    ax_list += [curr_on_axes, curr_off_axes]
    plot_rts_blur_summary(rf, curr_on_axes, curr_off_axes, cbar_axes, cmap=cmap)

    # PValues:
    row += 1
    curr_on_axes = fig.add_subplot(gsp[row, 0])
    curr_off_axes = fig.add_subplot(gsp[row, 1])
    cbar_axes = fig.add_subplot(gsp[row,2])
    ax_list += [curr_on_axes, curr_off_axes]
    plot_p_values(rf, curr_on_axes, curr_off_axes, cbar_axes, cmap=cmap)

    # Mask:
    row += 1
    curr_on_axes = fig.add_subplot(gsp[row, 0])
    curr_off_axes = fig.add_subplot(gsp[row, 1])
    cbar_axes = fig.add_subplot(gsp[row,2])
    ax_list += [curr_on_axes, curr_off_axes]
    plot_mask(rf, curr_on_axes, curr_off_axes, cbar_axes, cmap=cmap)

    # Gaussian fit:
    row += 1
    curr_on_axes = fig.add_subplot(gsp[row, 0])
    curr_off_axes = fig.add_subplot(gsp[row, 1])
    cbar_axes = fig.add_subplot(gsp[row,2])
    ax_list += [curr_on_axes, curr_off_axes]
    plot_gaussian_fit(rf, curr_on_axes, curr_off_axes, cbar_axes, cmap=cmap)

    # gs.tight_layout(fig)

    plt.subplots_adjust(top=0.95)

    for ax in ax_list:
        ax.set_adjustable('box-forced')

    if not save_file_name is None:
        fig.savefig(save_file_name)

    if show == True:
        plt.show()

    if close:
        plt.close(fig)
