"""Figure 1S6 pupil area plots, use run -i fig1S6.py"""

## Redo scatter plots from fig1.py, but only for movie experiments with insignificant
## effect of opto on pupil area, as determined by LMM from trial-wise pupil data:
assert EXPTYPE == 'pvmvis'
insigmsefname = os.path.join('stats', 'figure_1_S6_cde_insig_pupil_diam_mseus.csv')
df = pd.read_csv(insigmsefname)
insigmsestrs = list(df['mseustr'].values)
# filter out units from movie experiments with significant opto effects on pupil area:
insigmseustrs = []
for mseustr in mvimseustrs:
    mseustrlist = mseustr.split('_')
    assert len(mseustrlist) == 6
    msestr = '_'.join(mseustrlist[:5]) # drop the unit ID
    if msestr in insigmsestrs:
        insigmseustrs.append(mseustr)
mvimseustrs = insigmseustrs # overwrite
print("## WARNING: make sure not to save mvimseustrs to disk as a pickle. "
      "Its value is overwritten here as a hack")
## or instead of running this, which generates way too many plots, copy and paste only the
## meanrate, BR, spars and rel scatter plot code blocks into IPython:
print("Calling fig1.py to generate movie plots")
get_ipython().run_line_magic('run', '-i fig1.py')
plt.show()
print("Now save the first 4 scatter plot figures, but don't save any dataframes")


## scatter plot pupil area during feedback vs suppression, one point (mean over trials) per mse:
for stimtype, stimtypei in zip(STIMTYPESPLUSALL, STIMTYPESPLUSALLI):
    linmin, linmax = 0, 0.6
    ticks = 0, 0.5
    ctrlareas = ipos_opto['mean_area'][:, stimtypei, False].values
    optoareas = ipos_opto['mean_area'][:, stimtypei, True].values
    f, a = plt.subplots(figsize=DEFAULTFIGURESIZE)
    wintitle('pupil %s area opto scatter' % stimtype)
    # plot y=x line:
    xyline = [linmin, linmax], [linmin, linmax]
    a.plot(xyline[0], xyline[1], '--', color='gray', zorder=-1)
    # plot points:
    a.scatter(optoareas, ctrlareas, clip_on=False, marker='.',
              c='None', edgecolor='black', s=DEFSZ)
    a.set_xlabel('Suppression area (AU)')
    a.set_ylabel('Feedback area (AU)')
    a.set_xlim(linmin, linmax)
    a.set_ylim(linmin, linmax)
    a.set_xticks(ticks)
    a.set_yticks(ticks)
    a.set_aspect('equal')
    a.spines['left'].set_position(('outward', 4))
    a.spines['bottom'].set_position(('outward', 4))
    t, p = ttest_rel(optoareas, ctrlareas) # paired t-test
    a.add_artist(AnchoredText('p$=$%.2g' % p, loc='upper left', frameon=False))


## plot CDFs of pupil area during feedback vs suppression, use trial level precision, one
## plot per mse:
msestrs = np.unique(ipos_opto.reset_index()['mse'].values)
for msestr in ['PVCre_2017_0006_s03_e03', 'PVCre_2017_0015_s07_e05']:#msestrs:
    linmin, linmax = 0, 0.6
    ticks = 0, 0.5
    ctrlareas = ipos_opto['area_trialmean'][msestr, :, False]
    optoareas = ipos_opto['area_trialmean'][msestr, :, True]
    assert len(ctrlareas) == 1
    assert len(optoareas) == 1
    ctrlareas, optoareas = ctrlareas[0], optoareas[0] # pull out of series
    f, a = plt.subplots(figsize=DEFAULTFIGURESIZE)
    wintitle('pupil area opto CDF %s' % msestr)
    a.hist(ctrlareas, bins=np.unique(ctrlareas), density=True, cumulative=True, histtype='step',
           color='k', label='V1 Control')
    a.hist(optoareas, bins=np.unique(optoareas), density=True, cumulative=True, histtype='step',
           color=optoblue, label='V1 Suppressed')
    a.set_xlabel('Pupil Area (AU)')
    a.set_ylabel('Cumulative probability')
    a.set_yticks([0, 0.5, 1])
    #a.spines['left'].set_position(('outward', 4))
    #a.spines['bottom'].set_position(('outward', 4))
    _, ks_p = ks_2samp(optoareas, ctrlareas)
    a.add_artist(AnchoredText('p$=$%.2g' % ks_p, loc='upper left', frameon=False))
    #l = a.legend(frameon=False)
    #l.set_draggable(True)


## scatter plot firing rate FMI and burst ratio FMI (one value per mseu, y axis)
## vs pupil area FMI (one value per mse, x axis)
stimtype2FMI = {'mvi':mviFMI, 'grt':grtFMI, 'mvi+grt':pd.concat([mviFMI, grtFMI])}
for stimtype, stimtypei in zip(STIMTYPESPLUSALL, STIMTYPESPLUSALLI):
    # first calculate pupil area FMI for each mse:
    ipos_opto_stimtype = ipos_opto.loc[(slice(None), stimtypei, slice(None))]
    umses = np.unique(ipos_opto_stimtype.reset_index()['mse']) # unique mses for this stimtype
    mse2areafmi = {}
    for umse in umses:
        row = ipos_opto_stimtype.loc[umse]['mean_area'] # indexed by opto
        assert len(row) == 2
        if stimtype == 'mvi+grt':
            row = row.droplevel('stimtype') # drop unitary (redundant) index level
        feedback, suppress = row[False], row[True]
        areafmi = (feedback - suppress) / (feedback + suppress)
        mse2areafmi[umse] = areafmi
    # now collect neural FMIs for each mseu:
    FMI = stimtype2FMI[stimtype]
    st8 = 'none'
    mseus = np.unique(FMI.reset_index()['mseu'].values)
    rfmis = np.float64(FMI.loc[mseus]['meanrate'][:, st8].values)
    brfmis = np.float64(FMI.loc[mseus]['meanburstratio'][:, st8].values)
    # get mse string of every mseu in FMI:
    mses = np.array(['_'.join(mseu.split('_')[:-1]) for mseu in mseus])
    # assign appropriate pupil area FMI to each mseu:
    areafmis = np.tile(np.nan, len(mses)) # init to nans
    for msei, mse in enumerate(mses):
        if mse in mse2areafmi:
            areafmis[msei] = mse2areafmi[mse]
    ## scatter plot meanrate FMI vs pupil area FMI:
    figsize = DEFAULTFIGURESIZE[0]*1.05, DEFAULTFIGURESIZE[1] # tweak to make space for labels
    f, a = plt.subplots(figsize=figsize)
    wintitle('meanrate %s FMI vs pupil area FMI' % stimtype)
    # boolean indicating entries with non-NaN meanrate FMI & pupil area FMI:
    valid = ~np.isnan(areafmis) & ~np.isnan(rfmis)
    a.scatter(areafmis[valid], rfmis[valid], clip_on=False,
              marker='.', c='None', edgecolor='black', s=DEFSZ)
    # get fname of appropriate LMM .cvs file:
    fname = None
    if stimtype == 'mvi+grt':
        if EXPTYPE == 'pvmvis':
            fname = os.path.join('stats', 'figure_1_S6g_coefs.csv')
        elif EXPTYPE == 'ntsrmvis':
            fname = os.path.join('stats', 'figure_1_S6h_coefs.csv')
    if fname:
        try:
            df = pd.read_csv(fname)
            foundregression = True
        except FileNotFoundError:
            print('Missing file: %s' % fname)
            foundregression = False
        if foundregression:
            mm = df['slope'][0]
            b = df['intercept'][0]
            x = np.array([areafmis[valid].min(), areafmis[valid].max()])
            y = mm * x + b
            a.plot(x, y, '-', color='red') # plot linregress fit
    a.set_ylabel("Mean rate FMI")
    a.set_xlabel("Pupil area FMI")
    xmin, xmax = -0.15, 0.15
    ymin, ymax = -1, 1
    a.set_xlim(xmin, xmax)
    a.set_ylim(ymin, ymax)
    a.set_xticks([xmin, 0, xmax])
    a.set_yticks([ymin, 0, ymax])
    a.spines['left'].set_position(('outward', 4))
    a.spines['bottom'].set_position(('outward', 4))
    ## scatter plot meanburstratio FMI vs pupil area FMI:
    f, a = plt.subplots(figsize=figsize)
    wintitle('meanburstratio %s FMI vs pupil area FMI' % stimtype)
    # boolean indicating entries with non-NaN meanburstratio FMI & pupil area FMI:
    valid = ~np.isnan(areafmis) & ~np.isnan(brfmis)
    a.scatter(areafmis[valid], brfmis[valid], clip_on=False,
              marker='.', c='None', edgecolor='black', s=DEFSZ)
    # get fname of appropriate LMM .cvs file:
    if stimtype == 'mvi+grt':
        if EXPTYPE == 'pvmvis':
            fname = os.path.join('stats', 'figure_1_S6i_coefs.csv')
        elif EXPTYPE == 'ntsrmvis':
            fname = os.path.join('stats', 'figure_1_S6j_coefs.csv')
    if fname:
        try:
            df = pd.read_csv(fname)
            foundregression = True
        except FileNotFoundError:
            print('Missing file: %s' % fname)
            foundregression = False
        if foundregression:
            mm = df['slope'][0]
            b = df['intercept'][0]
            x = np.array([areafmis[valid].min(), areafmis[valid].max()])
            y = mm * x + b
            a.plot(x, y, '-', color='red') # plot linregress fit
    a.set_ylabel("Burst ratio FMI")
    a.set_xlabel("Pupil area FMI")
    xmin, xmax = -0.15, 0.15
    ymin, ymax = -1, 1
    a.set_xlim(xmin, xmax)
    a.set_ylim(ymin, ymax)
    a.set_xticks([xmin, 0, xmax])
    a.set_yticks([ymin, 0, ymax])
    a.spines['left'].set_position(('outward', 4))
    a.spines['bottom'].set_position(('outward', 4))
    # save to iposmi dataframe:
    if stimtype == 'mvi+grt':
        continue # don't overwrite stimtype column in iposmi, other values should be identical
    for mseu, areafmi, rfmi, brfmi in zip(mseus, areafmis, rfmis, brfmis):
        iposmi.loc[mseu, 'stimtype'] = stimtype # automatically adds a new mseu row if needed
        iposmi.loc[mseu, 'areafmi'] = areafmi
        iposmi.loc[mseu, 'meanratefmi'] = rfmi
        iposmi.loc[mseu, 'meanburstratiofmi'] = brfmi