# Copyright (c) Facebook, Inc. and its affiliates.
import json
import os
import re

import matplotlib.pyplot as plt
import numpy as np


def read_log_file(file_name, key_name, value_name, smooth=3):
    keys, values = [], []
    try:
        with open(file_name, "r") as f:
            for line in f:
                try:
                    e = json.loads(line.strip())
                    key, value = e[key_name], e[value_name]
                    keys.append(int(key))
                    values.append(float(value))
                except:
                    pass
    except:
        print("bad file: %s" % file_name)
        return None, None
    keys, values = np.array(keys), np.array(values)
    if smooth > 1 and values.shape[0] > 0:
        K = np.ones(smooth)
        ones = np.ones(values.shape[0])
        values = np.convolve(values, K, "same") / np.convolve(ones, K, "same")

    return keys, values


def parse_log_files(
    file_name_template, key_name, value_name, num_seeds, best_k=None, max_key=False
):
    all_values = []
    all_keys = []
    actual_keys = None
    for seed in range(1, num_seeds + 1):
        file_name = file_name_template % seed
        keys, values = read_log_file(file_name, key_name, value_name)
        if keys is None or keys.shape[0] == 0:
            continue
        all_keys.append(keys)
        all_values.append(values)

    if len(all_values) == 0:
        return None, None, None

    all_keys = sorted(all_keys, key=lambda x: x.shape[0])
    threshold = all_keys[-1].shape[0] if max_key else all_keys[0].shape[0]

    means, half_stds = [], []
    for i in range(threshold):
        vals = []

        for v in all_values:
            if i < v.shape[0]:
                vals.append(v[i])
        if best_k is not None:
            vals = sorted(vals)[-best_k:]
        means.append(np.mean(vals))
        half_stds.append(0.5 * np.std(vals))
    means = np.array(means)
    half_stds = np.array(half_stds)

    keys = all_keys[-1][:threshold]
    assert means.shape[0] == keys.shape[0]

    return keys, means, half_stds


def print_result(
    root,
    title,
    label=None,
    num_seeds=1,
    train=False,
    key_name="step",
    value_name="train_episode_reward",
    max_time=None,
    best_k=None,
    timescale=1,
    max_key=False,
    gen_gap=False,
):
    file_name = "train.log" if train else "eval.log"
    file_name_template = os.path.join(root, "seed_%d", file_name)
    keys, means, half_stds = parse_log_files(
        file_name_template,
        key_name,
        value_name,
        num_seeds,
        best_k=best_k,
        max_key=max_key,
    )

    label = label or root.split("/")[-1]
    if keys is None:
        return

    if gen_gap:
        _, eval_means, eval_half_stds = parse_log_files(
            file_name_template,
            key_name,
            "eval_episode_reward",
            num_seeds,
            best_k=best_k,
            max_key=max_key,
        )
        means = means - eval_means
        half_stds = half_stds - eval_half_stds

    if max_time is not None:
        idxs = np.where(keys <= max_time)
        keys = keys[idxs]
        means = means[idxs]
        half_stds = half_stds[idxs]

    keys *= timescale

    plt.plot(keys, means, label=label)
    plt.locator_params(nbins=10, axis="x")
    plt.locator_params(nbins=10, axis="y")
    plt.rcParams["figure.figsize"] = (10, 7)
    plt.rcParams["figure.dpi"] = 100
    plt.rcParams["font.size"] = 10
    plt.subplots_adjust(left=0.165, right=0.99, bottom=0.16, top=0.95)
    # plt.ylim(0, 1050)
    plt.tight_layout()

    plt.grid(alpha=0.8)
    plt.title(title)
    plt.fill_between(keys, means - half_stds, means + half_stds, alpha=0.2)
    plt.legend(loc="lower right", prop={"size": 6}).get_frame().set_edgecolor("0.1")
    plt.xlabel(key_name)
    if gen_gap:
        plt.ylabel("Generalization Gap")
    else:
        plt.ylabel(value_name)


def plot_seeds(
    task,
    exp_query,
    root,
    train=True,
    key_name="step",
    value_name="episode_reward",
    num_seeds=10,
):
    root = os.path.join(root, task)
    experiment = None
    for exp in os.listdir(root):
        if re.match(exp_query, exp):
            experiment = os.path.join(root, exp)
            break
    if experiment is None:
        return
    file_name = "train.log" if train else "eval.log"
    file_name_template = os.path.join(experiment, "seed_%d", file_name)

    plt.locator_params(nbins=10, axis="x")
    plt.locator_params(nbins=10, axis="y")
    plt.rcParams["figure.figsize"] = (10, 7)
    plt.rcParams["figure.dpi"] = 100
    plt.rcParams["font.size"] = 10
    plt.subplots_adjust(left=0.165, right=0.99, bottom=0.16, top=0.95)
    plt.grid(alpha=0.8)
    plt.tight_layout()
    plt.title(task)

    plt.xlabel(key_name)
    plt.ylabel(value_name)

    for seed in range(1, num_seeds + 1):
        file_name = file_name_template % seed
        keys, values = read_log_file(file_name, key_name, value_name)
        if keys is None or keys.shape[0] == 0:
            continue

        plt.plot(keys, values, label="seed_%d" % seed, linewidth=0.5)

    plt.legend(loc="lower right", prop={"size": 6}).get_frame().set_edgecolor("0.1")


def print_baseline(task, baseline, data, color):
    try:
        value = data[task][baseline]
    except:
        return

    plt.axhline(y=value, label=baseline, linestyle="--", color=color)
    plt.legend(loc="lower right", prop={"size": 6}).get_frame().set_edgecolor("0.1")


def print_planet_baseline(
    task, data, max_time=None, label="planet", color="black", offset=0
):
    try:
        keys, means, half_stds = data[task]
    except:
        return

    if max_time is not None:
        idx = np.searchsorted(keys, max_time)
        keys = keys[:idx]
        means = means[:idx]
        half_stds = half_stds[:idx]

    plt.plot(keys + offset, means, label=label, color=color)
    plt.fill_between(
        keys + offset, means - half_stds, means + half_stds, alpha=0.2, color=color
    )
    plt.legend(loc="lower right", prop={"size": 6}).get_frame().set_edgecolor("0.1")


def plot_experiment(
    task,
    exp_query,
    root="runs",
    exp_ids=None,
    train=False,
    key_name="step",
    value_name="train_episode_reward",
    baselines_data=None,
    num_seeds=10,
    planet_data=None,
    slac_data=None,
    max_time=None,
    best_k=None,
    timescale=1,
    max_key=False,
    gen_gap=False,
):
    root = os.path.join(root, task)

    experiments = set()
    for exp in os.listdir(root):
        if re.match(exp_query, exp):
            exp = os.path.join(root, exp)
            experiments.add(exp)

    exp_ids = list(range(len(experiments))) if exp_ids is None else exp_ids
    for exp_id, exp in enumerate(sorted(experiments)):
        if exp_id in exp_ids:
            print_result(
                exp,
                task,
                num_seeds=num_seeds,
                train=train,
                key_name=key_name,
                value_name=value_name,
                max_time=max_time,
                best_k=best_k,
                timescale=timescale,
                max_key=max_key,
                gen_gap=gen_gap,
            )

    if baselines_data is not None:
        print_baseline(task, "d4pg_pixels", baselines_data, color="gray")
        print_baseline(task, "d4pg", baselines_data, color="black")

    if planet_data is not None:
        print_planet_baseline(
            task, planet_data, max_time=max_time, label="planet", color="peru", offset=5
        )

    if slac_data is not None:
        action_repeat = {
            "ball_in_cup_catch": 4,
            "cartpole_swingup": 8,
            "cheetah_run": 4,
            "finger_spin": 2,
            "walker_walk": 2,
            "reacher_easy": 4,
        }
        offset = 10 * action_repeat[task]
        print_planet_baseline(
            task,
            slac_data,
            max_time=max_time,
            label="slac",
            color="black",
            offset=offset,
        )
