# code to plot Fig. S4

# import
import matplotlib.pyplot as plt
import scipy.io
import numpy as np
import spatint_utils
import os

# load plot params
spatint_utils.plot_params()


class FigS4:
    """Class for plotting Fig. S4"""

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

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

        # load data from matfile
        data = scipy.io.loadmat(parentdir + '/data/v1_spat_int.mat')

        # unfold data
        self.layerDepth = data['layerDepth'][0]
        self.center_sz_c_smooth = np.array([x[0] for x in data['center_sz_c_smooth']])
        self.si_c_smooth = np.array([x[0] for x in data['si_c_smooth']])
        self.layerY = data['layerY'][0]
        self.max_sz_plot = data['max_sz_plot'][0][0]
        self.rel_layerDepth = data['rel_layerDepth'][0]
        self.si_c = data['si_c'][0]
        self.relative_depth = data['relative_depth'][0]
        self.center_sz_c = data['center_sz_c'][0]

    def rfcs_x_depth(self, ax=None, figsize=(2.5, 2.5)):
        """Plot preferred size against cortical depth in V1 (Fig. S4a)

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

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

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

        # plot
        self.center_sz_c[self.center_sz_c > self.max_sz_plot] = self.max_sz_plot + 20
        ax.scatter(self.center_sz_c, self.relative_depth, color='k', facecolor='none')
        ax.plot(self.center_sz_c_smooth, self.relative_depth, color='r')
        for layer in self.layerY:
            ax.hlines(layer, 0, self.max_sz_plot, color='k', linestyle='--')

        # layout
        ax.set_xlabel('Preferred size ($\degree$)')
        ax.set_ylabel('Cortical depth to\nbase of L4 (mm)')
        ax.set_xlim(-2, self.max_sz_plot+25)
        ax.set_ylim([(self.rel_layerDepth[0] - self.layerDepth[0]) / 1000, self.rel_layerDepth[
            -1] / 1000])
        ax.invert_yaxis()
        ax.set_yticks(())
        ax.set_xticks((0, 25, 50, 75, 95))
        ax.set_xticklabels((0, 25, 50, 75, '>75'))
        ax.set_yticks((-0.4, 0, 0.4))
        ax.set_yticklabels(())
        ax.spines['bottom'].set_bounds(0, self.max_sz_plot)
        f = plt.gcf()
        f.tight_layout()

        # print n
        print('n = %d' % len(self.relative_depth))

        return ax

    def si_x_depth(self, ax=None, figsize=(2.5, 2.5)):
        """Plot suppression index against cortical depth in V1 (Fig. S4b)

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

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

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

        # plot si vs depth
        ax.scatter(self.si_c, self.relative_depth, color='k', s=15, facecolor='none')
        ax.plot(self.si_c_smooth, self.relative_depth, color='r')
        for layer in self.layerY:
            ax.hlines(layer, 0, 1, color='k', linestyle='--')

        # layout
        ax.set_xlabel('Suppression index (SI)')
        ax.set_ylabel('Cortical depth to\nbase of L4 (mm)')
        ax.set_xlim(-0.05, 1.05)
        ax.set_ylim([(self.rel_layerDepth[0] - self.layerDepth[0]) / 1000, self.rel_layerDepth[
            -1] / 1000])
        ax.invert_yaxis()
        ax.set_yticks((-0.4, 0, 0.4))
        ax.set_xticks((0, 0.5, 1))
        ax.spines['bottom'].set_bounds(0, 1)
        f = plt.gcf()
        f.tight_layout()

        return ax