import os
import torch

from .net import GreedyNet
from .play import play
from .utils import v_wrap


def evaluate_checkpoints(dir, env):
    n_s = env.observation_space.shape[0]
    n_a = env.action_space.n
    words_list = env.words
    word_width = len(env.words[0])
    net = GreedyNet(n_s, n_a, words_list, word_width)
    results = {}
    for checkpoint in os.listdir(dir):
        checkpoint_path = os.path.join(dir, checkpoint)
        if os.path.isfile(checkpoint_path):
            net.load_state_dict(torch.load(checkpoint_path))
            wins, guesses = evaluate(net, env)
            results[checkpoint] = wins, guesses
    return dict(sorted(results.items(), key=lambda x: (x[1][0], -x[1][1]), reverse=True))


def evaluate(net, env):
    n_wins = 0
    n_guesses = 0
    n_win_guesses = 0
    env = env.unwrapped
    N = env.allowable_words
    for goal_word in env.words[:N]:
        win, outcomes = play(net, env)
        if win:
            n_wins += 1
            n_win_guesses += len(outcomes)
        # else:
        #     print("Lost!", goal_word, outcomes)
        n_guesses += len(outcomes)
    print(f"Evaluation complete, won {n_wins/N*100}% and took {n_win_guesses/n_wins} guesses per win, "
          f"{n_guesses / N} including losses.")
    return n_wins/N*100, n_win_guesses/n_wins