from craftworld_env import Craftworld, Cookbook
from simple_rl.run_experiments import run_agents_on_mdp
from simple_rl.agents import QLearningAgent

import context
import rlang
from rlang.agents import RLangQLearningAgent

sys.path.append( path.dirname( path.dirname( path.abspath(__file__) ) ) )
from utils import translate_named, read_named_examples, unroll, parse_vocab, read_input, save_translation_results


def baseline_test():
    """
    This is just for ensuring that Craftworld actually works.
    We should be able to chuck a SimpleRL agent directly into it.
    It works

    """
    mdp = Craftworld(goal='gold', path_to_recipes='craftworld_env/recipes.yaml')
    agent = QLearningAgent(mdp.get_actions())
    run_agents_on_mdp([agent], mdp, steps=400, episodes=100)


def dev():
    cookbook = Cookbook('craftworld_env/recipes.yaml')
    knowledge = rlang.parse_file('craftworld.rlang')
    # print(knowledge)
    baseline_test()

def translate_input(user_input, vocab):
    # QUESTION:  how do we determine if we should pass effect vs. policy examples?
    examples_df = read_named_examples(
        path.join(script_dir,'data/effect_english.txt'), 
        path.join(script_dir, 'data/effect_names.txt'),
        path.join(script_dir, 'data/effect_rlang.txt')
    )

    return translate_named(user_input, examples_df, vocab)[1]

def append_to_rlang(phrase):
    shutil.copy('craftworld.rlang', 'enhanced_craftworld.rlang')
    with open(path.join(script_dir, 'enhanced_craftworld.rlang'), 'a') as out_f:
        out_f.write(f'\n{unroll(phrase)}\n')


def main(argv):
    if len(argv) != 2:
        print('Invalid number of arguments')
        print(f'Expected input: `python run_craftworld.py <PATH_TO_ENGLISH_INPUT>`')
        return
    elif not argv[1].lower().endswith('.txt'):
        print(f'Invalid argument "{argv[1]}"". Input file must be .txt file')
        return

    user_input = read_input(path.join(script_dir, argv[1]))
    vocab = parse_vocab('gridworld.rlang')
    print('vocab?', vocab)

    translation_results = []
    for input_line in user_input:
        t = translate_input(input_line, vocab)
        print('translation? ', t)
        append_to_rlang(t)
        translation_results.append([input_line, t])
    
    save_translation_results(translation_results)

    # TODO: save experiment per input?
    rlang_experiment()

if __name__ == '__main__':
    main(sys.argv)
