from flaky import flaky
from slm_lab.experiment.control import Trial
from slm_lab.experiment.monitor import InfoSpace
from slm_lab.lib import util
from slm_lab.spec import spec_util
import os
import pandas as pd
import pytest
import sys


# helper method to run all tests in test_spec
def run_trial_test(spec_file, spec_name=False):
    spec = spec_util.get(spec_file, spec_name)
    spec = spec_util.override_test_spec(spec)
    info_space = InfoSpace()
    info_space.tick('trial')
    trial = Trial(spec, info_space)
    trial_data = trial.run()
    assert isinstance(trial_data, pd.DataFrame)


@pytest.mark.parametrize('spec_file,spec_name', [
    ('experimental/reinforce.json', 'reinforce_mlp_cartpole'),
    ('experimental/reinforce.json', 'reinforce_rnn_cartpole'),
    # ('experimental/reinforce.json', 'reinforce_conv_breakout'),
])
def test_reinforce(spec_file, spec_name):
    run_trial_test(spec_file, spec_name)


@pytest.mark.parametrize('spec_file,spec_name', [
    ('experimental/reinforce.json', 'reinforce_mlp_pendulum'),
    ('experimental/reinforce.json', 'reinforce_rnn_pendulum'),
])
def test_reinforce_cont(spec_file, spec_name):
    run_trial_test(spec_file, spec_name)


@pytest.mark.parametrize('spec_file,spec_name', [
    ('experimental/a2c.json', 'a2c_mlp_shared_cartpole'),
    ('experimental/a2c.json', 'a2c_mlp_separate_cartpole'),
    ('experimental/a2c.json', 'a2c_rnn_shared_cartpole'),
    ('experimental/a2c.json', 'a2c_rnn_separate_cartpole'),
    # ('experimental/a2c.json', 'a2c_conv_shared_breakout'),
    # ('experimental/a2c.json', 'a2c_conv_separate_breakout'),
    ('experimental/a2c.json', 'a2c_mlp_concat_cartpole'),
])
def test_a2c(spec_file, spec_name):
    run_trial_test(spec_file, spec_name)


@pytest.mark.parametrize('spec_file,spec_name', [
    ('experimental/a2c.json', 'a2c_mlp_shared_pendulum'),
    ('experimental/a2c.json', 'a2c_mlp_separate_pendulum'),
    ('experimental/a2c.json', 'a2c_rnn_shared_pendulum'),
    ('experimental/a2c.json', 'a2c_rnn_separate_pendulum'),
])
def test_a2c_cont(spec_file, spec_name):
    run_trial_test(spec_file, spec_name)


@pytest.mark.parametrize('spec_file,spec_name', [
    ('experimental/ppo.json', 'ppo_mlp_shared_cartpole'),
    ('experimental/ppo.json', 'ppo_mlp_separate_cartpole'),
    ('experimental/ppo.json', 'ppo_rnn_shared_cartpole'),
    ('experimental/ppo.json', 'ppo_rnn_separate_cartpole'),
    # ('experimental/ppo.json', 'ppo_conv_shared_breakout'),
    # ('experimental/ppo.json', 'ppo_conv_separate_breakout'),
])
def test_ppo(spec_file, spec_name):
    run_trial_test(spec_file, spec_name)


@pytest.mark.parametrize('spec_file,spec_name', [
    ('experimental/ppo.json', 'ppo_mlp_shared_pendulum'),
    ('experimental/ppo.json', 'ppo_mlp_separate_pendulum'),
    ('experimental/ppo.json', 'ppo_rnn_shared_pendulum'),
    ('experimental/ppo.json', 'ppo_rnn_separate_pendulum'),
])
def test_ppo_cont(spec_file, spec_name):
    run_trial_test(spec_file, spec_name)


@flaky
@pytest.mark.parametrize('spec_file,spec_name', [
    ('experimental/ppo_sil.json', 'ppo_sil_mlp_shared_cartpole'),
    ('experimental/ppo_sil.json', 'ppo_sil_mlp_separate_cartpole'),
    ('experimental/ppo_sil.json', 'ppo_sil_rnn_shared_cartpole'),
    ('experimental/ppo_sil.json', 'ppo_sil_rnn_separate_cartpole'),
])
def test_ppo_sil(spec_file, spec_name):
    run_trial_test(spec_file, spec_name)


@flaky
@pytest.mark.parametrize('spec_file,spec_name', [
    ('experimental/ppo_sil.json', 'ppo_sil_mlp_shared_pendulum'),
    ('experimental/ppo_sil.json', 'ppo_sil_mlp_separate_pendulum'),
    ('experimental/ppo_sil.json', 'ppo_sil_rnn_shared_pendulum'),
    ('experimental/ppo_sil.json', 'ppo_sil_rnn_separate_pendulum'),
])
def test_ppo_sil_cont(spec_file, spec_name):
    run_trial_test(spec_file, spec_name)


@flaky
@pytest.mark.parametrize('spec_file,spec_name', [
    ('experimental/sil.json', 'sil_mlp_shared_cartpole'),
    ('experimental/sil.json', 'sil_mlp_separate_cartpole'),
    ('experimental/sil.json', 'sil_rnn_shared_cartpole'),
    ('experimental/sil.json', 'sil_rnn_separate_cartpole'),
    # ('experimental/sil.json', 'sil_conv_shared_breakout'),
    # ('experimental/sil.json', 'sil_conv_separate_breakout'),
])
def test_sil(spec_file, spec_name):
    run_trial_test(spec_file, spec_name)


@flaky
@pytest.mark.parametrize('spec_file,spec_name', [
    ('experimental/sil.json', 'sil_mlp_shared_pendulum'),
    ('experimental/sil.json', 'sil_mlp_separate_pendulum'),
    ('experimental/sil.json', 'sil_rnn_shared_pendulum'),
    ('experimental/sil.json', 'sil_rnn_separate_pendulum'),
])
def test_sil_cont(spec_file, spec_name):
    run_trial_test(spec_file, spec_name)


@pytest.mark.parametrize('spec_file,spec_name', [
    ('experimental/sarsa.json', 'sarsa_mlp_boltzmann_cartpole'),
    ('experimental/sarsa.json', 'sarsa_mlp_epsilon_greedy_cartpole'),
    ('experimental/sarsa.json', 'sarsa_rnn_boltzmann_cartpole'),
    ('experimental/sarsa.json', 'sarsa_rnn_epsilon_greedy_cartpole'),
    # ('experimental/sarsa.json', 'sarsa_conv_boltzmann_breakout'),
    # ('experimental/sarsa.json', 'sarsa_conv_epsilon_greedy_breakout'),
])
def test_sarsa(spec_file, spec_name):
    run_trial_test(spec_file, spec_name)


@pytest.mark.parametrize('spec_file,spec_name', [
    ('experimental/dqn.json', 'vanilla_dqn_cartpole'),
    ('experimental/dqn.json', 'dqn_boltzmann_cartpole'),
    ('experimental/dqn.json', 'dqn_epsilon_greedy_cartpole'),
    ('experimental/dqn.json', 'drqn_boltzmann_cartpole'),
    ('experimental/dqn.json', 'drqn_epsilon_greedy_cartpole'),
    # ('experimental/dqn.json', 'dqn_boltzmann_breakout'),
    # ('experimental/dqn.json', 'dqn_epsilon_greedy_breakout'),
    ('experimental/dqn.json', 'dqn_stack_epsilon_greedy_lunar'),
])
def test_dqn(spec_file, spec_name):
    run_trial_test(spec_file, spec_name)


@pytest.mark.parametrize('spec_file,spec_name', [
    ('experimental/ddqn.json', 'ddqn_boltzmann_cartpole'),
    ('experimental/ddqn.json', 'ddqn_epsilon_greedy_cartpole'),
    ('experimental/ddqn.json', 'ddrqn_boltzmann_cartpole'),
    ('experimental/ddqn.json', 'ddrqn_epsilon_greedy_cartpole'),
    # ('experimental/ddqn.json', 'ddqn_boltzmann_breakout'),
    # ('experimental/ddqn.json', 'ddqn_epsilon_greedy_breakout'),
])
def test_ddqn(spec_file, spec_name):
    run_trial_test(spec_file, spec_name)


@pytest.mark.parametrize('spec_file,spec_name', [
    ('experimental/dueling_dqn.json', 'dueling_dqn_boltzmann_cartpole'),
    ('experimental/dueling_dqn.json', 'dueling_dqn_epsilon_greedy_cartpole'),
    # ('experimental/dueling_dqn.json', 'dueling_dqn_boltzmann_breakout'),
    # ('experimental/dueling_dqn.json', 'dueling_dqn_epsilon_greedy_breakout'),
])
def test_dueling_dqn(spec_file, spec_name):
    run_trial_test(spec_file, spec_name)


@pytest.mark.parametrize('spec_file,spec_name', [
    ('experimental/hydra_dqn.json', 'hydra_dqn_boltzmann_cartpole'),
    ('experimental/hydra_dqn.json', 'hydra_dqn_epsilon_greedy_cartpole'),
    # ('experimental/hydra_dqn.json', 'hydra_dqn_epsilon_greedy_cartpole_2dball'),
])
def test_hydra_dqn(spec_file, spec_name):
    run_trial_test(spec_file, spec_name)


@flaky
@pytest.mark.parametrize('spec_file,spec_name', [
    ('experimental/dqn.json', 'dqn_pong'),
    # ('experimental/a2c.json', 'a2c_pong'),
])
def test_atari(spec_file, spec_name):
    run_trial_test(spec_file, spec_name)


@pytest.mark.parametrize('spec_file,spec_name', [
    ('experimental/reinforce.json', 'reinforce_conv_vizdoom'),
])
def test_reinforce_vizdoom(spec_file, spec_name):
    run_trial_test(spec_file, spec_name)


@pytest.mark.parametrize('spec_file,spec_name', [
    ('base.json', 'base_case_unity'),
    ('base.json', 'base_case_openai'),
    ('random.json', 'random_cartpole'),
    ('random.json', 'random_pendulum'),
    # ('base.json', 'multi_agent'),
    # ('base.json', 'multi_agent_multi_env'),
])
def test_base(spec_file, spec_name):
    run_trial_test(spec_file, spec_name)


@pytest.mark.parametrize('spec_file,spec_name', [
    ('base.json', 'multi_body'),
    ('base.json', 'multi_env'),
])
def test_base_multi(spec_file, spec_name):
    run_trial_test(spec_file, spec_name)

[(23, 'pytest.mark.parametrize', 'pytest.mark.parametrize', (['"""spec_file,spec_name"""', "[('experimental/reinforce.json', 'reinforce_mlp_cartpole'), (\n    'experimental/reinforce.json', 'reinforce_rnn_cartpole')]"], {}), False, 'import pytest\n'), (32, 'pytest.mark.parametrize', 'pytest.mark.parametrize', (['"""spec_file,spec_name"""', "[('experimental/reinforce.json', 'reinforce_mlp_pendulum'), (\n    'experimental/reinforce.json', 'reinforce_rnn_pendulum')]"], {}), False, 'import pytest\n'), (40, 'pytest.mark.parametrize', 'pytest.mark.parametrize', (['"""spec_file,spec_name"""', "[('experimental/a2c.json', 'a2c_mlp_shared_cartpole'), (\n    'experimental/a2c.json', 'a2c_mlp_separate_cartpole'), (\n    'experimental/a2c.json', 'a2c_rnn_shared_cartpole'), (\n    'experimental/a2c.json', 'a2c_rnn_separate_cartpole'), (\n    'experimental/a2c.json', 'a2c_mlp_concat_cartpole')]"], {}), False, 'import pytest\n'), (53, 'pytest.mark.parametrize', 'pytest.mark.parametrize', (['"""spec_file,spec_name"""', "[('experimental/a2c.json', 'a2c_mlp_shared_pendulum'), (\n    'experimental/a2c.json', 'a2c_mlp_separate_pendulum'), (\n    'experimental/a2c.json', 'a2c_rnn_shared_pendulum'), (\n    'experimental/a2c.json', 'a2c_rnn_separate_pendulum')]"], {}), False, 'import pytest\n'), (63, 'pytest.mark.parametrize', 'pytest.mark.parametrize', (['"""spec_file,spec_name"""', "[('experimental/ppo.json', 'ppo_mlp_shared_cartpole'), (\n    'experimental/ppo.json', 'ppo_mlp_separate_cartpole'), (\n    'experimental/ppo.json', 'ppo_rnn_shared_cartpole'), (\n    'experimental/ppo.json', 'ppo_rnn_separate_cartpole')]"], {}), False, 'import pytest\n'), (75, 'pytest.mark.parametrize', 'pytest.mark.parametrize', (['"""spec_file,spec_name"""', "[('experimental/ppo.json', 'ppo_mlp_shared_pendulum'), (\n    'experimental/ppo.json', 'ppo_mlp_separate_pendulum'), (\n    'experimental/ppo.json', 'ppo_rnn_shared_pendulum'), (\n    'experimental/ppo.json', 'ppo_rnn_separate_pendulum')]"], {}), False, 'import pytest\n'), (86, 'pytest.mark.parametrize', 'pytest.mark.parametrize', (['"""spec_file,spec_name"""', "[('experimental/ppo_sil.json', 'ppo_sil_mlp_shared_cartpole'), (\n    'experimental/ppo_sil.json', 'ppo_sil_mlp_separate_cartpole'), (\n    'experimental/ppo_sil.json', 'ppo_sil_rnn_shared_cartpole'), (\n    'experimental/ppo_sil.json', 'ppo_sil_rnn_separate_cartpole')]"], {}), False, 'import pytest\n'), (97, 'pytest.mark.parametrize', 'pytest.mark.parametrize', (['"""spec_file,spec_name"""', "[('experimental/ppo_sil.json', 'ppo_sil_mlp_shared_pendulum'), (\n    'experimental/ppo_sil.json', 'ppo_sil_mlp_separate_pendulum'), (\n    'experimental/ppo_sil.json', 'ppo_sil_rnn_shared_pendulum'), (\n    'experimental/ppo_sil.json', 'ppo_sil_rnn_separate_pendulum')]"], {}), False, 'import pytest\n'), (108, 'pytest.mark.parametrize', 'pytest.mark.parametrize', (['"""spec_file,spec_name"""', "[('experimental/sil.json', 'sil_mlp_shared_cartpole'), (\n    'experimental/sil.json', 'sil_mlp_separate_cartpole'), (\n    'experimental/sil.json', 'sil_rnn_shared_cartpole'), (\n    'experimental/sil.json', 'sil_rnn_separate_cartpole')]"], {}), False, 'import pytest\n'), (121, 'pytest.mark.parametrize', 'pytest.mark.parametrize', (['"""spec_file,spec_name"""', "[('experimental/sil.json', 'sil_mlp_shared_pendulum'), (\n    'experimental/sil.json', 'sil_mlp_separate_pendulum'), (\n    'experimental/sil.json', 'sil_rnn_shared_pendulum'), (\n    'experimental/sil.json', 'sil_rnn_separate_pendulum')]"], {}), False, 'import pytest\n'), (131, 'pytest.mark.parametrize', 'pytest.mark.parametrize', (['"""spec_file,spec_name"""', "[('experimental/sarsa.json', 'sarsa_mlp_boltzmann_cartpole'), (\n    'experimental/sarsa.json', 'sarsa_mlp_epsilon_greedy_cartpole'), (\n    'experimental/sarsa.json', 'sarsa_rnn_boltzmann_cartpole'), (\n    'experimental/sarsa.json', 'sarsa_rnn_epsilon_greedy_cartpole')]"], {}), False, 'import pytest\n'), (143, 'pytest.mark.parametrize', 'pytest.mark.parametrize', (['"""spec_file,spec_name"""', "[('experimental/dqn.json', 'vanilla_dqn_cartpole'), (\n    'experimental/dqn.json', 'dqn_boltzmann_cartpole'), (\n    'experimental/dqn.json', 'dqn_epsilon_greedy_cartpole'), (\n    'experimental/dqn.json', 'drqn_boltzmann_cartpole'), (\n    'experimental/dqn.json', 'drqn_epsilon_greedy_cartpole'), (\n    'experimental/dqn.json', 'dqn_stack_epsilon_greedy_lunar')]"], {}), False, 'import pytest\n'), (157, 'pytest.mark.parametrize', 'pytest.mark.parametrize', (['"""spec_file,spec_name"""', "[('experimental/ddqn.json', 'ddqn_boltzmann_cartpole'), (\n    'experimental/ddqn.json', 'ddqn_epsilon_greedy_cartpole'), (\n    'experimental/ddqn.json', 'ddrqn_boltzmann_cartpole'), (\n    'experimental/ddqn.json', 'ddrqn_epsilon_greedy_cartpole')]"], {}), False, 'import pytest\n'), (169, 'pytest.mark.parametrize', 'pytest.mark.parametrize', (['"""spec_file,spec_name"""', "[('experimental/dueling_dqn.json', 'dueling_dqn_boltzmann_cartpole'), (\n    'experimental/dueling_dqn.json', 'dueling_dqn_epsilon_greedy_cartpole')]"], {}), False, 'import pytest\n'), (179, 'pytest.mark.parametrize', 'pytest.mark.parametrize', (['"""spec_file,spec_name"""', "[('experimental/hydra_dqn.json', 'hydra_dqn_boltzmann_cartpole'), (\n    'experimental/hydra_dqn.json', 'hydra_dqn_epsilon_greedy_cartpole')]"], {}), False, 'import pytest\n'), (189, 'pytest.mark.parametrize', 'pytest.mark.parametrize', (['"""spec_file,spec_name"""', "[('experimental/dqn.json', 'dqn_pong')]"], {}), False, 'import pytest\n'), (197, 'pytest.mark.parametrize', 'pytest.mark.parametrize', (['"""spec_file,spec_name"""', "[('experimental/reinforce.json', 'reinforce_conv_vizdoom')]"], {}), False, 'import pytest\n'), (204, 'pytest.mark.parametrize', 'pytest.mark.parametrize', (['"""spec_file,spec_name"""', "[('base.json', 'base_case_unity'), ('base.json', 'base_case_openai'), (\n    'random.json', 'random_cartpole'), ('random.json', 'random_pendulum')]"], {}), False, 'import pytest\n'), (216, 'pytest.mark.parametrize', 'pytest.mark.parametrize', (['"""spec_file,spec_name"""', "[('base.json', 'multi_body'), ('base.json', 'multi_env')]"], {}), False, 'import pytest\n'), (14, 'slm_lab.spec.spec_util.get', 'spec_util.get', (['spec_file', 'spec_name'], {}), False, 'from slm_lab.spec import spec_util\n'), (15, 'slm_lab.spec.spec_util.override_test_spec', 'spec_util.override_test_spec', (['spec'], {}), False, 'from slm_lab.spec import spec_util\n'), (16, 'slm_lab.experiment.monitor.InfoSpace', 'InfoSpace', ([], {}), False, 'from slm_lab.experiment.monitor import InfoSpace\n'), (18, 'slm_lab.experiment.control.Trial', 'Trial', (['spec', 'info_space'], {}), False, 'from slm_lab.experiment.control import Trial\n')]