import os
import sys

from utils.logging import get_logger, Logger

import random
import matplotlib
import argparse
import yaml
from types import SimpleNamespace

from main import recursive_dict_update
from run import run_sequential
import logging

if __name__ == "__main__":
    parser = argparse.ArgumentParser()
    #parser.add_argument("--model_path", type=str)
    parser.add_argument("--config", type=str, default="mappo_egg")
    parser.add_argument("--env-config", type=str, default="egg_large_mo")
    parser.add_argument("--model", type=str, default="")
    parser.add_argument("--eff", type=float, default=0.4)
    test_args = parser.parse_args()

    # load env config as a name space
    config = yaml.safe_load(open("config/default.yaml"))
    recursive_dict_update(config, yaml.safe_load(open(f"config/envs/{test_args.env_config}.yaml")))
    recursive_dict_update(config, yaml.safe_load(open(f"config/algs/{test_args.config}.yaml")))

    args = SimpleNamespace(**config)

    args.env_args["efficiency"] = test_args.eff
    #args.env_args["survival_threshold"] = 40
    args.evaluate = True
    args.render = False
    args.use_cuda = False
    # Get current directory
    args.checkpoint_path = os.path.join(os.path.dirname(os.getcwd()), test_args.model)
    #last_ckpt = max([int(d) for d in os.listdir(args.checkpoint_path) if os.path.isdir(os.path.join(args.checkpoint_path, d))])
    #args.checkpoint_path = os.path.join(args.checkpoint_path, str(last_ckpt))

    args.env_args["seed"] = random.randint(0, 1000000)
    args.batch_size_run = 1
    args.device = "cpu"
    args.test_nepisode = 2000
    args.runner = "episode"
    args.action_selector = "soft_policies_filter_low"
    args.action_selection_filter_low = 0.1

    # Redirect stdout to file
    sys.stdout = open(args.checkpoint_path + "/test_output.txt", "w")

    runner, mac, learner = run_sequential(args, logger)

    env = runner.env._env.unwrapped
    env.print_results()
    env.plot_results("median", save_path=args.checkpoint_path+"/results.png")


