import matplotlib.pyplot as plt
import numpy as np
import pandas as pd
from c_inv.inference_paper_submission.layout import Frame, Figure
from c_inv.inference_paper_submission.panels import oricon_matrix, marginal_tuning, contrast_amp_distr
from c_inv.inference_paper_submission.util import (cm2inch, c_inv_paper_series, JNEUROSCIPARS, set_mpl_pars,
                   align_y_axis_labels, inhibcol, excitcol)
from pathlib import Path

set_mpl_pars()
plt.rcParams['font.sans-serif'] = ['Arial']

layoutstring = 'ABC'
layout = Frame(layoutstring, left=0.01, top=0.07, right=0.017, bot=0.15, h_space=0.03,
               w_space=0.015, panelpars={'A': {'left': 0.12, 'right': 0.09, 'width_ratio': 1.3,
                                               'h_space': 0.05, 'w_space': 0.05},
                                         'B': {'left': 0.26, 'h_space': 0.05, 'w_space': 0.02,
                                               'width_ratio': 0.95},
                                         'C': {'left': 0.2, 'width_ratio': 0.95}})
layout.panels['A'].new_panels('12\n34')
layout.panels['B'].new_panels('14\n23', panelpars={'4': {'invis': True},
                                                   '1': {'width_ratio': 2},
                                                   '2': {'width_ratio': 2,
                                                         'height_ratio': 3},
                                                   '3': {'height_ratio': 3}})
figure = Figure(layout, cm2inch(JNEUROSCIPARS['doublecolumn']), 2.2)
figure.annotate(fontweight='bold')


keys = [{'m': 'Ntsr1Cre_2020_0001', 's': 4, 'e': 9, 'u': 17},
        {'m': 'Ntsr1Cre_2020_0001', 's': 4, 'e': 9, 'u': 34},
        {'m': 'Ntsr1Cre_2020_0001', 's': 3, 'e': 9, 'u': 14},
        {'m': 'PVCre_2019_0009', 's': 4, 'e': 5, 'u': 60},
        {'m': 'PVCre_2019_0009', 's': 3, 'e': 6, 'u': 22},
        {'m': 'PVCre_2019_0009', 's': 3, 'e': 6, 'u': 34}]

###########
# load data
###########
pkg_dir = Path(__file__).parent
data_dir = pkg_dir.joinpath('Data')
oricon1 = pd.read_pickle(data_dir.joinpath('S2_oricon0.pkl'))
oricon2 = pd.read_pickle(data_dir.joinpath('S2_oricon1.pkl'))
oricon3 = pd.read_pickle(data_dir.joinpath('S2_oricon2.pkl'))
oricon4 = pd.read_pickle(data_dir.joinpath('S2_oricon3.pkl'))
marg_data = pd.read_pickle(data_dir.joinpath('S2_marginals_data.pkl'))
marg_pars = np.load(str(data_dir.joinpath('S2_marginals_pars.npy')), allow_pickle=True)
exc_amps = pd.read_pickle(data_dir.joinpath('S2_con_amps_exc.pkl'))
inh_amps = pd.read_pickle(data_dir.joinpath('S2_con_amps_inh.pkl'))


ax1 = figure.axes['A1']
_, cbar1 = oricon_matrix(ax1, oricon1, cbar=True)
cbar1.set_label('Firing (Hz)', labelpad=5)
ax2 = figure.axes['A2']
_, cbar2 = oricon_matrix(ax2, oricon2, cbar=True)
cbar2.set_label('Firing (Hz)', labelpad=4)
ax3 = figure.axes['A3']
_, cbar3 = oricon_matrix(ax3, oricon3, cbar=True)
cbar3.set_label('Firing (Hz)', labelpad=5)
cbar3.set_ticks([5, 10])
ax4 = figure.axes['A4']
_, cbar4 = oricon_matrix(ax4, oricon4, cbar=True)
cbar4.set_label('Firing (Hz)', labelpad=7.5)
for ax in [ax1, ax2, ax3, ax4]:
    ax.set_yticks([0, 6, 12])
    ax.set_yticklabels([0, 90, 180])
    ax.set_xticks([0, 4, 7])
    ax.set_xticklabels([0, 0.3, 1])
ax1.set_xlabel('')
ax1.set_xticklabels([])
ax2.set_xlabel('')
ax2.set_ylabel('')
ax2.set_xticklabels([])
ax2.set_yticklabels([])
ax4.set_ylabel('')
ax4.set_yticklabels([])

ax1 = figure.axes['B1']
ax2 = figure.axes['B2']
ax3 = figure.axes['B3']
plt.rcParams['lines.markersize'] = 3
marginal_tuning(marg_data, marg_pars, [ax1, ax2, ax3], tuning_model='wrap_gauss')
ax3.set_xlabel('Firing (Hz)')
ax3.set_xlim([12, 22])
ax1.set_ylabel('Firing (Hz)')
ax2.set_yticks([0, 3, 6, 9, 12])
ax2.set_yticklabels([0, 45, 90, 135, 180])
ax2.set_xticks([0, 4, 7])
ax2.set_xticklabels([0, 0.3, 1])
align_y_axis_labels(ax2, ax1)

ax4 = figure.axes['C']
contrast_amp_distr(exc_amps, inh_amps, ax4, excit_col=excitcol, inhib_col=inhibcol)