"""Fig 1S1 plots"""

# import
from djd.signal import rfs_for_cdata, optosupp_for_cdata
from djd.plot import simple_plot_rfs, simple_plot_optosupp
import matplotlib.pyplot as plt


# specify and update matplotlib parameters
fontsize = 6.4
params = {
            'font.family': 'FreeSans',
            'font.weight': 'normal',
            'font.size': fontsize,
            'pdf.fonttype': 42,
            'ps.fonttype': 42,
            'xtick.labelsize': fontsize,
            'ytick.labelsize': fontsize,
            'lines.linewidth' :0.734,
            'xtick.major.size' : 1.734,
            'ytick.major.size' : 1.734,
            'axes.linewidth' : 0.4,
            'axes.titlesize': fontsize,
            'ytick.major.width' :0.4,
            'xtick.major.width' :0.4,
            'savefig.dpi': 300,
            }

plt.rcParams.update(params)

# optosuppression panel

# define key and target channels
v1key = {'m': 'PVCre_2018_0003', 's': 1, 'e': 1}
exchanis = [24, 25, 26]

# get opto suppression data
ondata, offdata, t, stim_perc_red, light_diff_onset, light_dur, stim_dur = optosupp_for_cdata(
                                                                            key=v1key,
                                                                            pad=(-0.2, 1))

# init parameters
cmpin = 2.54
figwidth = 11.4 / cmpin
figheight = 15 / cmpin
bardistance = 1.2 # distance bar to axes
barwidth = 3
chanh = 0.15 # subplot height
interchanh = 0.038 # vertical space between subplots
figsize = (figwidth, figheight)
f = plt.figure(figsize=figsize)

# define first axis
l1 = 0.07
b1 = 0.61
w1 = 0.6
h = chanh

# create axis for each channel
axlist = []
for _ in exchanis:
    b1 = b1 - chanh
    ax = f.add_axes([l1, b1, w1, h])
    b1 = b1 - interchanh
    axlist.append(ax)

# fill axes with data
axs = simple_plot_optosupp(key=v1key, ondata=ondata, offdata=offdata, t=t, perc_red=stim_perc_red,
                           light_diff_onset=light_diff_onset, light_dur=light_dur,
                           stim_dur=stim_dur, barwidth=barwidth, bardistance=bardistance,
                           axs=axlist, chanis=exchanis)

# format figure
axs[1].set_ylabel('Normalized multi unit activity')
axs[-1].set_ylabel('')
axs[-1].set_xticklabels((0,1,2,3))

# rf panel

# define keys for the two mua rfs
ekeys = [{'m': 'PVCre_2017_0015', 's': 3, 'e':1},
         {'m': 'PVCre_2017_0015', 's': 7, 'e':1}]

# define analysis window
rf_offset = (0.05, 0.1)

# init parameters
chanh = 0.04 # subplot height
interchanh = 0.025 # vertical space between subplots
w = 0.18
h = chanh
start_chans = [7,14] # indices of channels to start from (one for each key)
nchans = 8 # number of channels to plot
sep = 2 # step: plot every second channel

# create one axis per channel
axlists=[]
for k in range(2):
    axlist=[]
    l = l1 + w1 + 0.1 + (0.1 * k)
    b2 = 0.61
    for i in range(nchans):
        b2 = b2 - chanh
        ax = f.add_axes([l, b2, w, h])
        b2 = b2 - interchanh
        axlist.append(ax)
    axlists.append(axlist)

# fill axes with data
for ekey, axlist, start_chan in zip(ekeys, axlists, start_chans):
    rfs, axes, chorder = rfs_for_cdata(ekey, offset=rf_offset)
    simple_plot_rfs(ekey, rfs[start_chan:start_chan+(nchans*sep):sep,:,:,:],
                    axes, interpolation='spline16', nrows=nchans,
                    chorder=chorder[start_chan:start_chan+(nchans*sep)],
                    axs=axlist, contrasts=2)

    for ax in axlist:
        ax.set_title((''))

# format figure
axlists[0][-1].yaxis.tick_left()
axlists[0][-1].set_ylabel('Elevation')
axlists[0][-1].set_xlabel('Azimuth')
axlists[1][-1].set_yticklabels((''))
axlists[0][0].set_title('Exp1')
axlists[1][0].set_title('Exp2')