import argparse
import os
import matplotlib.pyplot as plt
import json
import time
import numpy as np
import torch
from pathlib import Path

from utils.babyai_utils.baby_agent import load_agent
from utils.storage import get_status
from utils.env import make_env
from utils.other import seed
from utils.storage import get_model_dir
from models import *
from utils.env import env_args_str_to_dict
import gym
from termcolor import cprint

os.makedirs("./evaluation", exist_ok=True)

start = time.time()

# Parse arguments

parser = argparse.ArgumentParser()
parser.add_argument("--test-set-seed", type=int, default=0,
                    help="random seed (default: 0)")
parser.add_argument("--random-agent", action="store_true", default=False,
                    help="random actions")
parser.add_argument("--quiet", "-q", action="store_true", default=False,
                    help="quiet")
parser.add_argument("--eval-env", type=str, default=None,
                    help="env to evaluate on")
parser.add_argument("--model-to-evaluate", type=str, default=None,
                    help="model to evaluate")
parser.add_argument("--model-label", type=str, default=None,
                    help="model to evaluate")
parser.add_argument("--max-steps", type=int, default=None,
                    help="max num of steps")
parser.add_argument("--argmax", action="store_true", default=False,
                    help="select the action with highest probability (default: False)")
parser.add_argument("--episodes", type=int, default=1000,
                    help="number of episodes to test")
parser.add_argument("--test-p", type=float, default=0.05,
                    help="p value")
parser.add_argument("--n-seeds", type=int, default=8,
                    help="number of episodes to test")
parser.add_argument("--subsample-step", type=int, default=1,
                    help="subsample step")
parser.add_argument("--start-step", type=int, default=1,
                    help="at which step to start the curves")
parser.add_argument("--env_args", nargs='*', default=None)

args = parser.parse_args()

# Set seed for all randomness sources

seed(args.test_set_seed)

assert args.test_set_seed == 1 # turn on for testing
# assert not args.argmax

# assert args.num_frames == 28000000
# assert args.episodes == 1000

test_p = args.test_p
n_seeds = args.n_seeds
assert n_seeds in [16, 8, 4]
cprint("n seeds: {}".format(n_seeds), "red")
subsample_step = args.subsample_step
start_step = args.start_step

# Set device
def qprint(*a, **kwargs):
    if not args.quiet:
        print(*a, **kwargs)

device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
qprint(f"Device: {device}\n")

# what to load
if args.model_to_evaluate is None:
    models_to_evaluate = [
        "19-05_500K_HELP_env_MiniGrid-Exiter-8x8-v0_multi-modal-babyai11-agent_arch_original_endpool_res_custom-ppo-2"
    ]
    label_parser_dict = {
        "19-05_500K_HELP_env_MiniGrid-Exiter-8x8-v0_multi-modal-babyai11-agent_arch_original_endpool_res_custom-ppo-2": "Exiter_EB",
    }
else:
    model_name = args.model_to_evaluate.replace("./storage/", "").replace("storage/", "")
    models_to_evaluate = [
        model_name
    ]
    if args.model_label:
        label_parser_dict = {
            model_name: args.model_label,
        }
    else:
        label_parser_dict = {
            model_name: model_name,
        }
    qprint("evaluating models: ", models_to_evaluate)


# how do to stat tests
compare = {
    # "MH-BabyAI-ExpBonus": "Abl-MH-BabyAI-ExpBonus",
}

COLORS = ["red", "blue", "green", "black", "purpule", "brown", "orange", "gray"]
label_color_dict = {l: c for l, c in zip(label_parser_dict.values(), COLORS)}


test_set_check_path = Path("test_set_check_{}_nep_{}.json".format(args.test_set_seed, args.episodes))

def calc_perf_for_seed(i, model_name, seed, argmax, episodes, random_agent=False, num_frames=None):
    qprint("seed {}".format(i))
    model = Path(model_name) / str(i)
    model_dir = get_model_dir(model)

    if test_set_check_path.exists():
        with open(test_set_check_path, "r") as f:
            check_loaded = json.load(f)
        qprint("check loaded")
    else:
        qprint("check not loaded")
        check_loaded = None

    # Load environment
    with open(model_dir+"/config.json") as f:
        conf = json.load(f)

    if args.eval_env is None:
        qprint("evaluating on the original env")
        env_name = conf["env"]
    else:
        qprint("evaluating on a different env")
        env_name = args.eval_env

    env = gym.make(env_name, **env_args_str_to_dict(args.env_args))
    qprint("Environment loaded\n")

    # load agent
    agent = load_agent(env, model_dir, argmax)
    status = get_status(model_dir)
    qprint("Agent loaded at {} steps.".format(status.get("num_frames", -1)))

    check = {}

    seed_rewards = []
    seed_sr = []
    for episode in range(episodes):
        qprint("[{}/{}]: ".format(episode, episodes), end="", flush=True)

        obs = env.reset()

        # check envs are the same during seeds
        if episode in check:
            assert check[episode] == int(obs['image'].sum())
        else:
            check[episode] = int(obs['image'].sum())

        if check_loaded is not None:
            assert check[episode] == int(obs['image'].sum())
        i = 0
        tot_reward = 0
        while True:
            i+=1
            if random_agent:
                action = agent.get_random_action(obs)
            else:
                action = agent.get_action(obs)

            obs, reward, done, info = env.step(action)
            if reward:
                qprint("*", end="", flush=True)
            else:
                qprint(".", end="", flush=True)

            agent.analyze_feedback(reward, done)

            tot_reward += reward

            if done:
                seed_rewards.append(tot_reward)
                seed_sr.append(info["success"])
                break

            if args.max_steps is not None:
                if i > args.max_steps:
                    seed_rewards.append(tot_reward)
                    seed_sr.append(info["success"])
                    break

        qprint()

    seed_rewards = np.array(seed_rewards)
    seed_success_rates = np.array(seed_sr)

    if not test_set_check_path.exists():
        with open(test_set_check_path, "w") as f:
            json.dump(check, f)
            qprint("check saved")

    qprint("seed success rate:", seed_success_rates.mean())
    qprint("seed reward:", seed_rewards.mean())

    return seed_rewards.mean(), seed_success_rates.mean()


def get_available_steps(model):
    model_dir = Path(get_model_dir(model))
    per_seed_available_steps = {}
    for seed_dir in model_dir.glob("*"):
        per_seed_available_steps[seed_dir] = sorted([
           int(str(p.with_suffix("")).split("status_")[-1])
           for p in seed_dir.glob("status_*")
        ])

    num_steps = min([len(steps) for steps in per_seed_available_steps.values()])

    steps = list(per_seed_available_steps.values())[0][:num_steps]

    for available_steps in per_seed_available_steps.values():
        s_steps = available_steps[:num_steps]
        assert steps == s_steps

    return steps

def plot_with_shade(subplot_nb, ax, x, y, err, color, shade_color, label,
                    legend=False, leg_size=30, leg_loc='best', title=None,
                    ylim=[0, 100], xlim=[0, 40], leg_args={}, leg_linewidth=8.0, linewidth=7.0, ticksize=30,
                    zorder=None, xlabel='perf', ylabel='env steps', smooth_factor=1000):
    # plt.rcParams.update({'font.size': 15})
    ax.locator_params(axis='x', nbins=6)
    ax.locator_params(axis='y', nbins=5)
    ax.tick_params(axis='both', which='major', labelsize=ticksize)

    # smoothing
    def smooth(x_, n=50):
        return np.array([x_[max(i - n, 0):i + 1].mean() for i in range(len(x_))])

    if smooth_factor > 0:
        y = smooth(y, n=smooth_factor)
        err = smooth(err, n=smooth_factor)

    ax.plot(x, y, color=color, label=label, linewidth=linewidth, zorder=zorder)
    ax.fill_between(x, y - err, y + err, color=shade_color, alpha=0.2)
    if legend:
        leg = ax.legend(loc=leg_loc, fontsize=leg_size, **leg_args)  # 34
        for legobj in leg.legendHandles:
            legobj.set_linewidth(leg_linewidth)
    ax.set_xlabel(xlabel, fontsize=30)
    if subplot_nb == 0:
        ax.set_ylabel(ylabel, fontsize=30)
    ax.set_xlim(xmin=xlim[0], xmax=xlim[1])
    ax.set_ylim(bottom=ylim[0], top=ylim[1])
    if title:
        ax.set_title(title, fontsize=22)


def label_parser(label, label_parser_dict):
    if sum([1 for k, v in label_parser_dict.items() if k in label]) != 1:
        qprint("ERROR")
        qprint(label)
        exit()

    for k, v in label_parser_dict.items():
        if k in label: return v

    return label


f, ax = plt.subplots(1, 1, figsize=(10.0, 6.0))
ax = [ax]

performances = {}
per_seed_performances = {}
stds = {}


label_parser_dict_reverse = {v: k for k, v in label_parser_dict.items()}
assert len(label_parser_dict_reverse) == len(label_parser_dict)

label_to_model = {}
# evaluate and draw curves
for model in models_to_evaluate:
    label = label_parser(model, label_parser_dict)
    label_to_model[label] = model

    color = label_color_dict[label]
    performances[label] = []
    per_seed_performances[label] = []
    stds[label] = []

    final_perf = True

    if final_perf:

        results = []
        for s in range(n_seeds):
            results.append(calc_perf_for_seed(
                s,
                model_name=model,
                num_frames=None,
                seed=args.test_set_seed,
                argmax=args.argmax,
                episodes=args.episodes,
            ))
        rewards, success_rates = zip(*results)
        # dump per seed performance
        np.save("./evaluation/{}".format(label), success_rates)
        rewards = np.array(rewards)
        success_rates = np.array(success_rates)
        success_rate_mean = success_rates.mean()
        succes_rate_std = success_rates.std()

        label = label_parser(str(model), label_parser_dict)
        cprint("{}: {} +- std {}".format(label, success_rate_mean, succes_rate_std), "red")

    else:
        steps = get_available_steps(model)
        steps = steps[::subsample_step]
        steps = [s for s in steps if s > start_step]
        qprint("steps:", steps)

        for step in steps:
            results = []
            for s in range(n_seeds):
                results.append(calc_perf_for_seed(
                    s,
                    model_name=model,
                    num_frames=step,
                    seed=args.test_set_seed,
                    argmax=args.argmax,
                    episodes=args.episodes,
                ))

            rewards, success_rates = zip(*results)
            rewards = np.array(rewards)
            success_rates = np.array(success_rates)
            per_seed_performances[label].append(success_rates)
            performances[label].append(success_rates.mean())
            stds[label].append(success_rates.std())

        means = np.array(performances[label])
        err = np.array(stds[label])
        label = label_parser(str(model), label_parser_dict)
        max_steps = np.max(steps)
        min_steps = np.min(steps)
        min_y = 0.0
        max_y = 1.0
        ylabel = "performance"
        smooth_factor = 0

        plot_with_shade(0, ax[0], steps, means, err, color, color, label,
                        legend=True, xlim=[min_steps, max_steps], ylim=[min_y, max_y],
                        leg_size=20, xlabel="Env steps (millions)", ylabel=ylabel, linewidth=5.0, smooth_factor=smooth_factor)

assert len(label_to_model) == len(models_to_evaluate)


def get_compatible_steps(model1, model2, subsample_step):
    steps_1 = get_available_steps(model1)[::subsample_step]
    steps_2 = get_available_steps(model2)[::subsample_step]

    min_steps = min(len(steps_1), len(steps_2))
    steps_1 = steps_1[:min_steps]
    steps_2 = steps_2[:min_steps]
    assert steps_1 == steps_2

    return steps_1


# # stat tests
# for k, v in compare.items():
#     dist_1_steps = per_seed_performances[k]
#     dist_2_steps = per_seed_performances[v]
#
#     model_k = label_to_model[k]
#     model_v = label_to_model[v]
#     steps = get_compatible_steps(model_k, model_v, subsample_step)
#     steps = [s for s in steps if s > start_step]
#
#     for step, dist_1, dist_2 in zip(steps, dist_1_steps, dist_2_steps):
#         assert len(dist_1) == n_seeds
#         assert len(dist_2) == n_seeds
#
#         p = stats.ttest_ind(
#             dist_1,
#             dist_2,
#             equal_var=False
#         ).pvalue
#
#         if np.isnan(p):
#             from IPython import embed; embed()
#
#         if p < test_p:
#             plt.scatter(step, 0.8, color=label_color_dict[k], s=50, marker="x")
#
#         print("{} (m:{}) <---> {} (m:{}) = p: {}  result: {}".format(
#             k, np.mean(dist_1), v, np.mean(dist_2), p,
#             "Distributions different(p={})".format(test_p) if p < test_p else "Distributions same(p={})".format(test_p)
#         ))
#         print()
#
# f.savefig('graphics/test.png')
# f.savefig('graphics/test.svg')