import os
from argparse import ArgumentParser

import seaborn as sns
import matplotlib.pyplot as plt
import pandas as pd
import numpy as np
import pickle
from tensorboard.backend.event_processing.event_accumulator import EventAccumulator

from vis_paper_test import get_framework_name


seven_tab10_colors = [
    (0.12156862745098039, 0.4666666666666667, 0.7058823529411765),
    (1.0, 0.4980392156862745, 0.054901960784313725),
    (0.17254901960784313, 0.6274509803921569, 0.17254901960784313),
    (0.8392156862745098, 0.15294117647058825, 0.1568627450980392),
    (0.5803921568627451, 0.403921568627451, 0.7411764705882353),
    (0.5490196078431373, 0.33725490196078434, 0.29411764705882354),
    (0.8901960784313725, 0.4666666666666667, 0.7607843137254902),
    (0.4980392156862745, 0.4980392156862745, 0.4980392156862745),
    (0.7372549019607844, 0.7411764705882353, 0.13333333333333333),
    (0.09019607843137255, 0.7450980392156863, 0.8117647058823529),
]
plots = 2
palette = sns.color_palette(seven_tab10_colors[:4])

if plots == 2:
    first_col = [seven_tab10_colors[0]]
    seven_tab10_colors = first_col + seven_tab10_colors[-3:]
    palette = sns.color_palette(seven_tab10_colors)

sns.set(style="darkgrid", font_scale=1)
OUTPUT_DIR = "./output"


def get_scalars(
    tb_log,
    scalar_name,
    window_size=30,
):

    event_acc = EventAccumulator(tb_log)
    event_acc.Reload()

    try:
        _, time_steps, vals = zip(*event_acc.Scalars(scalar_name))
    except KeyError:
        print("Could not find scalar " + scalar_name)
        return None

    # convert time steps to episodes
    episodes = np.arange(1, len(time_steps) + 1)

    # the cumulative values must be divided by their respective episode length
    if "cum" in scalar_name:
        _, ep_lens_time_steps, ep_lens = zip(*event_acc.Scalars("rollout/ep_len_mean"))
        # access the episode lengths at the desired time steps (this will only make a difference if we access per object data)
        desired_ep_lens = []
        for i in range(len(time_steps)):
            time_steps = np.array(time_steps)
            ep_lens_time_steps = np.array(ep_lens_time_steps)
            index = np.where(time_steps[i] == ep_lens_time_steps)[0][0]
            desired_ep_lens.append(ep_lens[index])
        vals = np.array(vals) / np.array(desired_ep_lens)

    # smooth data with moving average filter
    vals = np.convolve(vals, np.ones(window_size), "valid") / window_size
    episodes = np.convolve(episodes, np.ones(window_size), "valid") / window_size

    # create pd dataframe
    data = {"episodes": episodes, scalar_name: vals}
    return pd.DataFrame(data)


def get_experiment_data(tb_log, args):

    event_acc = EventAccumulator(tb_log)
    event_acc.Reload()

    # the variable that you want to keep fixed
    # try:
    #     _, _, x = zip(*event_acc.Scalars(f"hparams/{args.keep_fixed}"))

    #     if x[0] != args.keep_fixed_value:
    #         print("none")
    #         return None

    # except Exception as e:
    #     print(e)
    #     return None

    df = pd.DataFrame()

    for s in args.scalar_names:
        df = df.append(get_scalars(tb_log, s[0]))

    # get some interesting hyperparameters
    des_hparams = ["reward_framework", "force_framework", "w_eps_torque", "w_delta", "gradient_steps", "learning_rate", "batch_size", "tau"]

    for hparam in des_hparams:
        try:
            _, _, x = zip(*event_acc.Scalars("hparams/" + hparam))
        except KeyError as e:
            print(e)
            continue
        if "framework" in hparam:
            force_framework = True if hparam == "force_framework" else False
            df[hparam] = get_framework_name(force_framework, x[0])
        else:
            df[hparam] = f"{x[0]:.5f}"

    # extract experiment id from log string
    id = tb_log.split("id")[1].split("_")[0]
    df["try_id"] = id
    df["tb_log"] = tb_log
    return df


def get_all_data(args, valid_tb_logs):
    print("Getting all data...")
    # get any directories
    all_dirs = [os.path.join(args.log_path, dir) for dir in os.listdir(args.log_path) if dir.endswith("_1")]

    # filter dirs for valid ones
    dirs = []
    for i in range(len(valid_tb_logs)):
        for j in range(len(all_dirs)):
            if all_dirs[j].__contains__(valid_tb_logs[i]):
                dirs.append(all_dirs[j])

    df = pd.DataFrame()

    m = 0
    for tb_log in dirs:
        m += 1
        print(f"({m}) - Getting " + tb_log)
        df = df.append(get_experiment_data(tb_log, args))
        # break

    # very infrequently we get data points with very high forces values, we treat them as outliers and remove them
    df_sum = df[df["episode/cum_sum_contact_forces"] <= 200]
    df = df[np.isnan(df["episode/cum_sum_contact_forces"])]
    df = df.append(df_sum)
    print(df)
    return df


def plot(args, df, title):
    hue_order = ["Full", "Normal", "Binary", "None"] if args.compare == "force_framework" else [r"$\epsilon$ and $\delta$", r"$\delta$", r"$\epsilon$", r"$\beta$"]
    num_frameworks = 4

    # all data should start at 15.5 episodes
    df = df[df["episodes"] >= 15.5]

    num_plots = len(args.scalar_names)
    fig = plt.figure(figsize=(10, 14.7))
    num_tries_per_exp = len(np.unique(df["tb_log"].values)) / num_frameworks
    print("Num tries per exp " + str(num_tries_per_exp))

    for i in range(num_plots):
        ax = fig.add_subplot(8, 3, i + 1)
        sns.lineplot(data=df, x="episodes", y=args.scalar_names[i][0], palette=palette, hue=args.compare, hue_order=hue_order, ax=ax)

        lines, labels = ax.get_legend_handles_labels()
        ax.get_legend().remove()
        ax.set_xlabel("Episodes")
        ax.set_xlim((0, 900))
        ax.xaxis.set_ticks(np.arange(0, 900, 250))
        ax.set_title("Plot (" + str(i + 1) + ") - All Objects")

        # a few metrics do not require setting of y ticks
        not_zero_to_one = ["rollout/ep_rew_mean", "episode/cum_num_contacts", "episode/cum_sum_contact_forces"]
        if args.scalar_names[i][0] not in not_zero_to_one:
            ax.yaxis.set_ticks(np.arange(0, 1.1, 0.5))

        if "box" in args.scalar_names[i][0]:
            ax.set_title("Plot (" + str(i + 1) + ") - Cuboids")
            ax.set_xlabel("Cuboid-Episodes")
            # todo remove
            ax.set_xlim((0, 300))
            ax.xaxis.set_ticks(np.arange(0, 290, 100))
        elif "cylinder" in args.scalar_names[i][0]:
            ax.set_title("Plot (" + str(i + 1) + ") - Cylinders")
            ax.set_xlabel("Cylinder-Episodes")
            # todo remove
            ax.set_xlim((0, 300))
            ax.xaxis.set_ticks(np.arange(0, 290, 100))
        elif "sphere" in args.scalar_names[i][0]:
            ax.set_title("Plot (" + str(i + 1) + ") - Spheres")
            ax.set_xlabel("Sphere-Episodes")
            # todo remove
            ax.set_xlim((0, 300))
            ax.xaxis.set_ticks(np.arange(0, 290, 100))
        ax.set_ylabel(args.scalar_names[i][1])
        if "sustained" in args.scalar_names[i][0]:
            ax.set_ylim((0, 1))

    fig.tight_layout(rect=[0, 0, 1, 0.96], pad=0.05, h_pad=1, w_pad=0.7)
    fig.legend(lines, labels, loc="upper center", ncol=4)
    fig.show()
    fig.savefig(f"{OUTPUT_DIR}/thesis_train_{args.prefix}_{args.compare}.png")
    fig.savefig(f"{OUTPUT_DIR}/thesis_train_{args.prefix}_{args.compare}.pdf")


if __name__ == "__main__":
    parser = ArgumentParser("Plots results of RL algorithm training.")
    parser.add_argument(
        "--log_path",
        type=str,
        default="/Users/koenig/Documents/master/semester_5/thesis/data/cluster_logs_sd6_14sep",
        help="Path to tensorboard log.",
    )
    parser.add_argument(
        "--prefix",
        type=str,
        default="23Aug_SmallTau25k_f",
        help="Prefix comment of your experiment.",
    )
    parser.add_argument(
        "--scalar_names",
        type=str,
        default="step/sustained_holding",
        help="Which metric to plot.",
    )
    parser.add_argument(
        "--compare",
        type=str,
        default="reward_framework",
        help="Which metric to compare.",
    )
    args = parser.parse_args()

    args.scalar_names = [
        ("step/sustained_holding", "Hold Success"),
        ("step/sustained_lifting", "Lift Success"),
        ("rollout/ep_rew_mean", "Episode Reward"),
        ("step/sustained_holding_box", "Hold Success"),
        ("step/sustained_holding_cylinder", "Hold Success"),
        ("step/sustained_holding_sphere", "Hold Success"),
        ("episode/cum_a_epsilon_force", r"$\epsilon_{f}$"),
        ("episode/cum_a_epsilon_torque", r"$\epsilon_{\tau}$"),
        ("episode/cum_num_contacts", "Num. Contacts"),
        ("episode/cum_a_delta_cur", r"$\delta_{cur}$"),
        ("episode/cum_a_delta_task", r"$\delta_{task}$"),
        ("episode/cum_sum_contact_forces", "Sum Forces in N"),
        ("episode/cum_a_epsilon_force_box", r"$\epsilon_{f}$"),
        ("episode/cum_a_epsilon_force_cylinder", r"$\epsilon_{f}$"),
        ("episode/cum_a_epsilon_force_sphere", r"$\epsilon_{f}$"),
        ("episode/cum_a_epsilon_torque_box", r"$\epsilon_{\tau}$"),
        ("episode/cum_a_epsilon_torque_cylinder", r"$\epsilon_{\tau}$"),
        ("episode/cum_a_epsilon_torque_sphere", r"$\epsilon_{\tau}$"),
        ("episode/cum_a_delta_cur_box", r"$\delta_{cur}$"),
        ("episode/cum_a_delta_cur_cylinder", r"$\delta_{cur}$"),
        ("episode/cum_a_delta_cur_sphere", r"$\delta_{cur}$"),
        ("episode/cum_a_delta_task_box", r"$\delta_{task}$"),
        ("episode/cum_a_delta_task_cylinder", r"$\delta_{task}$"),
        ("episode/cum_a_delta_task_sphere", r"$\delta_{task}$"),
    ]

    if plots == 1:
        args.compare = "reward_framework"
        frameworks = [r"$\epsilon$ and $\delta$", r"$\delta$", r"$\epsilon$", r"$\beta$"]
    elif plots == 2:
        args.compare = "force_framework"
        frameworks = ["Full", "Normal", "Binary", "None"]

    # we want to filter the frameworks that were successfully completeted
    done_frameworks_path = os.path.join(".", f"done_{args.compare}.pkl")
    with open(done_frameworks_path, "rb") as file:
        valid_tb_logs = pickle.load(file)

    title_path = os.path.join(".", f"title_{args.compare}.pkl")
    with open(title_path, "rb") as file:
        title = pickle.load(file)

    # get all data that we have a test result for
    df = get_all_data(args, valid_tb_logs)

    plot(args, df, title)

    print("done!")

    import pdb

    pdb.set_trace()
