"""Figure 1S3 and 3S1 unit classification plots, use run -i fig1S33S1.py"""

index = pd.Index(mvigrtmsustrs, name='msu')
columns = ['mvi_meanrate', 'mvi_meanburstratio', 'grt_meanrate', 'grt_meanburstratio',
           'sbc', 'depth', 'normdepth', 'dsi', 'rfdist',
           'mvi_meanrate_raw', 'mvi_meanburstratio_raw',
           'grt_meanrate_raw', 'grt_meanburstratio_raw']
fig1S33S1 = pd.DataFrame(index=index, columns=columns)

from scipy.stats import gmean

import matplotlib as mpl
from matplotlib.patches import Rectangle

TRANSTHRESH = 0.2
ONOFFTHRESH = 0.2


# stripplot FMI by SbC/non-SbC, mvi and grt, meanrate and meanburstratio:
np.random.seed(0) # to get identical horizontal jitter in strip plots on every run
figsize = DEFAULTFIGURESIZE
for stimtype in STIMTYPES:
    stimtypelabel = stimtype2axislabel[stimtype]
    for measure in ['meanrate', 'meanburstratio']:
        axislabel = measure2axislabel[measure]
        axislabel = short2longaxislabel.get(axislabel, axislabel)
        if axislabel.islower():
            axislabel = axislabel.capitalize()
        fmis, sbcs = [], []
        colname = '_'.join([stimtype, measure]) # e.g. 'mvi_meanrate'
        for msustr in mvigrtmsustrs:
            # save sbc to fig1S33S1 df, regardless of FMI val, will be overwritten multiple
            # times with the same value, but that's OK:
            sbc = celltype.loc[msustr]['sbc'] # bool
            fig1S33S1.loc[msustr]['sbc'] = sbc
            fmi = maxFMI.loc[msustr, 'none', stimtype][measure] # ignore run condition for FMI
            if np.isnan(fmi):
                continue
            fmis.append(fmi)
            sbcs.append(sbc)
            fig1S33S1.loc[msustr][colname] = fmi
        fmis = np.asarray(fmis)
        sbcs = np.asarray(sbcs)
        sbcfmis = fmis[sbcs == True]
        nonsbcfmis = fmis[sbcs == False]
        f, a = plt.subplots(figsize=figsize)
        wintitle('FMI SbC %s %s strip' % (stimtypelabel, measure))
        # plot y=0 line:
        a.axhline(y=0, ls='--', marker='', color='lightgray', zorder=-np.inf)
        data = pd.DataFrame.from_dict({'SbC':sbcfmis, 'Non-SbC':nonsbcfmis},
                                      orient='index').transpose()
        sns.stripplot(ax=a, data=data, clip_on=False, marker='.',
                      color='None', edgecolor='black', size=np.sqrt(50))
        # get fname of appropriate LMM .cvs file:
        fname = None # sanity check: clear from previous loop
        if stimtype == 'mvi':
            if measure == 'meanrate':
                fname = os.path.join('stats', 'figure_1_S3a_pred_means.csv')
            elif measure == 'meanburstratio':
                fname = os.path.join('stats', 'figure_1_S3f_pred_means.csv')
        elif stimtype == 'grt':
            if measure == 'meanrate':
                fname = os.path.join('stats', 'figure_3_S1a_pred_means.csv')
            elif measure == 'meanburstratio':
                fname = os.path.join('stats', 'figure_3_S1f_pred_means.csv')
        # fetch LMM means from .csv:
        df = pd.read_csv(fname)
        meannonsbcfmi = df['non_sbc'][0]
        meansbcfmi = df['sbc'][0]
        # plot mean with short horizontal lines:
        a.plot([-0.25, 0.25], [meansbcfmi, meansbcfmi], '-', lw=2, c='red', zorder=np.inf)
        a.plot([0.75, 1.25], [meannonsbcfmi, meannonsbcfmi], '-', lw=2, c='red', zorder=np.inf)
        a.set_ylabel('%s FMI' % axislabel)
        a.set_ylim(-1, 1)
        a.set_yticks([-1, 0, 1])
        a.tick_params(bottom=False)
        a.spines['bottom'].set_position(('outward', 5))
        a.spines['bottom'].set_visible(False)


# scatter plot FMI vs. depth by mvi and grt, meanrate and meanburstratio:
figsize = DEFAULTFIGURESIZE
for stimtype in STIMTYPES:
    stimtypelabel = stimtype2axislabel[stimtype]
    for measure in ['meanrate', 'meanburstratio']:
        axislabel = measure2axislabel[measure]
        axislabel = short2longaxislabel.get(axislabel, axislabel)
        if axislabel.islower():
            axislabel = axislabel.capitalize()
        fmis, depths = [], []
        colname = '_'.join([stimtype, measure]) # e.g. 'mvi_meanrate'
        for msustr in mvigrtmsustrs:
            # save depth to fig1S33S1 df, regardless of FMI val, will be overwritten multiple
            # times with the same value, but that's OK:
            depth = celltype.loc[msustr]['depth'] # float
            fig1S33S1.loc[msustr]['depth'] = depth
            fmi = maxFMI.loc[msustr, 'none', stimtype][measure] # ignore run condition for FMI
            if np.isnan(fmi):
                continue
            fmis.append(fmi)
            depths.append(depth)
            fig1S33S1.loc[msustr][colname] = fmi # might overwrite w/ identical values
        fmis = np.asarray(fmis)
        depths = np.asarray(depths)
        f, a = plt.subplots(figsize=figsize)
        wintitle('FMI depth %s %s' % (stimtypelabel, measure))
        # plot y=0 line:
        a.axhline(y=0, ls='--', marker='', color='lightgray', zorder=-np.inf)
        a.scatter(depths, fmis, clip_on=False, marker='.', c='None', edgecolor='k', s=DEFSZ)
        # get fname of appropriate LMM .cvs file:
        fname = None # sanity check: clear from previous loop
        if stimtype == 'mvi':
            if measure == 'meanrate':
                fname = os.path.join('stats', 'figure_1_S3b_coefs.csv')
            elif measure == 'meanburstratio':
                fname = os.path.join('stats', 'figure_1_S3g_coefs.csv')
        elif stimtype == 'grt':
            if measure == 'meanrate':
                fname = os.path.join('stats', 'figure_3_S1b_coefs.csv')
            elif measure == 'meanburstratio':
                fname = os.path.join('stats', 'figure_3_S1g_coefs.csv')
        # fetch LMM linregress fit params from .csv:
        df = pd.read_csv(fname)
        mm = df['slope'][0]
        b = df['intercept'][0]
        x = np.array([np.nanmin(depths), np.nanmax(depths)])
        y = mm * x + b
        a.plot(x, y, '-', color='red') # plot linregress fit
        a.set_xlabel('Depth ($\mathregular{\mu}$m)')
        a.set_ylabel('%s FMI' % axislabel)
        a.set_xlim(0, 500)
        a.set_ylim(-1, 1)
        a.set_yticks([-1, 0, 1])
        a.spines['left'].set_position(('outward', 4))
        a.spines['bottom'].set_position(('outward', 4))


# scatter plot FMI vs. DSI by mvi and grt, meanrate and meanburstratio:
figsize = DEFAULTFIGURESIZE
for stimtype in STIMTYPES:
    stimtypelabel = stimtype2axislabel[stimtype]
    for measure in ['meanrate', 'meanburstratio']:
        axislabel = measure2axislabel[measure]
        axislabel = short2longaxislabel.get(axislabel, axislabel)
        if axislabel.islower():
            axislabel = axislabel.capitalize()
        fmis, dsis = [], []
        colname = '_'.join([stimtype, measure]) # e.g. 'mvi_meanrate'
        for msustr in mvigrtmsustrs:
            # save DSI to fig1S33S1 df, regardless of FMI val, will be overwritten multiple
            # times with the same value, but that's OK:
            dsi = celltype.loc[msustr]['dsi'] # float
            fig1S33S1.loc[msustr]['dsi'] = dsi
            fmi = maxFMI.loc[msustr, 'none', stimtype][measure] # ignore run condition for FMI
            if np.isnan(fmi):
                continue
            fmis.append(fmi)
            dsis.append(dsi)
            fig1S33S1.loc[msustr][colname] = fmi # might overwrite w/ identical values
        fmis = np.asarray(fmis)
        dsis = np.asarray(dsis)
        f, a = plt.subplots(figsize=figsize)
        wintitle('FMI DSI %s %s' % (stimtypelabel, measure))
        # plot y=0 line:
        a.axhline(y=0, ls='--', marker='', color='lightgray', zorder=-np.inf)
        a.scatter(dsis, fmis, clip_on=False, marker='.', c='None', edgecolor='k', s=DEFSZ)
        # get fname of appropriate LMM .cvs file:
        fname = None # sanity check: clear from previous loop
        if stimtype == 'mvi':
            if measure == 'meanrate':
                fname = os.path.join('stats', 'figure_1_S3c_coefs.csv')
            elif measure == 'meanburstratio':
                fname = os.path.join('stats', 'figure_1_S3h_coefs.csv')
        elif stimtype == 'grt':
            if measure == 'meanrate':
                fname = os.path.join('stats', 'figure_3_S1c_coefs.csv')
            elif measure == 'meanburstratio':
                fname = os.path.join('stats', 'figure_3_S1h_coefs.csv')
        # fetch LMM linregress fit params from .csv:
        df = pd.read_csv(fname)
        mm = df['slope'][0]
        b = df['intercept'][0]
        x = np.array([np.nanmin(dsis), np.nanmax(dsis)])
        y = mm * x + b
        a.plot(x, y, '-', color='red') # plot linregress fit
        a.set_xlabel('DSI')
        a.set_ylabel('%s FMI' % axislabel)
        a.set_xlim(0, 1)
        a.set_ylim(-1, 1)
        a.set_yticks([-1, 0, 1])
        a.spines['left'].set_position(('outward', 4))
        a.spines['bottom'].set_position(('outward', 4))


# scatter plot FMI vs. distance of MUA envl RF from screen center, by mvi and grt,
# meanrate and meanburstratio:
figsize = DEFAULTFIGURESIZE
for stimtype in STIMTYPES:
    stimtypelabel = stimtype2axislabel[stimtype]
    for measure in ['meanrate', 'meanburstratio']:
        axislabel = measure2axislabel[measure]
        axislabel = short2longaxislabel.get(axislabel, axislabel)
        if axislabel.islower():
            axislabel = axislabel.capitalize()
        fmis, ds = [], []
        colname = '_'.join([stimtype, measure]) # e.g. 'mvi_meanrate'
        for msustr in mvigrtmsustrs:
            # save rfdist to fig1S33S1 df, regardless of FMI val, will be overwritten multiple
            # times with the same value, but that's OK:
            x0, y0 = cellscreenpos.loc[msustr]
            d = np.sqrt(x0**2 + y0**2) # distance from screen center, deg
            fig1S33S1.loc[msustr]['rfdist'] = d
            fmi = maxFMI.loc[msustr, 'none', stimtype][measure] # ignore run condition for FMI
            if np.isnan(fmi):
                continue
            fmis.append(fmi)
            ds.append(d)
            fig1S33S1.loc[msustr][colname] = fmi # might overwrite w/ identical values
        fmis = np.asarray(fmis)
        ds = np.asarray(ds)
        f, a = plt.subplots(figsize=figsize)
        wintitle('FMI rfdist %s %s' % (stimtypelabel, measure))
        # plot y=0 line:
        a.axhline(y=0, ls='--', marker='', color='lightgray', zorder=-np.inf)
        a.scatter(ds, fmis, clip_on=False, marker='.', c='None', edgecolor='k', s=DEFSZ)
        # get fname of appropriate LMM .cvs file:
        fname = None # sanity check: clear from previous loop
        if stimtype == 'mvi':
            if measure == 'meanrate':
                fname = os.path.join('stats', 'figure_1_S3d_coefs.csv')
            elif measure == 'meanburstratio':
                fname = os.path.join('stats', 'figure_1_S3i_coefs.csv')
        elif stimtype == 'grt':
            if measure == 'meanrate':
                fname = os.path.join('stats', 'figure_3_S1d_coefs.csv')
            elif measure == 'meanburstratio':
                fname = os.path.join('stats', 'figure_3_S1i_coefs.csv')
        # fetch LMM linregress fit params from .csv:
        df = pd.read_csv(fname)
        mm = df['slope'][0]
        b = df['intercept'][0]
        x = np.array([np.nanmin(ds), np.nanmax(ds)])
        y = mm * x + b
        a.plot(x, y, '-', color='red') # plot linregress fit
        a.set_xlabel('RF dist. from center ($\degree$)')
        a.set_ylabel('%s FMI' % axislabel)
        a.set_xlim(0, 40)
        a.set_ylim(-1, 1)
        a.set_yticks([-1, 0, 1])
        a.spines['left'].set_position(('outward', 4))
        a.spines['bottom'].set_position(('outward', 4))


# scatter plot FMI vs raw measure, for rate and burst ratio during control condition,
# by mvi and grt:
stimtype2resp = {'mvi':mviresp, 'grt':grtresp}
figsize = DEFAULTFIGURESIZE
for stimtype in STIMTYPES:
    stimtypelabel = stimtype2axislabel[stimtype]
    resp = stimtype2resp[stimtype]
    if stimtype == 'mvi':
        resp = resp.xs('nat', level='kind') # dereference movie 'kind' index level
    for measure in ['meanrate', 'meanburstratio']:
        axislabel = measure2axislabel[measure]
        axislabel = short2longaxislabel.get(axislabel, axislabel)
        axisunits = measure2axisunits.get(measure, '')
        if axislabel.islower():
            axislabel = axislabel.capitalize()
        fmis, msrs = [], []
        fmicolname = '_'.join([stimtype, measure]) # e.g. 'mvi_meanrate'
        msrcolname = '_'.join([stimtype, measure, 'raw']) # e.g. 'mvi_meanrate_raw'
        for msustr in mvigrtmsustrs:
            mseustr, fmi = maxFMI.loc[msustr, 'none', stimtype][['mseu', measure]]
            if pd.isna(mseustr) or pd.isna(fmi):
                continue
            msr = resp.loc[mseustr, 'none', False][measure]
            fmis.append(fmi)
            msrs.append(msr)
            fig1S33S1.loc[msustr][fmicolname] = fmi # might overwrite w/ identical values
            fig1S33S1.loc[msustr][msrcolname] = msr
        fmis, msrs = np.asarray(fmis), np.asarray(msrs)
        ## scatter plot FMI vs raw measure:
        f, a = plt.subplots(figsize=figsize)
        wintitle('FMI raw %s %s' % (stimtypelabel, measure))
        # plot y=0 line:
        a.axhline(y=0, ls='--', marker='', color='lightgray', zorder=-np.inf)
        a.scatter(msrs, fmis, clip_on=False, marker='.', c='None', edgecolor='k', s=DEFSZ)
        # get fname of appropriate LMM .cvs file:
        fname = None # sanity check: clear from previous loop
        if stimtype == 'mvi':
            if measure == 'meanrate':
                fname = os.path.join('stats', 'figure_1_S3e_coefs.csv')
            elif measure == 'meanburstratio':
                fname = os.path.join('stats', 'figure_1_S3j_coefs.csv')
        elif stimtype == 'grt':
            if measure == 'meanrate':
                fname = os.path.join('stats', 'figure_3_S1e_coefs.csv')
            elif measure == 'meanburstratio':
                fname = os.path.join('stats', 'figure_3_S1j_coefs.csv')
        # fetch LMM linregress fit params from .csv:
        df = pd.read_csv(fname)
        mm = df['slope'][0]
        b = df['intercept'][0]
        x = np.array([np.nanmin(msrs), np.nanmax(msrs)])
        y = mm * x + b
        a.plot(x, y, '-', color='red') # plot linregress fit
        a.set_xlabel('%s' % axislabel+axisunits)
        a.set_ylabel('%s FMI' % axislabel)
        a.set_xlim(xmin=0)
        a.set_ylim(-1, 1)
        #a.set_xticks(ticks)
        a.set_yticks([-1, 0, 1])
        a.spines['left'].set_position(('outward', 4))
        a.spines['bottom'].set_position(('outward', 4))