"""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)

# redefine these here, in case saved data was loaded from pickles and these definitions
# later in calc.py were therefore skipped:
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 and dumbbell 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
        '''
        else:
            # calc/plot linregress line:
            # exclude extreme +/- 1 FMI and RMI values from linregress:
            non1is = (abs(supprmis) != 1) & (abs(fbrmis) != 1)
            supprmisnon1 = supprmis[non1is]
            fbrmisnon1 = fbrmis[non1is]
            # robust least squares:
            result = least_squares(linear_loss, (0, 0), args=(supprmisnon1, fbrmisnon1),
                                   loss=loss, f_scale=f_scale)
            mm, b = result['x']
            residuals = result['fun']
            #rsq = residual_rsquared(fbrmisnon1, residuals)
            #mm, b, rr, p, stderr = linregress(rmisnon1, fmisnon1)
            #rsq = rr * rr
            x = np.array([supprmisnon1.min(), supprmisnon1.max()])
        '''
        #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))
        #txt = ('$\mathregular{R^{2}=%.2g}$' % rsq)
        #.add_artist(AnchoredText(txt, loc='lower right', frameon=False))
        #txt = ('$\mathregular{p=%.2g}$\n'
        #       '$\mathregular{R^{2}=%.2g}$' % (p, rsq)) # prob that slope is 0
        #a.add_artist(AnchoredText(txt, loc='upper left', frameon=False))
        ## dumbbell plot feedback vs suppression RMI
        '''
        f, a = plt.subplots(figsize=DUMBBELLFIGSIZE)
        wintitle('fbsupp RMI dumbbell %s %s' % (measure, stimtypelabel))
        # plot x=0 line:
        a.axvline(x=0, ls='--', marker='', color='lightgray', zorder=-np.inf)
        # plot horizontal lines:
        sortis = fbrmis.argsort() # y vals for hlines
        npairs = len(fbrmis)
        ys = np.arange(npairs)
        srtdfbrmis, srtdsupprmis = fbrmis[sortis], supprmis[sortis]
        a.hlines(ys, srtdfbrmis, srtdsupprmis, zorder=-100)
        sz = 50
        # plot supprmis points:
        a.scatter(srtdsupprmis, ys, marker='.', c='white', edgecolor='black', s=sz)
        # plot fbrmis points:
        a.scatter(srtdfbrmis, ys, marker='.', c='black', s=sz)
        a.set_xlabel('%s RMI' % axislabel)
        a.set_ylabel('Neurons')
        a.set_xlim(linmin, linmax)
        a.set_xticks(ticks)
        a.set_yticks([])
        '''


# 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
        '''
        else:
            # calc/plot linregress line:
            # exclude extreme +/- 1 FMI values from linregress:
            non1is = (abs(sitfmis) != 1) & (abs(runfmis) != 1)
            runfmisnon1 = runfmis[non1is]
            sitfmisnon1 = sitfmis[non1is]
            # robust least squares:
            result = least_squares(linear_loss, (0, 0), args=(sitfmisnon1, runfmisnon1),
                                   loss=loss, f_scale=f_scale)
            mm, b = result['x']
            residuals = result['fun']
            #rsq = residual_rsquared(runfmisnon1, residuals)
            #mm, b, rr, p, stderr = linregress(sitfmisnon1, runfmisnon1)
            #rsq = rr * rr
            x = np.array([sitfmisnon1.min(), sitfmisnon1.max()])
        '''
        #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))
        #txt = ('$\mathregular{R^{2}=%.2g}$' % rsq)
        #a.add_artist(AnchoredText(txt, loc='lower right', frameon=False))
        #txt = ('$\mathregular{p=%.2g}$\n'
        #       '$\mathregular{R^{2}=%.2g}$' % (p, rsq)) # prob that slope is 0
        #a.add_artist(AnchoredText(txt, loc='upper left', frameon=False))
        ## dumbbell plot run vs sit FMI
        '''
        f, a = plt.subplots(figsize=DUMBBELLFIGSIZE)
        wintitle('runsit FMI dumbbell %s %s' % (measure, stimtypelabel))
        # plot x=0 line:
        a.axvline(x=0, ls='--', marker='', color='lightgray', zorder=-np.inf)
        # plot horizontal lines:
        sortis = runfmis.argsort() # y vals for hlines
        npairs = len(runfmis)
        ys = np.arange(npairs)
        srtdrunfmis, srtdsitfmis = runfmis[sortis], sitfmis[sortis]
        a.hlines(ys, srtdrunfmis, srtdsitfmis, zorder=-100)
        sz = 50
        # plot sitfmis points:
        a.scatter(srtdsitfmis, ys, marker='.', c='white', edgecolor='black', s=sz)
        # plot runfmis points:
        a.scatter(srtdrunfmis, ys, marker='.', c='black', s=sz)
        a.set_xlabel('%s FMI' % short2longaxislabel.get(axislabel, axislabel))
        a.set_ylabel('Neurons')
        a.set_xlim(linmin, linmax)
        a.set_xticks(ticks)
        a.set_yticks([])
        '''


# 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
            """
            TODO: plot RMI calc'd across both opto conditions, but before doing that, you
            have to normalize by the relative mean measure (say rate) of the two opto
            conditions, and weight them appropriately, says Jackson Smith from Oxford. Should
            really do the equivalent for FMI, especially cuz there are not equal numbers of
            trials for run vs sit, but the analysis that plots FMI as a function of run vs sit
            trials more or less shows that you can pool across run condition to calculate 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
        '''
        else:
            # calc/plot linregress line:
            # exclude extreme +/- 1 FMI and RMI values from linregress:
            non1is = (abs(rmis) != 1) & (abs(fmis) != 1)
            rmisnon1 = rmis[non1is]
            fmisnon1 = fmis[non1is]
            # robust least squares:
            result = least_squares(linear_loss, (0, 0), args=(rmisnon1, fmisnon1),
                                   loss=loss, f_scale=f_scale)
            mm, b = result['x']
            residuals = result['fun']
            rsq = residual_rsquared(fmisnon1, residuals)
            mm, b, rr, p, stderr = linregress(rmisnon1, fmisnon1)
            rsq = rr * rr
            x = np.array([rmisnon1.min(), rmisnon1.max()])
        '''
        #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))
        #txt = ('$\mathregular{R^{2}=%.2g}$' % rsq)
        #a.add_artist(AnchoredText(txt, loc='lower right', frameon=False))
        #txt = ('$\mathregular{p=%.2g}$\n'
        #       '$\mathregular{R^{2}=%.2g}$' % (p, rsq)) # prob that slope is 0
        #a.add_artist(AnchoredText(txt, loc='upper left', frameon=False))