import matplotlib.pyplot as plt
import matplotlib as mpl
import matplotlib.patches as patches
import numpy as np

mpl.rcParams['figure.dpi'] = 300
mpl.rcParams.update({'font.size': 10})

'''
Graphs.py

A file containing a number of functions for generating/reproducing 
the experimental graphs of FRLC.
'''

def whisker_plot(x, y_bot, y_diff, ylabel="", xlabel="", title="", \
                 errbarlabel = r"$5^{th}$, $95^{th}$ percentiles", \
                 centerlabel=r"$50^{th}$ percentile", \
                line=True):
    if line:
        print(np.zeros_like(y_bot).shape)
        print(y_diff.shape)
        plt.errorbar(x, y_bot, yerr=(np.zeros_like(y_bot), y_diff), capsize=1, \
                     ecolor='black', ls='', lw=1, capthick=1, \
                     label=errbarlabel, alpha=0.6)
        plt.plot(x, y_bot + (1/2)*y_diff, marker='d', \
                 color='black', linestyle='--', \
                 label=centerlabel, alpha=0.6)
    else:
        plt.errorbar(x, y_bot, yerr=(np.zeros_like(y_bot), y_diff), capsize=1, \
                     ecolor='black', ls='', lw=1, capthick=1, \
                     label=errbarlabel, alpha=0.6)
        plt.plot(x, y_bot + (1/2)*y_diff, marker='d', color='black', \
                 ls='', label=centerlabel, alpha=0.6)
    
    plt.title(title)
    plt.xlabel(xlabel)
    plt.ylabel(ylabel)
    plt.legend()
    plt.show()
    return

def whisker_plot_combined(x, y_botA, y_diffA, y_botB, y_diffB, \
                          ylabel="", xlabel="", title="", \
                 errbarlabelA = r"$5^{th}$, $95^{th}$ percentiles", \
                 centerlabelA = r"$50^{th}$ percentile", \
                errbarlabelB = r"$5^{th}$, $95^{th}$ percentiles", \
                 centerlabelB = r"$50^{th}$ percentile", fighandle=None):
    plt.rcParams['text.usetex'] = True
    plt.rc('text.latex', preamble=r'\usepackage{amsmath}')
    # Plot A
    plt.errorbar(x, y_botA, yerr=(np.zeros_like(y_botA), y_diffA), capsize=1, \
                 ecolor='blue', ls='', lw=1, capthick=1, \
                 label=errbarlabelA, alpha=0.6)
    plt.plot(x, y_botA + (1/2)*y_diffA, marker='d', \
             color='blue', linestyle='--', \
             label=centerlabelA, alpha=0.6)
    # Plot B
    plt.errorbar(x, y_botB, yerr=(np.zeros_like(y_botB), y_diffB), capsize=1, \
                 ecolor='red', ls='', lw=1, capthick=1, \
                 label=errbarlabelB, alpha=0.6)
    plt.plot(x, y_botB + (1/2)*y_diffB, marker='d', \
             color='red', linestyle='--', \
             label=centerlabelB, alpha=0.6)
    
    plt.title(title)
    plt.xlabel(xlabel)
    plt.ylabel(ylabel)
    plt.legend()
    if fighandle is None:
        plt.show()
    else:
        plt.savefig(fighandle, transparent=True, bbox_inches="tight")
    return

def whisker_plot_combined_All(x, y_bots, y_diffs, errbarlabels, centerlabels, \
                           ylabel="", xlabel="", title="", \
                           fighandle=None, colors = ['orange', 'blue', 'green', 'red', 'gray']):
    
    plt.rcParams['text.usetex'] = True
    plt.rc('text.latex', preamble=r'\usepackage{amsmath}')
    for i in range(len(y_bots)):
        y_bot = y_bots[i]
        y_diff = y_diffs[i]
        errbarlabel = errbarlabels[i]
        centerlabel = centerlabels[i]
        color = colors[i]
        # Plot A
        plt.errorbar(x, y_bot, yerr=(np.zeros_like(y_bot), y_diff), capsize=1, \
                     ecolor=color, ls='', lw=1, capthick=1, \
                     label=errbarlabel, alpha=0.6)
        plt.plot(x, y_bot + (1/2)*y_diff, marker='d', \
                 color=color, linestyle='--', \
                 label=centerlabel, alpha=0.6)
    plt.title(title)
    plt.xlabel(xlabel)
    plt.ylabel(ylabel)
    plt.legend()
    if fighandle is None:
        plt.show()
    else:
        plt.savefig(fighandle, transparent=True, bbox_inches="tight")
    return