# code to plot figure S6

# import
import matplotlib.pyplot as plt
from importlib import reload
import pandas
import spatint_utils
import os

# reload modules
reload(spatint_utils)

# load plot params
spatint_utils.plot_params()


class FigS6:
    """Class for plotting panels of Fig. S6"""

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

        parentdir = os. path. dirname(os. getcwd())

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

    def ex_curves(self, figsize=(5, 2), axs=None):
        """Plot example dLGN size tuning curves (Fig. S6 a-j)

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

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

        if axs is None:
            # init figure
            f, axs = plt.subplots(2, 5, figsize=figsize)

        # determine example keys
        exkeys = [{'m': 'PVCre_2013_0046', 's': 8, 'e': 9, 'u': 3038},   # a
                  {'m': 'PVCre_2013_0046', 's': 6, 'e': 6, 'u': 4042},   # b
                  {'m': 'PVCre_2013_0046', 's': 8, 'e': 12, 'u': 3030},  # c
                  {'m': 'PVCre_2013_0054', 's': 10, 'e': 4, 'u': 2073},  # d
                  {'m': 'PVCre_2013_0046', 's': 8, 'e': 14, 'u': 3039},  # e
                  {'m': 'PVCre_2013_0054', 's': 10, 'e': 5, 'u': 1048},  # f
                  {'m': 'PVCre_2013_0046', 's': 5, 'e': 7, 'u': 3037},   # g
                  {'m': 'PVCre_2013_0054', 's': 11, 'e': 4, 'u': 3016},  # h
                  {'m': 'PVCre_2020_0002', 's': 6, 'e': 5, 'u': 4},      # i
                  {'m': 'PVCre_2013_0046', 's': 6, 'e': 6, 'u': 4040}]   # j

        # loop over example keys and plot
        for ax, exkey in zip(axs.flatten(), exkeys):

            ex_row = self.lgn_sztun_df[(self.lgn_sztun_df.m == exkey['m']) &
                                       (self.lgn_sztun_df.s == exkey['s']) &
                                       (self.lgn_sztun_df.e == exkey['e']) &
                                       (self.lgn_sztun_df.u == exkey['u'])].iloc[0]

            spatint_utils.plot_tun(means=ex_row['tun_mean'],
                                   sems=ex_row['tun_sem'],
                                   spons=ex_row['tun_spon_mean'],
                                   xs=ex_row['ti_axes'],
                                   c_fit=ex_row['c_sz_fit'],
                                   op_fit=ex_row['op_sz_fit'],
                                   c_prefsz=ex_row['rfcs_76'][0],
                                   op_prefsz=ex_row['rfcs_76'][1],
                                   ax=ax)

            # layout
            max_x = 76
            ax.set_xlim(-3, max_x)
            ax.set_xticks([0, 25, 50, 75])
            ax.spines['bottom'].set_bounds(0, max_x)
            yticks = ax.get_yticks()
            ax.set_yticks((yticks[0], int((yticks[-1] - 5) / 2), yticks[-1] - 5))
            ax.set_ylabel('')
            ax.set_xlabel('')
            ax.set_xticklabels([])

        axs.flatten()[5].set_ylabel('Firing rate (sp/s)')
        axs.flatten()[5].set_xlabel('Diameter ($\degree$)')#
        axs.flatten()[5].set_xticklabels([0, 25, 50, 75])
        f = plt.gcf()
        f.tight_layout()

        return axs