# code for Figure 2 panels

# import
from matplotlib import pyplot as plt
from matplotlib.ticker import ScalarFormatter, NullFormatter
import numpy as np
import pandas
from scipy import stats
from importlib import reload
import spatint_utils
import os

# reload
reload(spatint_utils)

spatint_utils.plot_params()
_, _, optocolor = spatint_utils.get_colors()


class Fig2:
    """Class to for plotting panels for Fig.2"""

    def __init__(self):
        """Init class"""

        # read lgn size tuning dataframe
        parentdir = os. path. dirname(os. getcwd())
        filename = parentdir + '/data/lgn_sztun_df.pkl'
        self.lgn_sztun_df = pandas.read_pickle(
            filepath_or_buffer=filename)

    def ex_sztun_curve(self, figsize=(2.5, 2.5), ax=None):
        """Plot example dLGN size-tuning curve (Fig. 2h)

        Parameters
        -------
        figsize: tuple
            Figure size (width, height)
        ax: mpl axis
            axis for plot

        Returns
        -------
        ax: mpl axis
            axis for plot
        """

        if ax is None:
            # create figure
            f, ax = plt.subplots(figsize=figsize)

        # define example index
        ex_idx = 22

        # plot curves
        spatint_utils.plot_tun(means=self.lgn_sztun_df.loc[ex_idx]['tun_mean'],
                               sems=self.lgn_sztun_df.loc[ex_idx]['tun_sem'],
                               spons=self.lgn_sztun_df.loc[ex_idx]['tun_spon_mean'],
                               xs=self.lgn_sztun_df.loc[ex_idx]['ti_axes'],
                               c_fit=self.lgn_sztun_df.loc[ex_idx]['c_sz_fit'],
                               op_fit=self.lgn_sztun_df.loc[ex_idx]['op_sz_fit'],
                               c_prefsz=self.lgn_sztun_df.loc[ex_idx]['rfcs_76'][0],
                               op_prefsz=self.lgn_sztun_df.loc[ex_idx]['rfcs_76'][1],
                               ax=ax)
        f = plt.gcf()
        f.tight_layout()

        return ax

    def fit_norm_curves(self, ax=None, figsize=(2.5, 2.5)):
        """Plots normalized dLGN size-tuning curves (Fig. 2i)

        Parameters
        -------
        figsize: tuple
            Figure size (width, height)
        ax: mpl axis
            axis for plot

        Returns
        -------
        ax: mpl axis
            axis for plot
        """

        if ax is None:
            # create figure if ax is none
            f, ax = plt.subplots(figsize=figsize)

        # get ydata
        yconts = self.lgn_sztun_df.c_sz_fit.to_list()
        yoptos = self.lgn_sztun_df.op_sz_fit.to_list()

        # normalize
        ycont_norm = np.vstack([cont / np.nanmax(np.concatenate((cont, opto))) for cont,
                                opto in zip(yconts, yoptos)])
        yopto_norm = np.vstack([opto / np.nanmax(np.concatenate((cont, opto))) for cont,
                               opto in zip(yconts, yoptos)])

        # compute mean for both conditions
        cont_mean = np.nanmean(ycont_norm, axis=0)
        opto_mean = np.nanmean(yopto_norm, axis=0)

        # compute sem for both conditions
        cont_sem = stats.sem(ycont_norm, axis=0)
        opto_sem = stats.sem(yopto_norm, axis=0)

        # plot curves and sem
        x_eval = np.arange(76)
        ax.plot(x_eval, cont_mean, color='k', linestyle='-')
        ax.fill_between(x_eval, cont_mean - cont_sem, cont_mean + cont_sem, color='k',
                        alpha=0.5, linewidth=0)
        ax.plot(x_eval, opto_mean, color=optocolor, linestyle='-')
        ax.fill_between(x_eval, opto_mean - opto_sem, opto_mean + opto_sem, color=optocolor,
                        alpha=0.5, linewidth=0)

        # layout
        ax.set_ylabel('Normalized firing rate')
        ax.set_xlabel('Diameter ($\degree$)')
        ax.set_xticks((0, 25, 50, 75))
        ax.set_yticks((0, 0.5, 1))
        ax.spines['bottom'].set_bounds(0, 75)
        ax.spines['left'].set_bounds(0, 1)

        f = plt.gcf()
        f.tight_layout()

        return ax

    def scatter(self, figsize=(2.5, 2.5), ax=None, alys=None):
        """Plot scatterplots to compare dLGN spatial integration with V1 intact vs suppressed (
        Fig. 2j-m)

        Parameters
        -------
        figsize: tuple
            Figure size (width, height)
        ax: mpl axis
            axis for plot
        alys: string: 'ropt', 'rsupp', 'rfcs', 'si'
            determines which parameter to analyze: modelled response to optimal stimulus (
            ropt), modelled response to large stimulus (rsupp), preferred size (rfcs),
            suppression index (si)

        Returns
        -------
        ax: mpl axis
            axis for plot
        """

        if ax is None:
            # create figure if ax is none
            f, ax = plt.subplots(figsize=figsize)

        if alys == 'ropt':
            # compare modelled response to optimal stimulus

            # get data
            cont = self.lgn_sztun_df['r_opt_c_76'].values
            supp = self.lgn_sztun_df['r_opt_op_crfcs_76'].values

            # compute statistics
            cont_mean, supp_mean = spatint_utils.compute_stats(cont=cont, supp=supp, alys=alys)

            # set threshold for plotting
            min_thres = 4
            min_dist = 1
            max_dist = 43
            cont_plot = cont
            cont_plot[cont_plot > 30] = max_dist
            cont_plot[cont_plot < min_thres] = min_thres - min_dist
            supp_plot = supp
            supp_plot[supp_plot > 30] = max_dist
            supp_plot[supp_plot < min_thres] = min_thres - min_dist

            # layout
            titlestr = 'Small size\nresponse (sp/s)'
            ax.set_yscale('log')
            ax.set_xscale('log')
            for axis in [ax.xaxis, ax.yaxis]:
                axis.set_major_formatter(ScalarFormatter())
                axis.set_minor_formatter(NullFormatter())
            ax.set_title(titlestr)
            ax.set_xlim(min_thres - (min_dist * 2), max_dist)
            ax.set_ylim(min_thres - (min_dist * 2), max_dist)
            ax.plot((min_thres, 30), (min_thres, 30), linestyle='-', color='grey', zorder=-1,
                    linewidth=0.35)
            ax.spines['left'].set_bounds(min_thres, 30)
            ax.spines['bottom'].set_bounds(min_thres, 30)
            ax.set_xticks((min_thres - min_dist, 10, 30, max_dist))
            ax.set_yticks((min_thres - min_dist, 10, 30, max_dist))
            ax.set_yticklabels(('<' + str(min_thres), 10, 30, '>30'))
            ax.set_xticklabels(('<' + str(min_thres), 10, 30, '   >30'))

        elif alys == 'rsupp':
            # compare modelled response to largest stimulus

            # get data
            cont = self.lgn_sztun_df['r_supp_201'].str[0].values
            supp = self.lgn_sztun_df['r_supp_201'].str[1].values

            # compute statistics
            cont_mean, supp_mean = spatint_utils.compute_stats(cont=cont, supp=supp, alys=alys)

            # set threshold for plotting
            min_thres = 4
            min_dist = 1
            max_dist = 43
            cont_plot = cont
            cont_plot[cont_plot > 30] = max_dist
            cont_plot[cont_plot < min_thres] = min_thres - min_dist
            supp_plot = supp
            supp_plot[supp_plot > 30] = max_dist
            supp_plot[supp_plot < min_thres] = min_thres - min_dist

            # layout
            titlestr = 'Large size\nresponse (sp/s)'
            ax.set_yscale('log')
            ax.set_xscale('log')
            for axis in [ax.xaxis, ax.yaxis]:
                axis.set_major_formatter(ScalarFormatter())
                axis.set_minor_formatter(NullFormatter())
            ax.set_title(titlestr)
            ax.set_xlim(min_thres - (min_dist * 2), max_dist)
            ax.set_ylim(min_thres - (min_dist * 2), max_dist)
            ax.plot((min_thres, 30), (min_thres, 30), linestyle='-', color='grey', zorder=-1,
                    linewidth=0.35)
            ax.spines['left'].set_bounds(min_thres, 30)
            ax.spines['bottom'].set_bounds(min_thres, 30)
            ax.set_xticks((min_thres - min_dist, 10, 30, max_dist))
            ax.set_yticks((min_thres - min_dist, 10, 30, max_dist))
            ax.set_yticklabels(('<' + str(min_thres), 10, 30, '>30'))
            ax.set_xticklabels(('<' + str(min_thres), 10, 30, '   >30'))

        elif alys == 'rfcs':
            # compare preferred size

            # get data
            cont = self.lgn_sztun_df['rfcs_76'].str[0].values
            supp = self.lgn_sztun_df['rfcs_76'].str[1].values

            # compute and calculate stats
            cont_mean, supp_mean = spatint_utils.compute_stats(cont=cont, supp=supp, alys=alys)

            # set threshold for plotting
            cont_plot = cont
            cont_plot[cont_plot > 30] = 35
            supp_plot = supp
            supp_plot[supp_plot > 30] = 35

            # layout
            titlestr = 'Preferred size ($\degree$)'
            ax.set_title(titlestr)
            ax.set_xlim(-1.75, 35)
            ax.set_ylim(-1.75, 35)
            ax.plot((0, 30), (0, 30), linestyle='-', color='grey', zorder=-1, linewidth=0.35)
            ax.spines['left'].set_bounds(0, 30)
            ax.spines['bottom'].set_bounds(0, 30)
            ax.set_xticks((0, 15, 30, 35))
            ax.set_yticks((0, 15, 30, 35))
            ax.set_yticklabels((0, 15, 30, '>30'))
            ax.set_xticklabels((0, 15, 30, '   >30'))

        elif alys == 'si':
            # compare suppression indices

            # get data
            cont_plot = self.lgn_sztun_df['si_76'].str[0].values
            supp_plot = self.lgn_sztun_df['si_76'].str[1].values

            # compute and calculate stats
            cont_mean, supp_mean = spatint_utils.compute_stats(cont=cont_plot, supp=supp_plot,
                                                               alys=alys)

            # layout
            titlestr = 'SI'
            ax.set_title(titlestr)
            ax.set_xlim(-0.05, 1.05)
            ax.set_ylim(-0.05, 1.05)
            ax.set_xticks((0, 0.5, 1))
            ax.set_yticks((0, 0.5, 1))
            ax.plot((0, 1), (0, 1), linestyle='-', color='grey', linewidth=0.35, zorder=-1)
            ax.spines['left'].set_bounds(0, 1)
            ax.spines['bottom'].set_bounds(0, 1)

        else:
            print('No proper analysis selected')
            return

        # general layout
        ax.set_title(titlestr)
        ax.scatter(cont_plot, supp_plot, s=15, facecolors='none', edgecolors='k',
                   linewidth=0.5, clip_on=False)
        ax.plot(cont_mean, supp_mean, linestyle='', marker='.', color='goldenrod', ms=15)
        # plot example
        ax.plot(cont_plot[22], supp_plot[22], linestyle='', marker='.', color='deeppink',
                ms=15)
        ax.set_ylabel('V1 suppression')
        ax.yaxis.label.set_color(optocolor)
        ax.set_xlabel('Control')

        f = plt.gcf()
        f.tight_layout()

        return ax