#!/usr/bin/env python3

import argparse
import os
import time
import matplotlib.pyplot as plt
from a3c.train import train
from a3c.eval import evaluate, evaluate_checkpoints
from a3c.play import suggest
from wordle_env.wordle import get_env


def training_mode(args, env, model_checkpoint_dir):
    max_ep = args.games
    start_time = time.time()
    pretrained_model_path = os.path.join(model_checkpoint_dir, args.model_name) if args.model_name else args.model_name
    global_ep, win_ep, gnet, res = train(env, max_ep, model_checkpoint_dir, args.gamma, args.seed, pretrained_model_path, args.save, args.min_reward, args.every_n_save)
    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 play_mode(args, env, model_checkpoint_dir):
    print("Play mode")
    words = [ word.strip() for word in args.words.split(',') ]
    states = [ state.strip() for state in args.states.split(',') ]
    pretrained_model_path = os.path.join(model_checkpoint_dir, args.model_name)
    word = suggest(env, words, states, pretrained_model_path)
    print(word)


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", "-m", help="If want to train from a pretrained model, the name of the pretrained model file")
    parser_train.add_argument(
        "--gamma", help="Gamma hyperparameter (discount factor) value", type=float, default=0.)
    parser_train.add_argument(
        "--seed", help="Seed used for random numbers generation", type=int, default=100)
    parser_train.add_argument(
        "--save", '-s', help="Save instances of the model while training", action='store_true')
    parser_train.add_argument(
        "--min_reward", help="The minimun global reward value achieved for saving the model", type=float, default=9.9)
    parser_train.add_argument(
        "--every_n_save", help="Check every n training steps to save the model", type=int, default=100)
    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)

    parser_play = subparsers.add_parser(
        'play', help='Give the model a word and the state result and the model will try to predict the goal word')
    parser_play.add_argument(
        "--words", "-w", help="List of words played in the wordle game", required=True)
    parser_play.add_argument(
        "--states", "-st", help="List of states returned by playing each of the words", required=True)
    parser_play.add_argument(
        "--model_name", "-m", help="Name of the pretrained model file thich will play the game", required=True)
    parser_play.set_defaults(func=play_mode)

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