import numpy as np
import pandas as pd
import seaborn as sns
from matplotlib import rc
import matplotlib.pyplot as plt

ROOT_PATH = 'data/'
EXTENSION = '.csv'
FONT_SIZE = 16
BINS = np.linspace(-1, 1, 5)

COLORS = ['steelblue', 'darkorange', 'green']


def VaR(values, alpha):
    """
    Compute the value at risk at alpha %.
    """
    n = len(values)
    b = int(n * alpha)
    values = np.sort(values)
    var = values[b]
    return var


def CVaR(values, alpha):
    """
    Compute the conditional value at risk at alpha %.
    """
    n = len(values)
    b = int(n * alpha)
    values = np.sort(values)
    return values[:b].mean()


def file_name(env_name, agent_name):
    return ROOT_PATH + env_name + '-' + agent_name + EXTENSION


def label_from_agent_name(agent_name):
    if agent_name == 'AsynDP-NSMDP':
        return 'DP-NSMDP'
    elif agent_name == 'AsynDP-snapshot':
        return 'DP-snapshot'
    else:
        return agent_name


def confidence_interval(v, p=0.7):
    m = v.mean()
    v = np.sort(v)
    n = len(v)
    perc = (1 - p) / 2
    return (abs(m - v[int(n * perc)]), abs(m - v[int(n * (1 - perc))]))


def analysis(env_name, agent_name_pool, alpha, doplot=True):
    print()
    values = []
    labels = []
    for agent_name in agent_name_pool:
        f = file_name(env_name, agent_name)
        df = pd.read_csv(f)
        discounted_return = df.discounted_return
        print(env_name,agent_name)
        print('    mean                      :', discounted_return.mean())
        print('    median                    :', discounted_return.median())
        print('    standard deviation        :', discounted_return.std())
        print('    VaR                       :', VaR(discounted_return.values, alpha), '(alpha =', alpha, ')')
        print('    CVaR (expected shortfall) :', CVaR(discounted_return.values, alpha), '(alpha =', alpha, ')')
        # plt.hist(discounted_return.values, BINS, alpha=0.7, label=agent_name)
        values.append(discounted_return.values)
        labels.append(agent_name)
    plt.hist(values, BINS, alpha=0.7, label=labels)
    if doplot:
        plt.legend(loc='upper right')
        plt.title(env_name)
        plt.xlabel('Discounted return')
        plt.ylabel('Frequency')
        plt.show()


def triptych(envs, agent_name_pool):
    titles = [r'$\epsilon = 0$', r'$\epsilon = 0.5$', r'$\epsilon = 1$']
    assert len(envs) == 3
    fig, (ax1, ax2, ax3) = plt.subplots(1, 3, sharey=True, figsize=(9, 4), facecolor='w', edgecolor='k')
    axes = [ax1, ax2, ax3]
    for i in range(3):
        values = []
        labels = []
        for agent_name in agent_name_pool:
            f = file_name(envs[i], agent_name)
            df = pd.read_csv(f)
            discounted_return = df.discounted_return
            # axes[i].hist(discounted_return.values, BINS, alpha=0.7, label=label_from_agent_name(agent_name))
            values.append(discounted_return.values)
            labels.append(label_from_agent_name(agent_name))

        axes[i].grid(True)  # , linestyle='--')
        # axes[i].xaxis.grid(True)

        axes[i].set_ylim((0, 3.5))
        # axes[i].hist(values, BINS, alpha=0.7, label=labels)
        '''
        for j in [0, 1, 2]:
            # sns.distplot(values[j], ax=axes[i], label=labels[j], kde=True, hist=False, kde_kws={'shade': True, 'linewidth': 1})
            sns.kdeplot(values[j], ax=axes[i], label=labels[j], bw=0.09, shade=True, legend=False, color=COLORS[j])
        '''
        sns.kdeplot(values[1], ax=axes[i], label=labels[1], bw=0.1, shade=True, legend=False, color=COLORS[1])
        sns.kdeplot(values[2], ax=axes[i], label=labels[2], bw=0.1, shade=True, legend=False, color=COLORS[2])
        sns.kdeplot(values[0], ax=axes[i], label=labels[0], bw=0.1, shade=False, legend=False, color=COLORS[0],
                    marker='o', markevery=5)

        axes[i].set_title(titles[i], fontsize=FONT_SIZE)
        axes[i].set_xlabel(r'Discounted return', fontsize=FONT_SIZE)
        axes[0].set_ylabel(r'Density', fontsize=FONT_SIZE)
        axes[i].set_xlim((-1.0, 1.0))
    fig.subplots_adjust(bottom=0.3, wspace=0.2)
    ax2.legend(loc='upper center', fontsize=FONT_SIZE, bbox_to_anchor=(0.5, -0.2), frameon=True, fancybox=False,
               shadow=False, ncol=3)
    plt.show()


def plot_return_vs_epsilon(envs, epsilons, agent_name_pool):
    stdeb = True
    fmts = ['o', '^', 's']
    linestyles = ['-', '--', ':']
    '''
    legend_elements = [
        Line2D([0], [0], linestyle='-', color='steelblue', lw=4, label='DP-NSMDP'),
        Line2D([0], [0], linestyle='--', color='darkorange', lw=4, label='DP-snapshot'),
        Line2D([0], [0], linestyle=':', color='green', lw=4, label='DP-NSMDP')
    ]
    '''
    plt.figure(figsize=(7,4))
    for i in range(len(agent_name_pool)):
        agent_name = agent_name_pool[i]
        means = []
        stds = []
        errinf = []
        errsup = []
        for j in range(len(envs)):
            f = file_name(envs[j], agent_name)
            df = pd.read_csv(f)
            discounted_return = df.discounted_return
            means.append(discounted_return.mean())
            stds.append(discounted_return.std())
            if stdeb:
                inf, sup = 0.5 * discounted_return.std(), 0.5 * discounted_return.std()
            else:
                inf, sup = confidence_interval(discounted_return)
            errinf.append(inf)
            errsup.append(sup)
        plt.errorbar(epsilons, means, yerr=[errinf, errsup], fmt=fmts[i], color=COLORS[i],
                     label=label_from_agent_name(agent_name))
        m = np.array(means)
        up = np.add(m, errsup)
        dw = np.add(m, [-e for e in errinf])
        plt.fill_between(epsilons, up, dw, color=COLORS[i], alpha=0.1)
        plt.plot(epsilons, up, color=COLORS[i], linestyle=linestyles[i])
        plt.plot(epsilons, dw, color=COLORS[i], linestyle=linestyles[i])
    plt.subplots_adjust(bottom=0.3, wspace=0.2)
    plt.legend(loc='upper center', fontsize=FONT_SIZE, bbox_to_anchor=(0.5, -0.2), frameon=True, fancybox=False,
               shadow=False, ncol=3)
    plt.xlabel(r'$\epsilon$', fontsize=FONT_SIZE)
    plt.ylabel('Discounted return', fontsize=FONT_SIZE)
    plt.ylim((-1.0, 1.0))
    plt.show()


def results_exploitation():
    rc('font', **{'family': 'sans-serif', 'sans-serif': ['Helvetica']})
    # for Palatino and other serif fonts use:
    # rc('font',**{'family':'serif','serif':['Palatino']})
    plt.rc('text', usetex=True)
    plt.rc('font', family='serif')

    #  Parameters
    alpha = 0.05 # Percentage value at risk
    epsilons = [0.0, 0.25, 0.5, 0.75, 1.0]
    env_name_pool = ['NSBridge-v0-' + str(epsilons[i]) for i in range(len(epsilons))]
    agent_name_pool = ['AsynDP-NSMDP', 'AsynDP-snapshot', 'RATS']

    plot_return_vs_epsilon(env_name_pool, epsilons, agent_name_pool)

    epsilons = [0.0, 0.5, 1.0]
    env_name_pool = ['NSBridge-v0-' + str(epsilons[i]) for i in range(len(epsilons))]
    triptych(env_name_pool, agent_name_pool)
    for env_name in env_name_pool:
        analysis(env_name, agent_name_pool, alpha)


if __name__ == '__main__':
    np.random.seed(199311)
    results_exploitation()
