"""Some, but not all Figure 4S1 plots, use run -i fig4S1.py. Companion to fig4.py.
fig1.py and fig3.py contribute some plots to 4S1 as well"""

mi = pd.MultiIndex.from_product([mvigrtmsustrs, STIMTYPES],
                                names=['msu', 'stimtype'])
fig4S1b = pd.DataFrame(index=mi, columns=['meanrate'])


"""Exporting all the various panels in fig4S1 to .csv is complicated.
Mostly they come from elsewhere:

a: maxFMI.csv : fig3.py
b: fig4S1b.csv : fig4S1.py
c--d: fig1.csv : fig1.py
e--f: maxFMI.csv, also duplicated in fig4.csv : fig4.py
g--l: left: fig1.csv; middle & right: fig3.csv : fig4S1.py
"""

# fig4S1b: scatter plot best movie vs best grating mean firing rate, during control condition,
# one point per msu:
figsize = DEFAULTFIGURESIZE
logmin, logmax = -1, 2
logticks = np.array([-1, 0, 1, 2])
mvivals, grtvals, exmplis, exmplmsustrs, normlis = [], [], [], [], []
keptmsui = 0 # manually init and increment instead of using enumerate()
for msustr in mvigrtmsustrs:
    try:
        mvival = bestmviresp['meanrate'][msustr, 'nat', 'none', False]
        grtval = bestgrtresp['meanrate'][msustr, 'none', False]
    except KeyError:
        continue # msustr doesn't exist in one of bestmviresp or bestgrtresp
    if pd.isna(mvival) or pd.isna(grtval): # missing one or both values
        continue
    mvivals.append(mvival)
    grtvals.append(grtval)
    # grt meanrate, meanrate02 and meanrate35 columns will all be identical:
    fig4S1b.loc[msustr, 'mvi']['meanrate'] = mvival # save
    fig4S1b.loc[msustr, 'grt']['meanrate'] = grtval # save
    if msustr in msu2exmpli:
        exmplis.append(keptmsui)
        exmplmsustrs.append(msustr)
    else:
        normlis.append(keptmsui)
    keptmsui += 1 # manually increment
mvivals = np.asarray(mvivals)
grtvals = np.asarray(grtvals)
f, a = plt.subplots(figsize=figsize)
wintitle('%s movie grating scatter %s %s' % ('meanrate', 'none', False))
# plot y=x line:
xyline = [10**logmin, 10**logmax], [10**logmin, 10**logmax]
a.plot(xyline[0], xyline[1], '--', color='gray', zorder=-1)
# plot normal (non-example) points:
c = desat(st82clr['none'], opto2alpha[False]) # do manual alpha mixing
a.scatter(grtvals[normlis], mvivals[normlis], clip_on=False,
          marker='.', c='None', edgecolor=c, s=DEFSZ)
# plot example points, one at a time:
for exmpli, msustr in zip(exmplis, exmplmsustrs):
    marker = exmpli2mrk[msu2exmpli[msustr]]
    c = exmpli2clr[msu2exmpli[msustr]]
    sz = exmpli2sz[msu2exmpli[msustr]]
    lw = exmpli2lw[msu2exmpli[msustr]]
    a.scatter(grtvals[exmpli], mvivals[exmpli], marker=marker, c=c, s=sz, lw=lw)
# plot mean:
#a.scatter(np.mean(grtvals), np.mean(mvivals),
#          c='red', edgecolor='red', s=50, marker='^')
a.set_xlabel('Grating FR (spk/s)')
a.set_ylabel('Movie FR (spk/s)')
a.set_xscale('log')
a.set_yscale('log')
a.set_xlim(10**logmin, 10**logmax)
a.set_ylim(10**logmin, 10**logmax)
a.set_xticks(10.0**logticks)
a.set_yticks(a.get_xticks()) # make log scale y ticks the same as x ticks
a.minorticks_off()
axes_disable_scientific(a)
a.set_aspect('equal')
a.spines['left'].set_position(('outward', 4))
a.spines['bottom'].set_position(('outward', 4))
#t, p = ttest_rel(grtvals, mvivals) # paired t-test
#a.add_artist(AnchoredText('p$=$%.2g' % p, loc='upper left', frameon=False))


# stripplot movie and grating mean firing rates during control condition, for all mseu:
np.random.seed(0) # to get identical horizontal jitter in strip plots on every run
figsize = DEFAULTFIGURESIZE
logmin, logmax = -2, 2
logticks = np.array([-2, -1, 0, 1, 2])
mvivals, grtvals = [], []
for mseustr in mvimseustrs:
    mvival = mviresp.loc[mseustr, 'nat', 'none', False]['meanrate']
    if pd.isna(mvival):
        continue
    mvivals.append(mvival)
for mseustr in grtmseustrs:
    grtval = grtresp.loc[mseustr, 'none', False]['meanrate']
    if pd.isna(grtval):
        continue
    grtvals.append(grtval)
mvivals = np.asarray(mvivals)
grtvals = np.asarray(grtvals)
f, a = plt.subplots(figsize=figsize)
wintitle('%s movie grating stripplot %s %s' % ('meanrate', 'none', False))
# plot y=0 line:
a.axhline(y=0, ls='--', marker='', color='lightgray', zorder=-np.inf)
data = pd.DataFrame.from_dict({'Movie':mvivals, 'Grating':grtvals},
                              orient='index').transpose()
sns.stripplot(ax=a, data=data, clip_on=False, marker='.',
              color='None', edgecolor='black', size=np.sqrt(50))
# plot mean with short horizontal lines:
#meanmvival, meangrtval = gmean(mvivals), gmean(grtvals)
#a.plot([-0.25, 0.25], [meanmvival, meanmvival], '-', lw=2, c='red', zorder=np.inf)
#a.plot([0.75, 1.25], [meangrtval, meangrtval], '-', lw=2, c='red', zorder=np.inf)
a.set_ylabel('Firing rate (spk/s)')
a.set_yscale('log')
a.set_ylim(10**logmin, 10**logmax)
a.set_yticks(10.0**logticks)
a.tick_params(bottom=False)
a.minorticks_off()
axes_disable_scientific(a, axiss=[a.yaxis]) # don't do it on x axis, messes up x label
a.spines['bottom'].set_position(('outward', 5))
a.spines['bottom'].set_visible(False)


# scatter plot blank movie meanrates:
figsize = DEFAULTFIGURESIZE[0]*1.02, DEFAULTFIGURESIZE[1] # tweak to make space for log units
logmin, logmax = -1.5, 2
logticks = np.array([-1, 0, 1, 2])
#log0min = logmin + 0.05
for kind in ['nat']:#MVIKINDS:
    for st8 in ['none']:#ALLST8S:
        f, a = plt.subplots(figsize=figsize)
        wintitle('opto meanrate blank movie %s %s' % (kind, st8))
        rons, roffs, exmplis, exmplmseustrs, normlis = [], [], [], [], []
        keptmseui = 0 # manually init and increment instead of using enumerate()
        for mseustr in mvimseustrs:
            meanrate = mviresp.loc[mseustr, kind, st8]['blankmeanrate']
            snr = mviresp.loc[mseustr, kind, st8]['snr']
            if meanrate.isna().any(): # missing one or both meanrates
                continue
            if (snr < SNRTHRESH).all(): # neither condition has decent SNR
                continue
            rons.append(meanrate[True])
            roffs.append(meanrate[False])
            rates = mviresp.loc[mseustr, kind, st8]['blankrates']
            #_, pval = ttest_ind(rates[False], rates[True], equal_var=False)
            #sgnfs.append(pval < SCATTERPTHRESH) # bool
            fig1.loc[mseustr, 'blankmeanrate'] = meanrate[False], meanrate[True] # save
            fig1.loc[mseustr]['blankrates'] = rates # save trial-wise values
            if mvimseu2exmpli.get(mseustr) == fig1exmpli:
                exmplis.append(keptmseui)
                exmplmseustrs.append(mseustr)
            else:
                normlis.append(keptmseui)
            keptmseui += 1 # manually increment
        rons = np.asarray(rons)
        roffs = np.asarray(roffs)
        # plot y=x line:
        xyline = [10**logmin, 10**logmax], [10**logmin, 10**logmax]
        a.plot(xyline[0], xyline[1], '--', color='gray', zorder=-1)
        # plot normal (non-example) points:
        a.scatter(rons[normlis], roffs[normlis], clip_on=False,
                  marker='.', c='None', edgecolor=st82clr[st8], s=DEFSZ)
        # plot example points, one at a time:
        for exmpli, mseustr in zip(exmplis, exmplmseustrs):
            marker = exmpli2mrk[mvimseu2exmpli[mseustr]]
            c = exmpli2clr[mvimseu2exmpli[mseustr]]
            sz = exmpli2sz[mvimseu2exmpli[mseustr]]
            lw = exmpli2lw[mvimseu2exmpli[mseustr]]
            a.scatter(rons[exmpli], roffs[exmpli], clip_on=False,
                      marker=marker, c=c, s=sz, lw=lw)
        a.set_xlabel('Suppression FR (spk/s)')
        a.set_ylabel('Feedback FR (spk/s)')
        a.set_xscale('log')
        a.set_yscale('log')
        a.set_xlim(10**logmin, 10**logmax)
        a.set_ylim(10**logmin, 10**logmax)
        a.set_xticks(10.0**logticks)
        a.set_yticks(a.get_xticks()) # make log scale y ticks the same as x ticks
        a.minorticks_off()
        axes_disable_scientific(a)
        a.set_aspect('equal')
        a.spines['left'].set_position(('outward', 4))
        a.spines['bottom'].set_position(('outward', 4))
        #t, p = ttest_rel(rons, roffs) # paired t-test
        #a.add_artist(AnchoredText('p$=$%.2g' % p, loc='lower right', frameon=False))
        #mu = rons.mean(), roffs.mean()
        #txt = '$\mathregular{\mu=%.1f, %.1f}$' % mu
        #a.add_artist(AnchoredText(txt, loc='upper left', frameon=False))


# scatter plot blank movie burst ratio:
figsize = DEFAULTFIGURESIZE[0]*1.05, DEFAULTFIGURESIZE[1] # tweak to make space for log units
logmin, logmax = -3, 0
logticks = np.array([-3, -2, -1, 0])
for kind in ['nat']:#MVIKINDS:
    for st8 in ['none']:#ALLST8S:
        f, a = plt.subplots(figsize=figsize)
        wintitle('opto burst ratio blank movie %s %s' % (kind, st8))
        brons, broffs, exmplis, exmplmseustrs, normlis = [], [], [], [], []
        keptmseui = 0 # manually init and increment instead of using enumerate()
        for mseustr in mvimseustrs:
            br = mviresp.loc[mseustr, kind, st8]['blankmeanburstratio']
            snr = mviresp.loc[mseustr, kind, st8]['snr']
            if br.isna().any(): # missing for at least one opto condition
                continue
            if (snr < SNRTHRESH).all(): # neither condition has decent SNR
                continue
            brons.append(br[True])
            broffs.append(br[False])
            burstratios = mviresp.loc[mseustr, kind, st8]['blankburstratios']
            #_, pval = ttest_ind(burstratios[False], burstratios[True], equal_var=False)
            #sgnfs.append(pval < SCATTERPTHRESH) # bool
            fig1.loc[mseustr, 'blankmeanburstratio'] = br[False], br[True] # save
            fig1.loc[mseustr]['blankburstratios'] = burstratios # save trial-wise values
            if mvimseu2exmpli.get(mseustr) == fig1exmpli:
                exmplis.append(keptmseui)
                exmplmseustrs.append(mseustr)
            else:
                normlis.append(keptmseui)
            keptmseui += 1 # manually increment
        brons, broffs = np.asarray(brons), np.asarray(broffs)
        # plot y=x line:
        xyline = [10**logmin, 10**logmax], [10**logmin, 10**logmax]
        a.plot(xyline[0], xyline[1], '--', color='gray', zorder=-1)
        # plot normal (non-example) points:
        a.scatter(brons[normlis], broffs[normlis], clip_on=True, # clip_on=False fails
                  marker='.', c='None', edgecolor=st82clr[st8], s=DEFSZ)
        # plot example points, one at a time:
        for exmpli, mseustr in zip(exmplis, exmplmseustrs):
            marker = exmpli2mrk[mvimseu2exmpli[mseustr]]
            c = exmpli2clr[mvimseu2exmpli[mseustr]]
            sz = exmpli2sz[mvimseu2exmpli[mseustr]]
            lw = exmpli2lw[mvimseu2exmpli[mseustr]]
            a.scatter(brons[exmpli], broffs[exmpli], clip_on=True, # clip_on=False fails
                      marker=marker, c=c, s=sz, lw=lw)
        a.set_xlabel('Suppression BR') # keep it short to maximize space for axes
        a.set_ylabel('Feedback BR')
        a.set_xscale('log')
        a.set_yscale('log')
        a.set_xlim(10**logmin, 10**logmax)
        a.set_ylim(10**logmin, 10**logmax)
        a.set_xticks(10.0**logticks)
        a.set_yticks(a.get_xticks()) # make log scale y ticks the same as x ticks
        a.minorticks_off()
        axes_disable_scientific(a)
        a.set_aspect('equal')
        a.spines['left'].set_position(('outward', 4))
        a.spines['bottom'].set_position(('outward', 4))
        #t, p = ttest_rel(brons, broffs) # paired t-test
        #a.add_artist(AnchoredText('p$=$%.2g' % p, loc='upper left', frameon=False))


# scatter plot blank and blankcond grating meanrates:
figsize = DEFAULTFIGURESIZE[0]*1.02, DEFAULTFIGURESIZE[1] # tweak to make space for log units
logmin, logmax = -1.5, 2
logticks = np.array([-1, 0, 1, 2])
#log0min = logmin + 0.05
for st8 in ['none']:#ALLST8S:
    for blnkname, blnksname in {'blankmeanrate':'blankrates',
                                'blankcondmeanrate':'blankcondrates'}.items():
        f, a = plt.subplots(figsize=figsize)
        wintitle('opto meanrate %s grating %s' % (blnkname, st8))
        rons, roffs, exmplis, exmplmseustrs, normlis = [], [], [], [], []
        keptmseui = 0 # manually init and increment instead of using enumerate()
        for mseustr in grtmseustrs:
            trialis = grtresp.loc[mseustr, st8]['trialis'] # non-blank & blank trialis
            if trialis.isna().any(): # missing for at least one opto condition
                continue
            ntrials = { opto:len(trialis[opto]) for opto in OPTOS } # should be equal
            blnkratesfull = pd.Series({ opto:np.full(ntrials[opto], np.nan) # pad
                                        for opto in [False, True] })
            meanrate = grtresp.loc[mseustr, st8][blnkname]
            if meanrate.isna().any(): # missing for at least one opto condition, can't plot
                fig3.loc[mseustr][blnksname] = blnkbrsfull # save nans for all trials
                continue
            rons.append(meanrate[True])
            roffs.append(meanrate[False])
            rates = grtresp.loc[mseustr, st8]['rates'] # trial-wise
            nnblnktrials = { opto:len(rates[opto]) for opto in OPTOS } # num non-blank trials
            blnkrates = grtresp.loc[mseustr, st8][blnksname]
            if blnksname == 'blankrates':
                for opto in OPTOS:
                    blnkratesfull[opto][:nnblnktrials[opto]] = blnkrates[opto]
            else: # blnksname == 'blankcondrates':
                for opto in OPTOS:
                    blnkratesfull[opto][nnblnktrials[opto]:] = blnkrates[opto]
            fig3.loc[mseustr, blnkname] = meanrate[False], meanrate[True] # save
            fig3.loc[mseustr][blnksname] = blnkratesfull # save padded trial-wise values
            if mseustr in grtmseu2exmpli:
                exmplis.append(keptmseui)
                exmplmseustrs.append(mseustr)
            else:
                normlis.append(keptmseui)
            keptmseui += 1 # manually increment
        rons = np.asarray(rons)
        roffs = np.asarray(roffs)
        # plot y=x line:
        xyline = [10**logmin, 10**logmax], [10**logmin, 10**logmax]
        a.plot(xyline[0], xyline[1], '--', color='gray', zorder=-1)
        # plot normal (non-example) points:
        a.scatter(rons[normlis], roffs[normlis], clip_on=True, # fails
                  marker='.', c='None', edgecolor=st82clr[st8], s=DEFSZ)
        # plot example points, one at a time:
        for exmpli, mseustr in zip(exmplis, exmplmseustrs):
            marker = exmpli2mrk[grtmseu2exmpli[mseustr]]
            c = exmpli2clr[grtmseu2exmpli[mseustr]]
            sz = exmpli2sz[grtmseu2exmpli[mseustr]]
            lw = exmpli2lw[grtmseu2exmpli[mseustr]]
            a.scatter(rons[exmpli], roffs[exmpli], clip_on=False,
                      marker=marker, c=c, s=sz, lw=lw)
        a.set_xlabel('Suppression FR (spk/s)')
        a.set_ylabel('Feedback FR (spk/s)')
        a.set_xscale('log')
        a.set_yscale('log')
        a.set_xlim(10**logmin, 10**logmax)
        a.set_ylim(10**logmin, 10**logmax)
        a.set_xticks(10.0**logticks)
        a.set_yticks(a.get_xticks()) # make log scale y ticks the same as x ticks
        a.minorticks_off()
        axes_disable_scientific(a)
        a.set_aspect('equal')
        a.spines['left'].set_position(('outward', 4))
        a.spines['bottom'].set_position(('outward', 4))
        #t, p = ttest_rel(rons, roffs) # paired t-test
        #a.add_artist(AnchoredText('p$=$%.2g' % p, loc='lower right', frameon=False))
        #mu = rons.mean(), roffs.mean()
        #txt = '$\mathregular{\mu=%.1f, %.1f}$' % mu
        #a.add_artist(AnchoredText(txt, loc='upper left', frameon=False))


# scatter plot blank and blank cond grating burst ratio:
figsize = DEFAULTFIGURESIZE[0]*1.05, DEFAULTFIGURESIZE[1] # tweak to make space for log units
logmin, logmax = -3, 0
logticks = np.array([-3, -2, -1, 0])
for st8 in ['none']:#ALLST8S:
    for blnkname, blnksname in {'blankmeanburstratio':'blankburstratios',
                                'blankcondmeanburstratio':'blankcondburstratios'}.items():
        f, a = plt.subplots(figsize=figsize)
        wintitle('opto burst ratio %s grating %s' % (blnkname, st8))
        brons, broffs, exmplis, exmplmseustrs, normlis = [], [], [], [], []
        keptmseui = 0 # manually init and increment instead of using enumerate()
        for mseustr in grtmseustrs:
            trialis = grtresp.loc[mseustr, st8]['trialis'] # non-blank & blank trialis
            if trialis.isna().any(): # missing for at least one opto condition
                continue
            ntrials = { opto:len(trialis[opto]) for opto in OPTOS } # should be equal
            blnkbrsfull = pd.Series({ opto:np.full(ntrials[opto], np.nan) # pad
                                      for opto in [False, True] })
            br = grtresp.loc[mseustr, st8][blnkname] # mean BR
            if br.isna().any(): # missing for at least one opto condition, can't plot
                fig3.loc[mseustr][blnksname] = blnkbrsfull # save nans for all trials
                continue
            brons.append(br[True])
            broffs.append(br[False])
            burstratios = grtresp.loc[mseustr, st8]['burstratios'] # trial-wise
            nnblnktrials = { opto:len(burstratios[opto]) for opto in OPTOS } # num non-blank trials
            blnkbrs = grtresp.loc[mseustr, st8][blnksname]
            if blnksname == 'blankburstratios':
                for opto in OPTOS:
                    blnkbrsfull[opto][:nnblnktrials[opto]] = blnkbrs[opto]
            else: # blnksname == 'blankcondburstratios':
                for opto in OPTOS:
                    blnkbrsfull[opto][nnblnktrials[opto]:] = blnkbrs[opto]
            fig3.loc[mseustr, blnkname] = br[False], br[True] # save
            fig3.loc[mseustr][blnksname] = blnkbrsfull # save padded trial-wise values
            if mseustr in grtmseu2exmpli:
                exmplis.append(keptmseui)
                exmplmseustrs.append(mseustr)
            else:
                normlis.append(keptmseui)
            keptmseui += 1 # manually increment
        brons, broffs = np.asarray(brons), np.asarray(broffs)
        # plot y=x line:
        xyline = [10**logmin, 10**logmax], [10**logmin, 10**logmax]
        a.plot(xyline[0], xyline[1], '--', color='gray', zorder=-1)
        # plot normal (non-example) points:
        a.scatter(brons[normlis], broffs[normlis], clip_on=True, # clip_on=False fails
                  marker='.', c='None', edgecolor=st82clr[st8], s=DEFSZ)
        # plot example points, one at a time:
        for exmpli, mseustr in zip(exmplis, exmplmseustrs):
            marker = exmpli2mrk[grtmseu2exmpli[mseustr]]
            c = exmpli2clr[grtmseu2exmpli[mseustr]]
            sz = exmpli2sz[grtmseu2exmpli[mseustr]]
            lw = exmpli2lw[grtmseu2exmpli[mseustr]]
            a.scatter(brons[exmpli], broffs[exmpli], clip_on=True, # clip_on=False fails
                      marker=marker, c=c, s=sz, lw=lw)
        a.set_xlabel('Suppression BR') # keep it short to maximize space for axes
        a.set_ylabel('Feedback BR')
        a.set_xscale('log')
        a.set_yscale('log')
        a.set_xlim(10**logmin, 10**logmax)
        a.set_ylim(10**logmin, 10**logmax)
        a.set_xticks(10.0**logticks)
        a.set_yticks(a.get_xticks()) # make log scale y ticks the same as x ticks
        a.minorticks_off()
        axes_disable_scientific(a)
        a.set_aspect('equal')
        a.spines['left'].set_position(('outward', 4))
        a.spines['bottom'].set_position(('outward', 4))
        #t, p = ttest_rel(brons, broffs) # paired t-test
        #a.add_artist(AnchoredText('p$=$%.2g' % p, loc='upper left', frameon=False))