"""Figure 1S2 and 5S1 eye position plots, use run -i ipos.py"""

index = pd.Index([], name='mseu') # init as empty
iposmi = pd.DataFrame(index=index, columns=['stimtype', 'iposfmi', 'relfmi', 'iposrmi',
                                            'relrmi', 'areafmi', 'meanratefmi',
                                            'meanburstratiofmi'])
maxiposedge, iposbinw = 10, 0.5
iposedges = np.arange(0, maxiposedge+iposbinw, iposbinw)


"""Figure 1S2"""

## plot cross trial stdev PDFs for control and opto trials:
for stimtype, stimtypei in zip(STIMTYPESPLUSALL, STIMTYPESPLUSALLI):
    f, a = plt.subplots(figsize=DEFAULTFIGURESIZE)
    wintitle('ipos_opto %s stdev cross trial PDF' % stimtype)
    stdctrl = ipos_opto['std_xpos_cross'][:, stimtypei, False]
    stdopto = ipos_opto['std_xpos_cross'][:, stimtypei, True]
    a.hist(stdctrl, bins=iposedges, histtype='step', lw=1.5,
           color=opto2clr[False], alpha=1, label=opto2fb[False])
    a.hist(stdopto, bins=iposedges, histtype='step', lw=1.5,
           color=opto2clr[True], alpha=1, label=opto2fb[True])
    a.set_xlabel('Eye position $\sigma$ ($\mathregular{\degree}$)')
    a.set_ylabel('Experiment count')
    #a.legend(frameon=False)
    print('opto %s std_cross:' % stimtype, wilcoxon(stdctrl, stdopto))


## plot cross trial stdev CDFs for control and opto trials:
for stimtype, stimtypei in zip(STIMTYPESPLUSALL, STIMTYPESPLUSALLI):
    f, a = plt.subplots(figsize=DEFAULTFIGURESIZE)
    wintitle('ipos_opto %s stdev cross trial CDF' % stimtype)
    stdctrl = ipos_opto['std_xpos_cross'][:, stimtypei, False]
    stdopto = ipos_opto['std_xpos_cross'][:, stimtypei, True]
    stdctrlbins = list(np.unique(stdctrl)) + [maxiposedge]
    stdoptobins = list(np.unique(stdopto)) + [maxiposedge]
    a.hist(stdctrl, bins=stdctrlbins, density=True, cumulative=True, histtype='step',
           lw=1.5, color=opto2clr[False], alpha=1, label=opto2fb[False])
    a.hist(stdopto, bins=stdoptobins, density=True, cumulative=True, histtype='step',
           lw=1.5, color=opto2clr[True], alpha=1, label=opto2fb[True])
    a.set_xlabel('Eye position $\sigma$ ($\mathregular{\degree}$)')
    a.set_ylabel('Cumulative probability')
    a.set_xlim(0, 9.5)
    a.set_xticks([0, 5])
    a.set_yticks([0, 0.5, 1])
    #a.legend(frameon=False)
    _, KS_p = ks_2samp(stdctrl, stdopto)
    txt = '$\mathregular{p_{KS}=%.2g}$' % KS_p # prob that distribs are the same
    a.add_artist(AnchoredText(txt, loc='upper left', frameon=False))


## scatter plot movie reliability FMI vs eye position stdev FMI:
# run reliability FMI for each movie mseu:
mseus, relfmis = zip(*mviFMI['rel'][:, 'run'].items()) # unpack indices and relfmi values
mseus, relfmis = np.array(mseus), np.array(relfmis)
# across-trial eye position stdev FMI for all movie experiments:
mseiposfmis = ((ipos_opto['std_xpos_cross'][:, 'mvi', False]
                - ipos_opto['std_xpos_cross'][:, 'mvi', True]) /
               (ipos_opto['std_xpos_cross'][:, 'mvi', False]
                + ipos_opto['std_xpos_cross'][:, 'mvi', True]))
# get mse string of every mseu in mviFMI:
mses = np.array(['_'.join(mseu.split('_')[:-1]) for mseu in mseus])
# assign appropriate iposfmi to each mseu:
iposfmis = np.tile(np.nan, len(mseus)) # init to nans
for msei, mse in enumerate(mses):
    if mse in mseiposfmis.index:
        iposfmis[msei] = mseiposfmis[mse]
# boolean indicating entries with non-NaN rel FMI & ipos FMI:
valid = ~np.isnan(iposfmis) & ~np.isnan(relfmis)
# scatter plot:
figsize = DEFAULTFIGURESIZE[0]*1.05, DEFAULTFIGURESIZE[1] # tweak to make space for labels
f, a = plt.subplots(figsize=figsize)
wintitle('rel FMI vs ipos stdev FMI')
a.scatter(iposfmis[valid], relfmis[valid], clip_on=False,
          marker='.', c='None', edgecolor='black', s=DEFSZ)
# plot regression:
fname = os.path.join('stats', 'figure_1_S2i_coefs.csv')
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([iposfmis[valid].min(), iposfmis[valid].max()])
    y = mm * x + b
    a.plot(x, y, '-', color='red') # plot linregress fit
a.set_ylabel("Reliability FMI")
a.set_xlabel("Eye position $\sigma$ 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:
for mseu, iposfmi, relfmi in zip(mseus, iposfmis, relfmis):
    iposmi.loc[mseu, 'stimtype'] = 'mvi'
    iposmi.loc[mseu, 'iposfmi'] = iposfmi
    iposmi.loc[mseu, 'relfmi'] = relfmi


"""Figure 5S1"""

## plot eye position stdev distributions:

## plot cross trial stdev PDFs for run and sit trials:
for stimtype, stimtypei in zip(STIMTYPESPLUSALL, STIMTYPESPLUSALLI):
    f, a = plt.subplots(figsize=DEFAULTFIGURESIZE)
    wintitle('ipos_st8 %s stdev cross trial PDF' % stimtype)
    stdrun = ipos_st8['std_xpos_cross'][:, stimtypei, 'run']
    stdsit = ipos_st8['std_xpos_cross'][:, stimtypei, 'sit']
    a.hist(stdrun, bins=iposedges, histtype='step', lw=1.5,
           color=st82clr['run'], alpha=1, label='Run')
    a.hist(stdsit, bins=iposedges, histtype='step', lw=1.5,
           color=st82clr['sit'], alpha=1, label='Sit')
    a.set_xlabel('Eye position $\sigma$ ($\mathregular{\degree}$)')
    a.set_ylabel('Experiment count')
    #a.legend(frameon=False)
    print('st8 %s std_cross:' % stimtype, wilcoxon(stdrun, stdsit))


## plot cross trial stdev CDFs for run and sit trials:
for stimtype, stimtypei in zip(STIMTYPESPLUSALL, STIMTYPESPLUSALLI):
    f, a = plt.subplots(figsize=DEFAULTFIGURESIZE)
    wintitle('ipos_st8 %s stdev cross trial CDF' % stimtype)
    stdrun = ipos_st8['std_xpos_cross'][:, stimtypei, 'run']
    stdsit = ipos_st8['std_xpos_cross'][:, stimtypei, 'sit']
    stdrunbins = list(np.unique(stdrun)) + [maxiposedge]
    stdsitbins = list(np.unique(stdsit)) + [maxiposedge]
    a.hist(stdrun, bins=stdrunbins, density=True, cumulative=True, histtype='step',
           lw=1.5, color=st82clr['run'], alpha=1, label='Run')
    a.hist(stdsit, bins=stdsitbins, density=True, cumulative=True, histtype='step',
           lw=1.5, color=st82clr['sit'], alpha=1, label='Sit')
    a.set_xlabel('Eye position $\sigma$ ($\mathregular{\degree}$)')
    a.set_ylabel('Cumulative probability')
    a.set_xlim(0, 8)
    a.set_xticks([0, 5])
    a.set_yticks([0, 0.5, 1])
    #a.legend(frameon=False)
    _, KS_p = ks_2samp(stdrun, stdsit)
    #_, MW_p = scipy.stats.mannwhitneyu(stdrun, stdsit)
    txt = '$\mathregular{p_{KS}=%.2g}$' % KS_p # prob that distribs are the same
    a.add_artist(AnchoredText(txt, loc='upper left', frameon=False))


## scatter plot movie reliability RMI vs eye position stdev FMI:
# feedback reliability RMI for each movie mseu:
mseus, relrmis = zip(*mviRMI['rel'][:, False].items()) # unpack indices and relrmi values
mseus, relrmis = np.array(mseus), np.array(relrmis)
# across-trial eye position stdev RMI for each experiment, 27 vals, 1 per exp:
mseiposrmis = ((ipos_st8['std_xpos_cross'][:, 'mvi', 'run']
                - ipos_st8['std_xpos_cross'][:, 'mvi', 'sit']) /
               (ipos_st8['std_xpos_cross'][:, 'mvi', 'run']
                + ipos_st8['std_xpos_cross'][:, 'mvi', 'sit']))
# get mse string of every mseu in mviRMI:
mses = np.array(['_'.join(mseu.split('_')[:-1]) for mseu in mseus])
# assign appropriate iposrmi to each mseu:
iposrmis = np.tile(np.nan, len(mseus)) # init to nans
for msei, mse in enumerate(mses):
    if mse in mseiposrmis.index:
        iposrmis[msei] = mseiposrmis[mse]
# boolean indicating entries with non-NaN rel RMI & ipos RMI:
valid = ~np.isnan(iposrmis) & ~np.isnan(relrmis)
# scatter plot:
figsize = DEFAULTFIGURESIZE[0]*1.05, DEFAULTFIGURESIZE[1] # tweak to make space for labels
f, a = plt.subplots(figsize=figsize)
wintitle('rel RMI vs ipos stdev RMI')
a.scatter(iposrmis[valid], relrmis[valid], clip_on=False,
          marker='.', c='None', edgecolor='black', s=DEFSZ)
# plot regression:
fname = os.path.join('stats', 'figure_5_S1i_coefs.csv')
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([iposrmis[valid].min(), iposrmis[valid].max()])
    y = mm * x + b
    a.plot(x, y, '-', color='red') # plot linregress fit
a.set_ylabel("Reliability RMI")
a.set_xlabel("Eye position $\sigma$ RMI")
ymin, ymax = -1, 1
#a.set_xlim(-0.1, 0.6)
a.set_ylim(ymin, ymax)
a.set_xticks([0, 0.5])
a.set_yticks([ymin, 0, ymax])
a.spines['left'].set_position(('outward', 4))
a.spines['bottom'].set_position(('outward', 4))
# save to iposmi dataframe:
for mseu, iposrmi, relrmi in zip(mseus, iposrmis, relrmis):
    iposmi.loc[mseu, 'stimtype'] = 'mvi' # automatically adds a new mseu row if needed
    iposmi.loc[mseu, 'iposrmi'] = iposrmi
    iposmi.loc[mseu, 'relrmi'] = relrmi