# code to plot modelled data in figure 3

# import libs
import numpy as np
import matplotlib.pyplot as plt
import matplotlib.patches as patch
import scipy.stats as stats
from importlib import reload
import spatint_utils
import sim_edog

# reload modules
reload(spatint_utils)

# load params
spatint_utils.plot_params()
_, _, optocolor = spatint_utils.get_colors()
cmpin = 2.54


class Fig3:
    """Class for plotting Figure 3"""

    def __init__(self, figsize=np.array((18.3, 8)), edog=None, inhib_fb=(1, 3, 9, 40)):
        """Init class

        Parameters
        -----
        figsize: tuple len 2
            figuresize (width, heigth)
        edog: class object
            edog with simulated data
        inhib_fb: tuple
            example inhibitory feedback kernel widths to be plotted
        """

        # store init parameters
        self.figsize = figsize / cmpin                     # figure size
        self.sids = edog.sids                              # change in suppression indices
        self.rfcsds = edog.rfcsds                          # change in preferred size
        self.smallrateds = edog.smallrateds                # change in response to small stim
        self.largerateds = edog.largerateds                # change in response to large stim
        self.lowbounds = edog.lowbounds                    # dict with low bounds for matching
        self.highbounds = edog.highbounds                  # dict with high bounds for matching
        self.inhib_fb = inhib_fb                           # ex inhibitory fb kernel widths
        self.start = edog.start                            # first inhib fb kernel width
        self.stop = edog.stop                              # last inhib fb kernel width
        self.inhib_fbi = np.array(inhib_fb) - edog.start   # index for example kernels

    def plot(self):
        """Plot figure"""

        # init figure
        f = plt.figure(figsize=self.figsize)

        # init plotting variables
        interplotspace = 0.5 / cmpin / self.figsize[0]
        axdict = {}
        axdict['b'] = 3 / cmpin / self.figsize[1]
        axdict['h'] = 1.5 / cmpin / self.figsize[1]
        axdict['w'] = 1.5 / cmpin / self.figsize[0]
        l = 5 / cmpin / self.figsize[0]

        # plot modelled size tuning curves
        inhib_fbs = self.inhib_fb
        inhib_fbi = self.inhib_fbi
        start = self.start
        stop = self.stop
        edog_model = sim_edog.SimEdog()
        sb_label = False

        # get color for inhib kernels
        cmap = plt.cm.get_cmap('Reds_r')
        coli = np.linspace(0, 0.5, 4)

        # loop over example inhibitory width
        for i, inhib_fb in enumerate(inhib_fbs):

            # add axis
            ax = f.add_axes([l, axdict['b'], axdict['w'], axdict['h']])
            # add curves
            edog_model.plot(ax=ax, inhib_fb=inhib_fb)
            # add kernel schemas
            self._plot_kernel(inhib_fb=inhib_fb, l=l, f=f, col=cmap(coli[i]), i=i,
                              sb_label=sb_label)

            # label according to panel position
            if i > 0:
                # remove labels
                ax.set_ylabel('')
                ax.set_yticklabels('')
                ax.set_xlabel('')
                ax.set_xticklabels('')

            if i < (len(inhib_fbs) - 1):
                # add length
                l += interplotspace + axdict['w']

            if i == (len(inhib_fbs) - 2):
                # add sb_label
                sb_label = True

            elif i == (len(inhib_fbs) - 1):
                ax.text(10, 0.35, 'FB weight = 1', color='k')
                ax.text(10, 0.15, 'FB weight = 0', color=optocolor)

        # get x values
        x = np.arange(start, stop, 1)

        # define markersize
        ms = 8

        # add change in receptive field center size
        # prepare axis
        interplotspace = 1.5 / cmpin / self.figsize[0]
        axdict['h'] = 1.2 / cmpin / self.figsize[1]
        axdict['w'] = 1.2 / cmpin / self.figsize[0]
        l += axdict['w'] + interplotspace + 0.02
        ax = f.add_axes([l, axdict['b'], axdict['w'], axdict['h']])

        # add plot
        ax.plot(x, self.rfcsds, c='k')

        # edit layout
        xlims = ax.get_xlim()
        ax.hlines(0, xlims[0], xlims[1], colors='grey', linestyle='--', linewidth=0.5)
        ax.set_xticks((1, 20, 40))
        ax.set_yticks((0, -5, -10))
        ax.spines['bottom'].set_bounds(1, 40)
        ax.spines['left'].set_bounds(-10, 0)
        ax.set_ylabel('$\Delta$ preferred\nsize (%)')
        ax.set_xlabel('Inh FB kernel width ($\degree$)')

        # add quality rectangles
        ylims = ax.get_ylim()
        bottom = ylims[0]
        height = ylims[1] - bottom
        width_qual = self.highbounds['all'] - self.lowbounds['all']

        # overall
        rect_all = patch.Rectangle((self.lowbounds['all'], bottom), width_qual, height,
                                   facecolor='gold', alpha=0.7, zorder=0)
        ax.add_patch(rect_all)
        # preferred size specific
        rfcs_width = self.highbounds['rfcsds'] - self.lowbounds['rfcsds']
        rect_rfcs = patch.Rectangle((self.lowbounds['rfcsds'], bottom), rfcs_width, height,
                                    facecolor='gold', alpha=0.2, zorder=0)
        ax.add_patch(rect_rfcs)

        # add example points
        ax.scatter(inhib_fbs, np.array(self.rfcsds)[inhib_fbi], c=np.array(cmap(coli)), s=ms,
                   zorder=3)

        # add changes in suppression index
        # prepare axis
        l += axdict['w'] + interplotspace
        ax = f.add_axes([l, axdict['b'], axdict['w'], axdict['h']])

        # add plot
        ax.plot(x, self.sids, c='k')

        # edit layout
        xlims = ax.get_xlim()
        ax.hlines(0, xlims[0], xlims[1], colors='grey', linestyle='--', linewidth=0.5)
        ax.set_xticks((1, 20, 40))
        ax.set_yticks((0, 100, 200))
        ax.spines['bottom'].set_bounds(1, 40)
        ax.spines['left'].set_bounds(0, 200)
        ax.set_ylabel('$\Delta$ suppression\nindex (%)')

        # add quality rectangle
        ylims = ax.get_ylim()
        bottom = ylims[0]
        height = ylims[1] - bottom

        # overall
        rect_all = patch.Rectangle((self.lowbounds['all'], bottom), width_qual, height,
                                   facecolor='gold', alpha=0.7, zorder=0)
        ax.add_patch(rect_all)
        # si specific
        si_width = self.highbounds['sids'] - self.lowbounds['sids']
        rect_si = patch.Rectangle((self.lowbounds['sids'], bottom), si_width, height,
                                  facecolor='gold', alpha=0.2, zorder=0)
        ax.add_patch(rect_si)

        # add example points
        ax.scatter(inhib_fbs, np.array(self.sids)[inhib_fbi], c=np.array(cmap(coli)), s=ms,
                   zorder=3)

        # add changes in smallrates
        # prepare axis
        axdict['b'] += axdict['h'] + (1.5 / cmpin / self.figsize[0])
        l -= axdict['w'] + interplotspace
        ax = f.add_axes([l, axdict['b'], axdict['w'], axdict['h']])

        # add plot
        ax.plot(x, self.smallrateds, c='k')

        # layout
        xlims = ax.get_xlim()
        ax.hlines(0, xlims[0], xlims[1], colors='grey', linestyle='--', linewidth=0.5)
        ax.set_xticks((1, 20, 40))
        ax.set_xticklabels('')
        ax.set_yticks((-25, 0, 25))
        ax.spines['bottom'].set_bounds(1, 40)
        ylims = ax.get_ylim()
        ax.spines['left'].set_bounds(-25, ylims[1])
        ax.set_ylabel('$\Delta$ response to\npreferred size (%)')

        # add quality rectangle
        ylims = ax.get_ylim()
        bottom = ylims[0]
        height = ylims[1] - bottom
        # overall
        rect_all = patch.Rectangle((self.lowbounds['all'], bottom), width_qual, height,
                                   facecolor='gold', alpha=0.7, zorder=0)
        ax.add_patch(rect_all)
        # smallrate specific
        smallrate_width = self.highbounds['smallrateds'] - self.lowbounds['smallrateds']
        rect_smallrate = patch.Rectangle((self.lowbounds['smallrateds'], bottom),
                                         smallrate_width, height, facecolor='gold',
                                         alpha=0.2, zorder=0)
        ax.add_patch(rect_smallrate)
        # add example points
        ax.scatter(inhib_fbs, np.array(self.smallrateds)[inhib_fbi], c=np.array(cmap(coli)),
                   s=ms, zorder=3)

        # add changes in large rates
        # add axis
        l += axdict['w'] + interplotspace
        ax = f.add_axes([l, axdict['b'], axdict['w'], axdict['h']])

        # add plot
        ax.plot(x, self.largerateds, c='k')

        # edit layout
        xlims = ax.get_xlim()
        ax.hlines(0, xlims[0], xlims[1], colors='grey', linestyle='--', linewidth=0.5)
        ax.set_xticks((1, 20, 40))
        ax.set_xticklabels('')
        ax.set_yticks((-20, 0))
        ax.spines['bottom'].set_bounds(1, 40)
        ylims = ax.get_ylim()
        ax.spines['left'].set_bounds(ylims)
        ax.set_ylabel('$\Delta$ response to\n largest size (%)')

        # add quality rectangle
        ylims = ax.get_ylim()
        bottom = ylims[0]
        height = ylims[1] - bottom
        # overall
        rect_all = patch.Rectangle((self.lowbounds['all'], bottom), width_qual, height,
                                   facecolor='gold', alpha=0.7, zorder=0)
        ax.add_patch(rect_all)
        # largerate specific
        largerate_width = self.highbounds['largerateds'] - self.lowbounds['largerateds']
        rect_largerate = patch.Rectangle((self.lowbounds['largerateds'], bottom),
                                         largerate_width, height, facecolor='gold', alpha=0.2,
                                         zorder=0)
        ax.add_patch(rect_largerate)
        # add example points
        ax.scatter(inhib_fbs, np.array(self.largerateds)[inhib_fbi], c=np.array(cmap(coli)),
                   s=ms, zorder=3)

    def _plot_kernel(self, inhib_fb, l, f, col, i, sb_label=True):
        """plot schematic kernels

        Parameters
        -----
        inhib_fb: int
            widht of inhibitory feedback kernel
        l: float
            left edge of plot
        f: mpl figure
            figure
        col: tuple
            color for inhibitory kernel
        i: int
            panel index
        sb_label: bool
            if true plot scalebar
        """

        # define axis for inhibitory gaussian
        b = 5.65 / cmpin / self.figsize[1]
        h = 0.2 / cmpin / self.figsize[1]
        w = 2 / cmpin / self.figsize[0]
        l -= 0.2 / cmpin / self.figsize[0]
        zz = 0.1
        xlims = (-40, 40)

        # compute inhibtory gaussian
        x = np.arange(-30, 31, 1)
        mu = 0
        sigma_c = inhib_fb
        amp_c = -0.6
        normdist = stats.norm.pdf(x, mu, sigma_c)
        y_inhib = normdist * amp_c

        # plot inhibitory gaussian
        ax = f.add_axes([l, b, w, h])
        ax.plot(x,y_inhib, color=col, clip_on=False)
        ax.set_ylim((0, zz))
        ax.set_xlim(xlims)
        plt.axis('off')

        # add x axis label
        if i == 0:
            ax.text(-45, -0.1, 'inh', rotation='vertical', c='r')

        # define axis for excitatory gaussian
        b = 6 / cmpin / self.figsize[1]

        # compute excitatory gaussian
        mu = 0
        sigma_c = 1
        amp_c = 0.3
        normdist = stats.norm.pdf(x, mu, sigma_c)
        y_ex = normdist * amp_c

        # plot excitatory gaussian
        ax = f.add_axes([l, b, w, h])
        ax.plot(x, y_ex, color='g', clip_on=False)
        ax.set_ylim((0, zz))
        ax.set_xlim(xlims)
        plt.axis('off')

        # add x axis label
        if i == 0:
            ax.text(-45, 0, 'exc', rotation='vertical', c='g')

        # add scalebar
        x1 = 20
        x2 = 30
        xdiff = np.abs(x2 - x1)
        ax.plot((x1, x2), (0.13, 0.13), color='k', clip_on=False)

        if sb_label:
            scalebar = '%d$\degree$' % xdiff
            ax.text(x1, 0.18, scalebar, color='k', clip_on=False)

        # define axis for sum of gaussians
        b = 5 / cmpin / self.figsize[1]

        # compute sum of two gaussians
        y_sum = y_ex + y_inhib

        # plot sum of gaussians
        ax = f.add_axes([l, b, w, h])
        ax.plot(x, y_sum, color='k', clip_on=False)
        ax.set_ylim((0, zz))
        ax.set_xlim(xlims)
        plt.axis('off')

        # add x axis label
        if i == 0:
            ax.text(-45, 0, 'sum', rotation='vertical')
            ax.text(-70, 0, 'CT feedback\n kernel ($K_{RC}$)', rotation='vertical')