"""Figure 5S2 pupil area plots, use run -i fig5S2.py"""

## collect PSTH trial matrices with same timebase as pupil area signal, masked by runspeed
## to include only sitting periods, then also masked by pupil constriction/dilation upper and
## lower quartiles. Store in arearate_mviresp DataFrame indexed by [mseu, opto]
allmseustrs = mviresp.index.levels[0].values # get all mice, series, experiments, units
levels = [allmseustrs, OPTOS]
names = ['mseu', 'opto']
mi = pd.MultiIndex.from_product(levels, names=names)
columns = ['trialmat', 'bins', 'binw', 'tq_trialmat', 'bq_trialmat',
           'psth', 'tq_psth', 'bq_psth', 'pct_changes']
arearate_mviresp = pd.DataFrame(index=mi, columns=columns)
arearates = {} # mse to (arearate, arearate_ts) mapping
# for each row (mse) in ipos_opto df, find 1st derivative of pupil area trial matrix:
for index, row in ipos_opto.iterrows():
    msestr, stimtype, opto = index
    if not stimtype == 'mvi':
        continue
    print(msestr, stimtype, opto)
    DX = 1 # decimation factor, e.g. 5 increases sampling period by 5X, from ~20 ms to ~100 ms
    area_trialmat = row['area_trialmat'][:, ::DX]
    area_trialts = row['area_trialts'][::DX]
    ntrials, nt = area_trialmat.shape
    arearate = np.diff(area_trialmat, axis=1) # 2D: ntrials, nt-1
    arearate_ts = area_trialts[:-1] # drop the last timepoint due to diff
    arearates[msestr] = arearate.copy(), arearate_ts.copy()
    # get locomotion info and mask out all instances of running:
    runrow = runspeed.loc[msestr, 'nat', opto]
    assert (row['trialis'] == runrow['trialis']).all() # sanity check
    run_ts, speed = runrow['t'], runrow['speed']
    # each runspeed trial can have slightly different number of timepoints, so have to
    # resample each runspeed trial separately:
    assert len(speed) == ntrials
    for i in range(ntrials): # iterate over trials
        f = scipy.interpolate.interp1d(run_ts[i], speed[i], axis=-1, kind='linear',
                                       fill_value='extrapolate')
        rs_speed = f(arearate_ts) # runspeed signal resampled to arearate_ts
        # mask out arearate during running periods:
        runmask = rs_speed > 0.25 # cm/s
        arearate[i, runmask] = np.nan # these points shouldn't influence the percentile calc
    # get top and bottom quartiles for each timepoint, averaged across all trials:
    tq, bq = np.nanpercentile(arearate, [75, 25], axis=0) # each is nt-1
    assert len(tq) == nt - 1
    assert len(bq) == nt - 1
    # mask out timepoints per trial in which arearate falls in top or bottom quartile:
    tqmask = np.int8(arearate >= tq) # 2D binary: ntrials, nt-1
    bqmask = np.int8(arearate <= bq)
    # create top and bottom quartile arearate arrays:
    arearatetq = np.full(arearate.shape, np.nan)
    arearatetq[tqmask] = arearate[tqmask]
    arearatebq = np.full(arearate.shape, np.nan)
    arearatebq[bqmask] = arearate[bqmask]
    # get all mseu for this mse:
    hits = findmse(allmseustrs, msestr)
    mseustrs = allmseustrs[hits] # all mseu that belong to this mse
    # get bins for calculating PSTHs using arearate timepoints:
    binw = np.median(np.diff(arearate_ts))
    ledges = arearate_ts # left edges
    redges = np.hstack([arearate_ts[1:], arearate_ts[-1]+binw]) # right edges
    bins = np.column_stack((ledges, redges))
    # calculate PSTHs and binned % rate change for each rate percentile range for each mseu:
    for mseustr in mseustrs:
        trialmat = np.full(arearate.shape, np.nan)
        mviresprow = mviresp.loc[mseustr, 'nat', 'none', opto]
        if np.isnan(mviresprow['dt']) or mviresprow['snr'] < SNRTHRESH:
            continue
        #raster, braster = mviresprow['raster'], mviresprow['braster']
        raster = mviresprow['raster']
        # generate trial-based psth matrix, binned at bins:
        for triali, trspkts in enumerate(raster):
            # tres is ignored when kernel == None
            trpsth = raster2psth([trspkts], bins, binw, None, kernel=None)
            trialmat[triali] = trpsth
        # create top quartile trialmat by masking out all entries in trialmat which are not
        # in top quartile:
        tq_trialmat = trialmat.copy()
        tq_trialmat[tqmask == 0] = np.nan
        # create bottom quartile trialmat by masking out all entries in trialmat which are not
        # in bottom quartile:
        bq_trialmat = trialmat.copy()
        bq_trialmat[bqmask == 0] = np.nan
        # compute mean across trials, only timepoints remain:
        psth = np.nanmean(trialmat, axis=0)
        tq_psth = np.nanmean(tq_trialmat, axis=0)
        bq_psth = np.nanmean(bq_trialmat, axis=0)
        # get firing rate percentile threshold ranges of the psth:
        lo, bq, mid, tq, hi = np.percentile(psth, [0, 25, 50, 75, 100])
        rpctileranges = [[lo, bq], [bq, mid], [mid, tq], [tq, hi]]
        # calculate binned percent change of PSTH of top quartile vs bottom quartile pupil
        # arearate, for each firing rate percentile range:
        pct_changes = []
        for rpctilerange in rpctileranges:
            # find PSTH timepoints that fall in this rate percentile range:
            tis = (rpctilerange[0] < psth) & (psth <= rpctilerange[1]) # bool array
            if (tis == False).all(): # no timepoints in this rpctilerange
                pct_change = np.nan
            else:
                tqmean = tq_psth[tis].mean()
                bqmean = bq_psth[tis].mean()
                if bqmean > 0:
                    pct_change = (tqmean - bqmean) / bqmean * 100
                else:
                    #print('bqmean', bqmean)
                    pct_change = np.nan
            pct_changes.append(pct_change)
        pct_changes = np.array(pct_changes)
        arearate_mviresprow = arearate_mviresp.loc[mseustr, opto]
        arearate_mviresprow['trialmat'] = trialmat
        arearate_mviresprow['tq_trialmat'] = tq_trialmat
        arearate_mviresprow['bq_trialmat'] = bq_trialmat
        arearate_mviresprow['bins'] = bins
        arearate_mviresprow['binw'] = binw
        arearate_mviresprow['psth'] = psth
        arearate_mviresprow['tq_psth'] = tq_psth
        arearate_mviresprow['bq_psth'] = bq_psth
        arearate_mviresprow['pct_changes'] = pct_changes # length 4, 1 value per FR quartile


## collect % change values for both opto conditions, and resample them to allow
## confidence interval calculations; keep only rows that have non-nan entries in pct_changes:
arearate_mviresp_notna = arearate_mviresp.dropna(subset=['pct_changes'])
pcoffs = np.vstack(arearate_mviresp_notna.loc[(slice(None), False), 'pct_changes'].values)
pcons = np.vstack(arearate_mviresp_notna.loc[(slice(None), True), 'pct_changes'].values)
noffs, nons = len(pcoffs), len(pcons) # not necessarily the same, some can be unpaired
NSAMPLES = 5000
rspcoffmeds = np.full((NSAMPLES, 4), np.nan) # resampled % change opto off medians
rspconmeds = np.full((NSAMPLES, 4), np.nan)
rng = np.random.default_rng(0) # init random number generator with seed 0
for samplei in range(NSAMPLES):
    rspcoffs = rng.choice(pcoffs, noffs, replace=True, axis=0)
    rspcons = rng.choice(pcons, nons, replace=True, axis=0)
    rspcoffmeds[samplei] = np.nanmedian(rspcoffs, axis=0)
    rspconmeds[samplei] = np.nanmedian(rspcons, axis=0)
# calculate medians and percentiles for CIs:
med_pcoffs = np.nanmedian(pcoffs, axis=0)
med_rspcoffmeds = np.nanmedian(rspcoffmeds, axis=0)
lci_rspcoffmeds = np.nanpercentile(rspcoffmeds, 2.5, axis=0)
uci_rspcoffmeds = np.nanpercentile(rspcoffmeds, 97.5, axis=0)
med_pcons = np.nanmedian(pcons, axis=0)
med_rspconmeds = np.nanmedian(rspconmeds, axis=0)
lci_rspconmeds = np.nanpercentile(rspconmeds, 2.5, axis=0)
uci_rspconmeds = np.nanpercentile(rspconmeds, 97.5, axis=0)
print('median pcoffs:     ', med_pcoffs)
print('median rspcoffmeds:', med_rspcoffmeds)
print('  2.5% rspcoffmeds:', lci_rspcoffmeds)
print(' 97.5% rspcoffmeds:', uci_rspcoffmeds)
print('median pcons:      ', med_pcons)
print('median rspconmeds: ', med_rspconmeds)
print('  2.5% rspconmeds: ', lci_rspconmeds)
print(' 97.5% rspconmeds: ', uci_rspconmeds)


## plot arearate trial matrix for example experiment:
exmplmsestrs = ['PVCre_2017_0006_s03_e03',
                'PVCre_2017_0015_s07_e04',
                'PVCre_2019_0002_s08_e04']
for msestr in exmplmsestrs:#list(arearates):
    arearate, arearate_ts = arearates[msestr]
    binw = np.median(np.diff(arearate_ts)) # s
    arearate = arearate / binw # normalized pupil area per sec
    ntrials = len(arearate)
    tmin, tmax = arearate_ts[0], np.ceil(arearate_ts[-1])
    extent = tmin, tmax, ntrials, 1
    f, a = plt.subplots(figsize=(6, 3))
    wintitle('arearate trialmat %s' % msestr)
    #std = np.nanstd(arearate)
    #vlim = 3*std
    vlim = 0.3 # might result in a bit of clipping of outlier values
    im = a.imshow(arearate, extent=extent, aspect='auto', origin='upper', cmap='seismic',
                  interpolation='nearest', vmin=-vlim, vmax=vlim)
    plt.colorbar(im, ticks=[-vlim, 0, vlim])
    #a.set_xticks(np.arange(5))
    a.set_xlabel('Time (s)')
    a.set_ylabel('Trial')
    #a.set_xticks(tmin, tmax, 50)
    yticks = np.arange(0, ntrials+1, 50)
    yticks[0] = 1
    a.set_yticks(yticks)


## plot example raster plot:
exmplmseustr = 'PVCre_2019_0002_s08_e04_u32'
mviresprow = mviresp.loc[exmplmseustr, 'nat', 'none', False]
raster = mviresprow['raster']
dt = mviresprow['dt'] # stimulus duration (s)
SCATTER = False # True: low quality, but fast; False: high quality, but slow
title = 'movie raster %s %s' % (mseustr, False)
a = simpletraster(raster=raster, alpha=0.5, s=1, dt=dt, offsets=[0, 0],
                  scatter=SCATTER,
                  title=title, inchespersec=1.5*RASTERSCALEX,
                  inchespertrial=1/50*RASTERSCALEX,
                  xaxis=True)
a.spines['left'].set_position(('outward', 4))
a.spines['bottom'].set_position(('outward', 4))
yticks = np.arange(0, ntrials+1, 50)
yticks[0] = 1
a.set_yticks(yticks)


## plot example PSTHs:
exmplmsestr = 'PVCre_2019_0002_s08_e04'
hits = findmse(allmseustrs, exmplmsestr)
exmplmseustrs = allmseustrs[hits]
alpha = 0.5
for mseustr in [exmplmseustr]:#exmplmseustrs:
    psthoptos = arearate_mviresp.loc[mseustr][['tq_psth', 'bq_psth']]
    #maxpsth = np.vstack(psthoptos.values).max() # max of mean PSTH across both opto conditions
    maxpsth = np.hstack(psthoptos.values.flatten()).max()
    for opto in OPTOS:
        row = arearate_mviresp.loc[(mseustr, opto)]
        bins, psth, tq_psth, bq_psth = row[['bins', 'psth', 'tq_psth', 'bq_psth']]
        if np.any(np.isnan(bins)):
            continue
        t = bins.mean(axis=1)
        f, a = plt.subplots(figsize=(5, 3))
        wintitle('top bottom qrt PSTH %s opto=%s' % (mseustr, opto))
        #maxpsth = psth.max()
        #a.plot(t, psth/maxpsth, '-', color='black', alpha=alpha, label='Mean')
        a.plot(t, tq_psth, '-', color='red', alpha=alpha, label='Top quartile')
        a.plot(t, bq_psth, '-', color='blue', alpha=alpha, label='Bottom quartile')
        a.spines['left'].set_position(('outward', 4))
        a.spines['bottom'].set_position(('outward', 4))
        a.set_xlabel('Time (s)')
        a.set_ylabel('Normalized firing rate')
        a.set_xlim(0, 5)
        a.set_ylim(-1, maxpsth)
        a.set_xticks([0, 1, 2, 3, 4, 5])
        #a.set_yticks([0, 1])
        #a.legend()


## bar plot % change of firing for top vs bottom quartile pupil arearate, median
## across population, one value per firing rate quartile. Separate plots for
## each opto condition. See Reimer2014 fig5e:
opto2height = {False:med_pcoffs, True:med_pcons}
opto2yerr = {False:np.array([med_pcoffs-lci_rspcoffmeds, uci_rspcoffmeds-med_pcoffs]),
             True:np.array([med_pcons-lci_rspconmeds, uci_rspconmeds-med_pcons])}
for opto in OPTOS:
    f, a = plt.subplots(figsize=(3, 4))
    wintitle('top bottom qrt pct FR change bar opto=%s' % opto)
    x = np.arange(4)
    height = opto2height[opto]
    yerr = opto2yerr[opto]
    color = opto2clr[opto]
    a.bar(x, height, yerr=yerr, color=color, ecolor='gray')
    #a.hlines(0, 0, 3)
    #a.axhline(y=0, c='k', ls='--', alpha=0.5)
    a.set_xlabel('Firing rate quartile')
    a.set_ylabel('% Change')
    a.spines['left'].set_position(('outward', 4))
    a.spines['bottom'].set_position(('outward', 4))
    #a.spines['bottom'].set_visible(False)
    a.set_xticks([0, 1, 2, 3])
    a.set_xticklabels(['0-25%', '25-50%', '50-75%', '75-100%'], rotation=45, ha='center')


## bar plot % change of firing for top vs bottom quartile pupil arearate, median
## across population, one value per firing rate quartile. One plot for both opto conditions.
## See Reimer2014 fig5e:
w = 0.4 # bar width
opto2x = {False:np.arange(-w/2, 4-w/2, 1), True:np.arange(w/2, 4+w/2, 1)}
opto2height = {False:med_pcoffs, True:med_pcons}
opto2yerr = {False:np.array([med_pcoffs-lci_rspcoffmeds, uci_rspcoffmeds-med_pcoffs]),
             True:np.array([med_pcons-lci_rspconmeds, uci_rspconmeds-med_pcons])}
opto2label = {False:'V1 Control', True:'V1 Suppressed'}
f, a = plt.subplots(figsize=(3, 3))
wintitle('top bottom qrt pct FR change bar opto=both')
for opto in OPTOS:
    x = opto2x[opto]
    height = opto2height[opto]
    yerr = opto2yerr[opto]
    color = opto2clr[opto]
    label = opto2label[opto]
    a.bar(x, height, width=w, yerr=yerr, color=color, label=label, ecolor='gray')
    #a.hlines(0, 0, 3)
    #a.axhline(y=0, c='k', ls='--', alpha=0.5)
a.set_xlabel('Firing rate quartile')
a.set_ylabel('% Firing rate change')
a.spines['left'].set_position(('outward', 4))
a.spines['bottom'].set_position(('outward', 4))
#a.spines['bottom'].set_visible(False)
a.set_xticks([0, 1, 2, 3])
a.set_xticklabels(['0-25%', '25-50%', '50-75%', '75-100%'], rotation=45, ha='center')
#ymin = np.floor(min(lci_rspcoffmeds.min(), lci_rspconmeds.min()))
ymin = -1.5
ymax = np.ceil(max(uci_rspcoffmeds.max(), uci_rspconmeds.max()))
a.set_ylim(ymin=ymin, ymax=ymax)
#a.legend()


## calculate mean FR, reliability, and SNR for top and bottom pupil arearate quartiles:
columns = ['mseu', 'opto', 'tqr', 'bqr', 'tqrel', 'bqrel', 'tqsnr', 'bqsnr']
fig5S2 = []
for mseustr in allmseustrs:
    arearate_mviresprow = arearate_mviresp.loc[mseustr]
    tq_trialmats = arearate_mviresprow['tq_trialmat']
    bq_trialmats = arearate_mviresprow['bq_trialmat']
    if np.any(pd.isna(tq_trialmats)) or np.any(pd.isna(bq_trialmats)):
        continue
    for opto in OPTOS:
        tq_trialmat = tq_trialmats[opto]
        bq_trialmat = bq_trialmats[opto]
        tqr = np.nanmean(tq_trialmat) # top quartile, mean rate
        bqr = np.nanmean(bq_trialmat) # bottom quartile, mean rate
        tqrel, _ = reliability(tq_trialmat, average='mean', ignore_nan=True)
        bqrel, _ = reliability(bq_trialmat, average='mean', ignore_nan=True)
        tqsnr = snr_baden2016(tq_trialmat)
        bqsnr = snr_baden2016(bq_trialmat)
        # skip mseu that have mean rates that fall below thresh due to being limited
        # by periods of pupil area dynamics:
        if tqr < RATETHRESH and bqr < RATETHRESH:
            continue
        row = {}
        row['mseu'] = mseustr
        row['opto'] = opto
        row['tqr'] = tqr # top quartile mean firing rate
        row['bqr'] = bqr
        row['tqrel'] = tqrel # top quartile reliability
        row['bqrel'] = bqrel
        row['tqsnr'] = tqsnr # top quartile SNR
        row['bqsnr'] = bqsnr
        fig5S2.append(row)
fig5S2 = pd.DataFrame(fig5S2).set_index(['mseu', 'opto'])


## scatter plot top vs bottom pupil arearate quartiles for meanrate, reliability, and SNR.
## See Reimer2014 fig5fghi:
figsize = DEFAULTFIGURESIZE[0]*1.02, DEFAULTFIGURESIZE[1] # tweak to make space for log units
for opto in OPTOS:
    rows = fig5S2.loc[(slice(None), opto), :]
    normlis, = np.where(rows.reset_index()['mseu'] != exmplmseustr)
    exmpli, = np.where(rows.reset_index()['mseu'] == exmplmseustr)
    ## scatter plot mean firing rates during top vs. bottom quartile pupil area
    ## dilation/constriction rate:
    f, a = plt.subplots(figsize=figsize)
    wintitle('top bottom qrt scatter meanrate opto=%s' % opto)
    # replace any FRs < MINFRVAL with MINFRVAL, and any > MAXFRVAL with MAXFRVAL, for plotting:
    logmin, logmax = -0.9, 1.9
    logticks = np.array([0, 1])
    MINFRVAL, MAXFRVAL = 10**logmin, 10**logmax
    tqr, bqr = rows['tqr'].values, rows['bqr'].values
    tqr[tqr < MINFRVAL] = MINFRVAL
    bqr[bqr < MINFRVAL] = MINFRVAL
    tqr[tqr > MAXFRVAL] = MAXFRVAL
    bqr[bqr > MAXFRVAL] = MAXFRVAL
    # plot y=x line:
    linmin, linmax = MINFRVAL, MAXFRVAL
    xyline = [linmin, linmax], [linmin, linmax]
    a.plot(xyline[0], xyline[1], '--', color='gray', zorder=-1)
    # plot normal points:
    a.scatter(bqr[normlis], tqr[normlis], clip_on=False, marker='.',
              c='None', edgecolor='black', s=DEFSZ)
    # plot example point:
    a.scatter(bqr[exmpli], tqr[exmpli], clip_on=False, marker='x',
              c=green, s=60)
    a.set_title('FR (Hz), opto=%s' % opto)
    a.set_ylabel("Top 1/4 arearate")
    a.set_xlabel("Bottom 1/4 arearate")
    a.set_xscale('log')
    a.set_yscale('log')
    a.set_xlim(MINFRVAL, MAXFRVAL)
    a.set_ylim(MINFRVAL, MAXFRVAL)
    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(bqr, tqr) # paired t-test
    a.add_artist(AnchoredText('p$=$%.2g' % p, loc='lower right', frameon=False))
    ## scatter plot reliability during top vs. bottom quartile pupil area
    ## dilation/constriction rate:
    f, a = plt.subplots(figsize=figsize)
    wintitle('top bottom qrt scatter rel opto=%s' % opto)
    tqrel, bqrel = rows['tqrel'].values, rows['bqrel'].values
    # plot y=x line:
    linmin, linmax = -0.002, 0.06
    xyline = [linmin, linmax], [linmin, linmax]
    a.plot(xyline[0], xyline[1], '--', color='gray', zorder=-1)
    # plot normal points:
    a.scatter(bqrel[normlis], tqrel[normlis], clip_on=False, marker='.',
              c='None', edgecolor='black', s=DEFSZ)
    # plot example point:
    a.scatter(bqrel[exmpli], tqrel[exmpli], clip_on=False, marker='x',
              c=green, s=60)
    a.set_title('Reliability, opto=%s' % opto)
    a.set_ylabel("Top 1/4 arearate")
    a.set_xlabel("Bottom 1/4 arearate")
    #a.set_xscale('log')
    #a.set_yscale('log')
    a.set_xlim(linmin, linmax)
    a.set_ylim(linmin, linmax)
    #a.set_xticks(10.0**logticks)
    a.set_xticks([0, linmax])
    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))
    # filter out nans for t-test:
    t, p = ttest_rel(bqrel, tqrel) # paired t-test
    a.add_artist(AnchoredText('p$=$%.2g' % p, loc='lower right', frameon=False))
    ## scatter plot SNR during top vs. bottom quartile pupil area
    ## dilation/constriction rate:
    f, a = plt.subplots(figsize=figsize)
    wintitle('top bottom qrt scatter snr opto=%s' % opto)
    tqsnr, bqsnr = rows['tqsnr'].values, rows['bqsnr'].values
    # plot y=x line:
    linmin, linmax = 0, 0.7
    xyline = [linmin, linmax], [linmin, linmax]
    a.plot(xyline[0], xyline[1], '--', color='gray', zorder=-1)
    # plot normal points:
    a.scatter(bqsnr[normlis], tqsnr[normlis], clip_on=False, marker='.',
              c='None', edgecolor='black', s=DEFSZ)
    # plot example point:
    a.scatter(bqsnr[exmpli], tqsnr[exmpli], clip_on=False, marker='x',
              c=green, s=60)
    a.set_title('SNR, opto=%s' % opto)
    a.set_ylabel("Top 1/4 arearate")
    a.set_xlabel("Bottom 1/4 arearate")
    #a.set_xscale('log')
    #a.set_yscale('log')
    a.set_xlim(linmin, linmax)
    a.set_ylim(linmin, linmax)
    #a.set_xticks(10.0**logticks)
    a.set_xticks([0, linmax])
    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))
    # filter out nans for t-test:
    t, p = ttest_rel(bqsnr, tqsnr) # paired t-test
    a.add_artist(AnchoredText('p$=$%.2g' % p, loc='lower right', frameon=False))