import os
from libensemble.resources.env_resources import EnvResources


def setup_standalone_run():
    os.environ["LIBE_RESOURCES_TEST_NODE_LIST"] = ""


def teardown_standalone_run():
    os.environ["LIBE_RESOURCES_TEST_NODE_LIST"] = ""


def setup_function(function):
    print("setup_function    function:%s" % function.__name__)
    os.environ["LIBE_RESOURCES_TEST_NODE_LIST"] = ""


def teardown_function(function):
    print("teardown_function    function:%s" % function.__name__)
    os.environ["LIBE_RESOURCES_TEST_NODE_LIST"] = ""


# Tests ========================================================================================

# Tests for obtaining nodelist from environment variables


def test_slurm_nodelist_empty():
    os.environ["LIBE_RESOURCES_TEST_NODE_LIST"] = ""
    exp_out = []  # empty
    nodelist = EnvResources.get_slurm_nodelist(node_list_env="LIBE_RESOURCES_TEST_NODE_LIST")
    assert nodelist == exp_out, "Nodelist returned does not match expected"


def test_slurm_nodelist_single():
    os.environ["LIBE_RESOURCES_TEST_NODE_LIST"] = "knl-0056"
    exp_out = ["knl-0056"]
    nodelist = EnvResources.get_slurm_nodelist(node_list_env="LIBE_RESOURCES_TEST_NODE_LIST")
    assert nodelist == exp_out, "Nodelist returned does not match expected"


def test_slurm_nodelist_single_nodash():
    os.environ["LIBE_RESOURCES_TEST_NODE_LIST"] = "nid00056"
    exp_out = ["nid00056"]
    nodelist = EnvResources.get_slurm_nodelist(node_list_env="LIBE_RESOURCES_TEST_NODE_LIST")
    assert nodelist == exp_out, "Nodelist returned does not match expected"


def test_slurm_nodelist_knl_seq():
    os.environ["LIBE_RESOURCES_TEST_NODE_LIST"] = "knl-[0009-0012]"
    exp_out = ["knl-0009", "knl-0010", "knl-0011", "knl-0012"]
    nodelist = EnvResources.get_slurm_nodelist(node_list_env="LIBE_RESOURCES_TEST_NODE_LIST")
    assert nodelist == exp_out, "Nodelist returned does not match expected"


def test_slurm_nodelist_bdw_seq():
    os.environ["LIBE_RESOURCES_TEST_NODE_LIST"] = "bdw-[0009-0012]"
    exp_out = ["bdw-0009", "bdw-0010", "bdw-0011", "bdw-0012"]
    nodelist = EnvResources.get_slurm_nodelist(node_list_env="LIBE_RESOURCES_TEST_NODE_LIST")
    assert nodelist == exp_out, "Nodelist returned does not match expected"


def test_slurm_nodelist_groups():
    os.environ["LIBE_RESOURCES_TEST_NODE_LIST"] = "knl-[0020-0022,0137-0139,1234]"
    exp_out = ["knl-0020", "knl-0021", "knl-0022", "knl-0137", "knl-0138", "knl-0139", "knl-1234"]
    nodelist = EnvResources.get_slurm_nodelist(node_list_env="LIBE_RESOURCES_TEST_NODE_LIST")
    assert nodelist == exp_out, "Nodelist returned does not match expected"


def test_slurm_nodelist_groups_partitions():
    os.environ["LIBE_RESOURCES_TEST_NODE_LIST"] = "bdw-[0254,0384,0565-0568],bdwd-[0004,0009]"
    exp_out = ["bdw-0254", "bdw-0384", "bdw-0565", "bdw-0566", "bdw-0567", "bdw-0568", "bdwd-0004", "bdwd-0009"]
    nodelist = EnvResources.get_slurm_nodelist(node_list_env="LIBE_RESOURCES_TEST_NODE_LIST")
    assert nodelist == exp_out, "Nodelist returned does not match expected"


def test_slurm_nodelist_groups_nodash():
    os.environ["LIBE_RESOURCES_TEST_NODE_LIST"] = "nid0[0020-0022,0137-0139,1234]"
    exp_out = ["nid00020", "nid00021", "nid00022", "nid00137", "nid00138", "nid00139", "nid01234"]
    nodelist = EnvResources.get_slurm_nodelist(node_list_env="LIBE_RESOURCES_TEST_NODE_LIST")
    assert nodelist == exp_out, "Nodelist returned does not match expected"


def test_slurm_nodelist_groups_longprefix():
    os.environ["LIBE_RESOURCES_TEST_NODE_LIST"] = "super-[000020-000022,000137-000139,001234,023456-023458]"
    exp_out = [
        "super-000020",
        "super-000021",
        "super-000022",
        "super-000137",
        "super-000138",
        "super-000139",
        "super-001234",
        "super-023456",
        "super-023457",
        "super-023458",
    ]
    nodelist = EnvResources.get_slurm_nodelist(node_list_env="LIBE_RESOURCES_TEST_NODE_LIST")
    assert nodelist == exp_out, "Nodelist returned does not match expected"


def test_slurm_nodelist_reverse_grp():
    os.environ["LIBE_RESOURCES_TEST_NODE_LIST"] = "knl-[0020-0022,0139-0137,1234]"
    exp_out = ["knl-0020", "knl-0021", "knl-0022", "knl-0137", "knl-0138", "knl-0139", "knl-1234"]
    nodelist = EnvResources.get_slurm_nodelist(node_list_env="LIBE_RESOURCES_TEST_NODE_LIST")
    assert nodelist == exp_out, "Nodelist returned does not match expected"


def test_cobalt_nodelist_empty():
    os.environ["LIBE_RESOURCES_TEST_NODE_LIST"] = ""
    exp_out = []  # empty
    nodelist = EnvResources.get_cobalt_nodelist(node_list_env="LIBE_RESOURCES_TEST_NODE_LIST")
    assert nodelist == exp_out, "Nodelist returned does not match expected"


def test_cobalt_nodelist_single():
    os.environ["LIBE_RESOURCES_TEST_NODE_LIST"] = "56"
    exp_out = ["56"]
    nodelist = EnvResources.get_cobalt_nodelist(node_list_env="LIBE_RESOURCES_TEST_NODE_LIST")
    assert nodelist == exp_out, "Nodelist returned does not match expected"


def test_cobalt_nodelist_seq():
    os.environ["LIBE_RESOURCES_TEST_NODE_LIST"] = "9-12"
    exp_out = ["9", "10", "11", "12"]
    nodelist = EnvResources.get_cobalt_nodelist(node_list_env="LIBE_RESOURCES_TEST_NODE_LIST")
    assert nodelist == exp_out, "Nodelist returned does not match expected"


def test_cobalt_nodelist_groups():
    os.environ["LIBE_RESOURCES_TEST_NODE_LIST"] = "20-22,137-139,1234"
    exp_out = ["20", "21", "22", "137", "138", "139", "1234"]
    nodelist = EnvResources.get_cobalt_nodelist(node_list_env="LIBE_RESOURCES_TEST_NODE_LIST")
    assert nodelist == exp_out, "Nodelist returned does not match expected"


def test_cobalt_nodelist_reverse_grp():
    os.environ["LIBE_RESOURCES_TEST_NODE_LIST"] = "20-22,139-137,1234"
    exp_out = ["20", "21", "22", "137", "138", "139", "1234"]
    nodelist = EnvResources.get_cobalt_nodelist(node_list_env="LIBE_RESOURCES_TEST_NODE_LIST")
    assert nodelist == exp_out, "Nodelist returned does not match expected"


def test_lsf_nodelist_empty():
    os.environ["LIBE_RESOURCES_TEST_NODE_LIST"] = ""
    exp_out = []  # empty
    nodelist = EnvResources.get_lsf_nodelist(node_list_env="LIBE_RESOURCES_TEST_NODE_LIST")
    assert nodelist == exp_out, "Nodelist returned does not match expected"


def test_lsf_nodelist_single():
    os.environ["LIBE_RESOURCES_TEST_NODE_LIST"] = "batch5" + " g06n02" * 42
    exp_out = ["g06n02"]
    nodelist = EnvResources.get_lsf_nodelist(node_list_env="LIBE_RESOURCES_TEST_NODE_LIST")
    assert nodelist == exp_out, "Nodelist returned does not match expected"


def test_lsf_nodelist_seq():
    os.environ["LIBE_RESOURCES_TEST_NODE_LIST"] = "batch5" + " g06n02" * 42 + " h21n18" * 42
    exp_out = ["g06n02", "h21n18"]
    nodelist = EnvResources.get_lsf_nodelist(node_list_env="LIBE_RESOURCES_TEST_NODE_LIST")
    assert nodelist == exp_out, "Nodelist returned does not match expected"


# These dont apply to the lsf lists as they are listed in full
# def test_lsf_nodelist_groups():
# def test_lsf_nodelist_reverse_grp():


def test_lsf_nodelist_shortform_empty():
    os.environ["LIBE_RESOURCES_TEST_NODE_LIST"] = ""
    exp_out = []  # empty
    nodelist = EnvResources.get_lsf_nodelist_frm_shortform(node_list_env="LIBE_RESOURCES_TEST_NODE_LIST")
    assert nodelist == exp_out, "Nodelist returned does not match expected"


def test_lsf_nodelist_shortform_single():
    os.environ["LIBE_RESOURCES_TEST_NODE_LIST"] = "batch5 1 g06n02 42"
    exp_out = ["g06n02"]
    nodelist = EnvResources.get_lsf_nodelist_frm_shortform(node_list_env="LIBE_RESOURCES_TEST_NODE_LIST")
    assert nodelist == exp_out, "Nodelist returned does not match expected"


def test_lsf_nodelist_shortform_seq():
    os.environ["LIBE_RESOURCES_TEST_NODE_LIST"] = "batch5 1 g06n02 42 h21n18 42"
    exp_out = ["g06n02", "h21n18"]
    nodelist = EnvResources.get_lsf_nodelist_frm_shortform(node_list_env="LIBE_RESOURCES_TEST_NODE_LIST")
    assert nodelist == exp_out, "Nodelist returned does not match expected"


def test_shortnames_nochange_slurm():
    env_resources = EnvResources()
    # Test Slurm abbrev
    exp_names = ["knl-0019", "knl-0021", "knl-0022", "knl-0137", "knl-0138", "knl-0139", "knl-2345"]
    env_resources.scheduler = "Slurm"
    abbrev_names = env_resources.shortnames(exp_names)
    assert abbrev_names == exp_names, "Abbreviated names returned do not match expected"
    del env_resources


def test_shortnames_slurm():
    env_resources = EnvResources()
    # Test Slurm abbrev
    exp_names = ["knl-0019", "knl-0021", "knl-0022"]
    full_names = ["knl-0019.some.suffix", "knl-0021.some.suffix", "knl-0022.diff_suffix"]
    env_resources.scheduler = "Slurm"
    abbrev_names = env_resources.shortnames(full_names)
    assert abbrev_names == exp_names, "Abbreviated names returned do not match expected"
    del env_resources


def test_shortnames_nochange_cobalt():
    env_resources = EnvResources()
    # Test Cobalt abbrev
    exp_names = ["21", "22", "137", "138", "1234", "11234"]
    env_resources.scheduler = "Cobalt"
    abbrev_names = env_resources.shortnames(exp_names)
    assert abbrev_names == exp_names, "Abbreviated names returned do not match expected"
    del env_resources


def test_shortnames_cobalt():
    env_resources = EnvResources()
    # Test Cobalt abbrev
    exp_names = ["20", "21", "22", "137", "138", "1234", "11234"]
    full_names = ["nid00020", "nid00021", "nid00022", "nid00137", "nid00138", "nid01234", "nid11234"]
    env_resources.scheduler = "Cobalt"
    abbrev_names = env_resources.shortnames(full_names)
    assert abbrev_names == exp_names, "Abbreviated names returned do not match expected"
    del env_resources


if __name__ == "__main__":
    setup_standalone_run()

    test_slurm_nodelist_empty()
    test_slurm_nodelist_single()
    test_slurm_nodelist_single_nodash()
    test_slurm_nodelist_knl_seq()
    test_slurm_nodelist_bdw_seq()
    test_slurm_nodelist_groups()
    test_slurm_nodelist_groups_partitions()
    test_slurm_nodelist_groups_nodash()
    test_slurm_nodelist_groups_longprefix()
    test_slurm_nodelist_reverse_grp()

    test_cobalt_nodelist_empty()
    test_cobalt_nodelist_single()
    test_cobalt_nodelist_seq()
    test_cobalt_nodelist_groups()
    test_cobalt_nodelist_reverse_grp()

    test_lsf_nodelist_empty()
    test_lsf_nodelist_single()
    test_lsf_nodelist_seq()

    test_lsf_nodelist_shortform_empty()
    test_lsf_nodelist_shortform_single()
    test_lsf_nodelist_shortform_seq()

    test_shortnames_nochange_slurm()
    test_shortnames_slurm()
    test_shortnames_nochange_cobalt()
    test_shortnames_cobalt()

    teardown_standalone_run()
