"""Figure 6 plots, use run -i fig6.py"""

allmseustrs = np.union1d(mvimseustrs, grtmseustrs)
modis = ['fmi', 'rmi', 'suppressionrmi', 'feedbackrmi', 'runfmi', 'sitfmi']
mi = pd.MultiIndex.from_product([allmseustrs, STIMTYPES, modis],
                                 names=['mseu', 'stimtype', 'mi'])
columns = ['meanrate', 'meanburstratio', 'spars', 'rel', 'meanpkw', 'snr']
fig6 = pd.DataFrame(index=mi, columns=columns)

stimtype2FMI = {'mvi': mviFMI, 'grt': grtFMI}
stimtype2RMI = {'mvi': mviRMI, 'grt': grtRMI}

loss = 'soft_l1' # loss function, less sensitive to outliers than ordinary least squares
f_scale = 0.25 # residual value at which to start treating points as outliers


# scatter plot run modulation index for feedback vs. suppression trials,
# for movies and gratings, for all measures:
figsize = DEFAULTFIGURESIZE
linmin, linmax, linstep = -1, 1, 1
ticks = np.arange(linmin, linmax+linstep, linstep)
for measure in modmeasuresnoblank:
    axislabel = measure2axislabel[measure]
    axislabel = short2longaxislabel.get(axislabel, axislabel)
    if axislabel.islower():
        axislabel = axislabel.capitalize()
    for stimtype in STIMTYPES:
        FMI = stimtype2FMI[stimtype]
        RMI = stimtype2RMI[stimtype]
        mseustrs = {'mvi':mvimseustrs, 'grt':grtmseustrs}[stimtype]
        mseu2exmpli = {'mvi':mvimseu2exmpli, 'grt':grtmseu2exmpli}[stimtype]
        stimtypelabel = stimtype2axislabel[stimtype]
        supprmis, fbrmis, exmplis, exmplmseustrs, normlis = [], [], [], [], []
        keptmseui = 0 # manually init and increment instead of using enumerate()
        for mseustr in mseustrs:
            supprmi = RMI[measure][mseustr, True]
            fbrmi = RMI[measure][mseustr, False]
            if np.isnan(supprmi) or np.isnan(fbrmi): # missing one or both mod values
                continue
            supprmis.append(supprmi)
            fbrmis.append(fbrmi)
            fig6.loc[mseustr, stimtype, 'suppressionrmi'][measure] = supprmi
            fig6.loc[mseustr, stimtype, 'feedbackrmi'][measure] = fbrmi
            if mseustr in mseu2exmpli:
                exmplis.append(keptmseui)
                exmplmseustrs.append(mseustr)
            else:
                normlis.append(keptmseui)
            keptmseui += 1 # manually increment
        supprmis = np.asarray(supprmis)
        fbrmis = np.asarray(fbrmis)
        if len(supprmis) == 0 or len(fbrmis) == 0:
            continue # nothing to plot
        f, a = plt.subplots(figsize=figsize)
        ## scatter plot feedback vs suppression RMI
        wintitle('fbsupp RMI %s %s' % (measure, stimtypelabel))
        # plot x=0 and y=0 lines:
        a.axhline(y=0, ls='--', marker='', color='lightgray', zorder=-np.inf)
        a.axvline(x=0, ls='--', marker='', color='lightgray', zorder=-np.inf)
        # plot y=x line:
        xyline = [linmin, linmax], [linmin, linmax]
        a.plot(xyline[0], xyline[1], '--', color='gray', zorder=-1)
        # plot normal (non-example) points:
        a.scatter(supprmis[normlis], fbrmis[normlis], clip_on=False,
                  marker='.', c='None', edgecolor='black', s=DEFSZ)
        # plot example points, one at a time:
        for exmpli, mseustr in zip(exmplis, exmplmseustrs):
            marker = exmpli2mrk[mseu2exmpli[mseustr]]
            c = exmpli2clr[mseu2exmpli[mseustr]]
            sz = exmpli2sz[mseu2exmpli[mseustr]]
            lw = exmpli2lw[mseu2exmpli[mseustr]]
            a.scatter(supprmis[exmpli], fbrmis[exmpli], clip_on=False,
                      marker=marker, c=c, s=sz, lw=lw)
        # get fname for relevant figure panels:
        relev_measures = ['meanrate', 'meanburstratio', 'spars', 'rel']
        if (measure in relev_measures) and stimtype == 'mvi':
            # if measure is relevant, get parameters for model from csv file:
            if measure == 'meanrate':
                fname = os.path.join('stats', 'figure_6a1_coefs.csv')
            elif measure == 'meanburstratio':
                fname = os.path.join('stats', 'figure_6a2_coefs.csv')
            elif measure == 'spars':
                fname = os.path.join('stats', 'figure_6a3_coefs.csv')
            elif measure == 'rel':
                fname = os.path.join('stats', 'figure_6a4_coefs.csv')
            try:
                df = pd.read_csv(fname)
                foundregression = True
            except FileNotFoundError:
                print('Missing file: %s' % fname)
                foundregression = False
            if foundregression:
                df = pd.read_csv(fname)
                mm = df['slope'][0]
                b = df['intercept'][0]
                x = np.array([np.nanmin(supprmis), np.nanmax(supprmis)])
                y = mm * x + b
                a.plot(x, y, '-', color='red') # plot linregress fit
        #a.set_title('%s RMI' % axislabel) # typeface too big by default, decreases panel size
        a.set_xlabel('Suppression')
        a.set_ylabel('Feedback')
        a.set_xlim(linmin, linmax)
        a.set_ylim(linmin, linmax)
        a.set_xticks(ticks)
        a.set_yticks(a.get_xticks()) # make log scale y ticks the same as x ticks
        a.minorticks_off()
        a.set_aspect('equal')
        a.spines['left'].set_position(('outward', 4))
        a.spines['bottom'].set_position(('outward', 4))


# scatter plot feedback modulation index for run vs. sit trials,
# for movies and gratings, for all measures:
figsize = DEFAULTFIGURESIZE
linmin, linmax, linstep = -1, 1, 1
ticks = np.arange(linmin, linmax+linstep, linstep)
for measure in modmeasuresnoblank:
    axislabel = measure2axislabel[measure]
    if axislabel.islower():
        axislabel = axislabel.capitalize()
    for stimtype in STIMTYPES:
        FMI = stimtype2FMI[stimtype]
        mseustrs = {'mvi':mvimseustrs, 'grt':grtmseustrs}[stimtype]
        mseu2exmpli = {'mvi':mvimseu2exmpli, 'grt':grtmseu2exmpli}[stimtype]
        stimtypelabel = stimtype2axislabel[stimtype]
        runfmis, sitfmis, exmplis, exmplmseustrs, normlis = [], [], [], [], []
        keptmseui = 0 # manually init and increment instead of using enumerate()
        for mseustr in mseustrs:
            runfmi = FMI[measure][mseustr, 'run']
            sitfmi = FMI[measure][mseustr, 'sit']
            if np.isnan(runfmi) or np.isnan(sitfmi): # missing one or both mod values
                continue
            runfmis.append(runfmi)
            sitfmis.append(sitfmi)
            fig6.loc[mseustr, stimtype, 'runfmi'][measure] = runfmi
            fig6.loc[mseustr, stimtype, 'sitfmi'][measure] = sitfmi
            if mseustr in mseu2exmpli:
                exmplis.append(keptmseui)
                exmplmseustrs.append(mseustr)
            else:
                normlis.append(keptmseui)
            keptmseui += 1 # manually increment
        runfmis = np.asarray(runfmis)
        sitfmis = np.asarray(sitfmis)
        if len(runfmis) == 0 or len(sitfmis) == 0:
            continue # nothing to plot
        f, a = plt.subplots(figsize=figsize)
        wintitle('runsit FMI %s %s' % (measure, stimtypelabel))
        # plot x=0 and y=0 lines:
        a.axhline(y=0, ls='--', marker='', color='lightgray', zorder=-np.inf)
        a.axvline(x=0, ls='--', marker='', color='lightgray', zorder=-np.inf)
        # plot y=x line:
        xyline = [linmin, linmax], [linmin, linmax]
        a.plot(xyline[0], xyline[1], '--', color='gray', zorder=-1)
        # plot normal (non-example) points:
        a.scatter(sitfmis[normlis], runfmis[normlis], clip_on=False,
                  marker='.', c='None', edgecolor='black', s=DEFSZ)
        # plot example points, one at a time:
        for exmpli, mseustr in zip(exmplis, exmplmseustrs):
            marker = exmpli2mrk[mseu2exmpli[mseustr]]
            c = exmpli2clr[mseu2exmpli[mseustr]]
            sz = exmpli2sz[mseu2exmpli[mseustr]]
            lw = exmpli2lw[mseu2exmpli[mseustr]]
            a.scatter(sitfmis[exmpli], runfmis[exmpli], clip_on=False,
                      marker=marker, c=c, s=sz, lw=lw)
        # get fname for relevant figure panels:
        if stimtype == 'mvi':
            relev_measures = ['meanrate', 'meanburstratio', 'spars', 'rel']
        elif stimtype == 'grt':
            relev_measures = ['meanrate', 'meanburstratio']
        if measure in relev_measures:
            # if measure is relevant, get parameters for model from csv file:
            if stimtype == 'mvi':
                if measure == 'meanrate':
                    fname = os.path.join('stats', 'figure_6b1_coefs.csv')
                elif measure == 'meanburstratio':
                    fname = os.path.join('stats', 'figure_6b2_coefs.csv')
                elif measure == 'spars':
                    fname = os.path.join('stats', 'figure_6b3_coefs.csv')
                elif measure == 'rel':
                    fname = os.path.join('stats', 'figure_6b4_coefs.csv')
            elif stimtype == 'grt':
                if measure == 'meanrate':
                    fname = os.path.join('stats', 'figure_6_S1b1_coefs.csv')
                elif measure == 'meanburstratio':
                    fname = os.path.join('stats', 'figure_6_S1b2_coefs.csv')
            try:
                df = pd.read_csv(fname)
                foundregression = True
            except FileNotFoundError:
                print('Missing file: %s' % fname)
                foundregression = False
            if foundregression:
                df = pd.read_csv(fname)
                mm = df['slope'][0]
                b = df['intercept'][0]
                x = np.array([np.nanmin(sitfmis), np.nanmax(sitfmis)])
                y = mm * x + b
                a.plot(x, y, '-', color='red') # plot linregress fit
        #a.set_title('%s FMI' % axislabel) # typeface too big by default, decreases panel size
        a.set_xlabel('Sit')
        a.set_ylabel('Run')
        a.set_xlim(linmin, linmax)
        a.set_ylim(linmin, linmax)
        a.set_xticks(ticks)
        a.set_yticks(a.get_xticks()) # make log scale y ticks the same as x ticks
        a.minorticks_off()
        a.set_aspect('equal')
        a.spines['left'].set_position(('outward', 4))
        a.spines['bottom'].set_position(('outward', 4))


# scatter plot feedback modulation index vs. run modulation index,
# for movies and gratings, for all measures:
figsize = DEFAULTFIGURESIZE
linmin, linmax, linstep = -1, 1, 1
ticks = np.arange(linmin, linmax+linstep, linstep)
for measure in modmeasuresnoblank:
    axislabel = measure2axislabel[measure]
    if axislabel.islower():
        axislabel = axislabel.capitalize()
    for stimtype in STIMTYPES:
        FMI = stimtype2FMI[stimtype]
        RMI = stimtype2RMI[stimtype]
        mseustrs = {'mvi':mvimseustrs, 'grt':grtmseustrs}[stimtype]
        mseu2exmpli = {'mvi':mvimseu2exmpli, 'grt':grtmseu2exmpli}[stimtype]
        stimtypelabel = stimtype2axislabel[stimtype]
        fmis, rmis, exmplis, exmplmseustrs, normlis = [], [], [], [], []
        keptmseui = 0 # manually init and increment instead of using enumerate()
        for mseustr in mseustrs:
            fmi = FMI[measure][mseustr, 'none'] # ignore run condition for FMI
            rmi = RMI[measure][mseustr, False] # use only feedback condition for RMI
            if np.isnan(fmi) or np.isnan(rmi): # missing one or both mod values
                continue
            fmis.append(fmi)
            rmis.append(rmi)
            fig6.loc[mseustr, stimtype, 'fmi'][measure] = fmi
            fig6.loc[mseustr, stimtype, 'rmi'][measure] = rmi
            if mseustr in mseu2exmpli:
                exmplis.append(keptmseui)
                exmplmseustrs.append(mseustr)
            else:
                normlis.append(keptmseui)
            keptmseui += 1 # manually increment
        fmis = np.asarray(fmis)
        rmis = np.asarray(rmis)
        if len(fmis) == 0 or len(rmis) == 0:
            continue # nothing to plot
        f, a = plt.subplots(figsize=figsize)
        wintitle('FMI vs RMI %s %s none False' % (measure, stimtypelabel))
        # plot x=0 and y=0 lines:
        a.axhline(y=0, ls='--', marker='', color='lightgray', zorder=-np.inf)
        a.axvline(x=0, ls='--', marker='', color='lightgray', zorder=-np.inf)
        # plot y=x line:
        xyline = [linmin, linmax], [linmin, linmax]
        a.plot(xyline[0], xyline[1], '--', color='gray', zorder=-1)
        # plot normal (non-example) points:
        a.scatter(rmis[normlis], fmis[normlis], clip_on=False,
                  marker='.', c='None', edgecolor='black', s=DEFSZ)
        # plot example points, one at a time:
        for exmpli, mseustr in zip(exmplis, exmplmseustrs):
            marker = exmpli2mrk[mseu2exmpli[mseustr]]
            c = exmpli2clr[mseu2exmpli[mseustr]]
            sz = exmpli2sz[mseu2exmpli[mseustr]]
            lw = exmpli2lw[mseu2exmpli[mseustr]]
            a.scatter(rmis[exmpli], fmis[exmpli], clip_on=False,
                      marker=marker, c=c, s=sz, lw=lw)
        # get fname for relevant figure panels:
        if stimtype == 'mvi':
            relev_measures = ['meanrate', 'meanburstratio', 'spars', 'rel']
        elif stimtype == 'grt':
            relev_measures = ['meanrate', 'meanburstratio']
        if measure in relev_measures:
            # if measure is relevant, get parameters for model from csv file:
            if stimtype == 'mvi':
                if measure == 'meanrate':
                    fname = os.path.join('stats', 'figure_6c1_coefs.csv')
                elif measure == 'meanburstratio':
                    fname = os.path.join('stats', 'figure_6c2_coefs.csv')
                elif measure == 'spars':
                    fname = os.path.join('stats', 'figure_6c3_coefs.csv')
                elif measure == 'rel':
                    fname = os.path.join('stats', 'figure_6c4_coefs.csv')
            elif stimtype == 'grt':
                if measure == 'meanrate':
                    fname = os.path.join('stats', 'figure_6_S1c1_coefs.csv')
                elif measure == 'meanburstratio':
                    fname = os.path.join('stats', 'figure_6_S1c2_coefs.csv')
            try:
                df = pd.read_csv(fname)
                foundregression = True
            except FileNotFoundError:
                print('Missing file: %s' % fname)
                foundregression = False
            if foundregression:
                df = pd.read_csv(fname)
                mm = df['slope'][0]
                b = df['intercept'][0]
                x = np.array([np.nanmin(rmis), np.nanmax(rmis)])
                y = mm * x + b
                a.plot(x, y, '-', color='red') # plot linregress fit
        #a.set_title('%s' % axislabel) # typeface too big by default, decreases panel size
        a.set_xlabel('RMI')
        a.set_ylabel('FMI')
        a.set_xlim(linmin, linmax)
        a.set_ylim(linmin, linmax)
        a.set_xticks(ticks)
        a.set_yticks(a.get_xticks())
        a.minorticks_off()
        a.set_aspect('equal')
        a.spines['left'].set_position(('outward', 4))
        a.spines['bottom'].set_position(('outward', 4))