"""
Figure 5 S1
"""

import os
import sys
import pickle as pkl
import numpy as np
from matplotlib import pyplot as plt
from matplotlib.ticker import MaxNLocator
import helper
import seaborn as sns

assert sys.version_info < (3, 0), "Please run this script using Python2.7"

def plot_stats(data_, filename):
    """

    :param data_: pickled data from simulation
    :param filename: where to save
    :return:
    """
    modularity = np.arange(0.75, 1.01, 0.05)
    populations = ['E{}'.format(x) for x in range(6)]
    stats = ['ccs_pearson', 'cvs', 'mean_rates']
    labels = ['Pearson CC', 'CV', 'spks/sec']
    # plotting params
    helper.usetex_font()
    color_palette = sns.cubehelix_palette(20, start=.5, rot=-.75)[::2]

    # figure params
    fig_ratio = 6 / 10.  # raster plot ratio for 1000 Neurons
    fig_width = 8
    labelsize = 24
    ticksize = 18
    legendsize = 16
    lw = 3.

    figure = plt.figure(figsize=(15, 4))  # this is to match the raster plot size

    for stat_idx, stat in enumerate(stats):
        ax = plt.subplot2grid((1, 3), (0, stat_idx))

        for pop_idx, pop in enumerate(populations):
            mean = data_[pop][stat]['mean'][:, 0]
            color = color_palette[pop_idx]
            ax.plot(modularity, mean, '-', color=color, label=r'$\mathrm{SSN}_' + str(pop_idx) + '$', lw=lw)

        # lim = {
        #     'ccs_pearson': (0., 0.6) if not supplementary else (0., 0.8),
        #     'mean_rates': (1., 16.),
        #     'corrected_rates': (1., 16.)
        # }
        # ax.set_ylim(*lim[stat])
        ax.yaxis.set_major_locator(MaxNLocator(5))
        # ax.xaxis.set_major_locator(MaxNLocator(6))
        ax.set_ylabel(labels[stat_idx], fontsize=labelsize)
        ax.set_xlabel('m', fontsize=labelsize)
        ax.set_xticks([0.75, 0.85, 1.])

        ax.grid(False)
        ax.spines['right'].set_visible(False)
        ax.spines['top'].set_visible(False)
        ax.tick_params(axis='both', labelsize=ticksize, direction='out')

        if stat_idx == 0:
            plt.legend(loc='best', prop={'size': legendsize}, handleheight=1.2, handlelength=2.6, frameon=False)

    figure.tight_layout()
    figure.subplots_adjust(wspace=0.5)
    figure.savefig(os.path.join('plots/', filename))
    print("Finished plotting Figure 5S1.")


if __name__ == "__main__":
    with open('data/fig_5s1_conductance_stats.pkl', 'r') as f:
        data = pkl.load(f)

    plot_stats(data, 'fig5_s1.pdf')
