#!/usr/bin/env python3

import argparse
import gym
import os
import sys
import time
import matplotlib.pyplot as plt
from a3c.discrete_A3C import train, evaluate, evaluate_checkpoints
from wordle_env.wordle import WordleEnvBase


def training_mode(args, env, model_checkpoint_dir):
    max_ep = args.games
    start_time = time.time()
    if args.model_name:
        pretrained_model_path = os.path.join(
            model_checkpoint_dir, args.model_name)
        global_ep, win_ep, gnet, res = train(
            env, max_ep, model_checkpoint_dir, pretrained_model_path)
    else:
        global_ep, win_ep, gnet, res = train(env, max_ep, model_checkpoint_dir)
    print("--- %.0f seconds ---" % (time.time() - start_time))
    print_results(global_ep, win_ep, res)
    evaluate(gnet, env)


def evaluation_mode(args, env, model_checkpoint_dir):
    print("Evaluation mode")
    results = evaluate_checkpoints(model_checkpoint_dir, env)
    print(results)


def print_results(global_ep, win_ep, res):
    print("Jugadas:", global_ep.value)
    print("Ganadas:", win_ep.value)
    plt.plot(res)
    plt.ylabel('Moving average ep reward')
    plt.xlabel('Step')
    plt.show()


if __name__ == "__main__":
    parser = argparse.ArgumentParser()
    parser.add_argument(
        "enviroment", help="Enviroment (type of wordle game) used for training, example: WordleEnvFull-v0")
    parser.add_argument(
        "--models_dir", help="Directory where models are saved (default=checkpoints)", default='checkpoints')
    subparsers = parser.add_subparsers(help='sub-command help')

    parser_train = subparsers.add_parser(
        'train', help='Train a model from scratch or train from pretrained model')
    parser_train.add_argument(
        "--games", "-g", help="Number of games to train", type=int, required=True)
    parser_train.add_argument(
        "--model_name", "-n", help="If want to train from a pretrained model, the name of the pretrained model file")
    parser_train.add_argument(
        "--gamma", help="Gamma hyperparameter value", type=float, default=0.)
    parser_train.set_defaults(func=training_mode)

    parser_eval = subparsers.add_parser(
        'eval', help='Evaluate saved models for the enviroment')
    parser_eval.set_defaults(func=evaluation_mode)

    args = parser.parse_args()
    env_id = args.enviroment
    env = gym.make(env_id)
    model_checkpoint_dir = os.path.join(args.models_dir, env.unwrapped.spec.id)
    args.func(args, env, model_checkpoint_dir)