import os
import shutil
import numpy as np
from libensemble.output_directory import EnsembleDirectory
from libensemble.utils.loc_stack import LocationStack


def test_range_single_element():
    """Single H_row labeling"""

    work = {"H_fields": ["x", "num_nodes", "procs_per_node"], "libE_info": {"H_rows": np.array([5]), "workerID": 1}}
    assert EnsembleDirectory.extract_H_ranges(work) == "5", "Failed to correctly parse single H row"


def test_range_two_separate_elements():
    """Multiple H_rows, non-sequential"""

    work = {"H_fields": ["x", "num_nodes", "procs_per_node"], "libE_info": {"H_rows": np.array([2, 8]), "workerID": 1}}
    assert EnsembleDirectory.extract_H_ranges(work) == "2_8", "Failed to correctly parse nonsequential H rows"


def test_range_two_ranges():
    """Multiple sequences of H_rows"""

    work = {
        "H_fields": ["x", "num_nodes", "procs_per_node"],
        "libE_info": {"H_rows": np.array([0, 1, 2, 3, 7, 8]), "workerID": 1},
    }
    assert EnsembleDirectory.extract_H_ranges(work) == "0-3_7-8", "Failed to correctly parse multiple H ranges"


def test_range_mixes():
    """Mix of single rows and sequences of H_rows"""

    work = {
        "H_fields": ["x", "num_nodes", "procs_per_node"],
        "libE_info": {"H_rows": np.array([2, 3, 4, 6, 8, 9, 11, 14]), "workerID": 1},
    }
    assert (
        EnsembleDirectory.extract_H_ranges(work) == "2-4_6_8-9_11_14"
    ), "Failed to correctly parse H row single elements and ranges."


def test_copy_back():
    """When workers conclude their work, workers have the option of copying
    back their work into a directory created by the manager."""

    inputdir = "./calc"
    copybackdir = "./calc_back"
    inputfile = "./calc/file"

    for dir in [inputdir, copybackdir]:
        os.makedirs(dir, exist_ok=True)

    libE_specs = {"sim_dirs_make": True, "ensemble_dir_path": inputdir, "ensemble_copy_back": True}

    ls = LocationStack()
    ls.register_loc("test", inputfile)
    # fake_worker = FakeWorker(libE_specs, inputdir, '.', ls)
    ed = EnsembleDirectory(libE_specs, ls)
    EnsembleDirectory.copy_back(ed)
    assert "file" in os.listdir(copybackdir), "File not copied back to starting dir"

    for dir in [inputdir, copybackdir]:
        shutil.rmtree(dir)


def test_worker_dirs_but_no_sim_dirs():
    """Test Worker._make_calc_dir() directory structure without sim_dirs"""
    inputdir = "./calc"
    inputfile = "./calc/file"
    ensemble_dir = "./test_ens"

    for dir in [inputdir, inputfile, ensemble_dir]:
        os.makedirs(dir, exist_ok=True)

    libE_specs = {"ensemble_dir_path": ensemble_dir, "use_worker_dirs": True, "sim_input_dir": inputdir}

    ls = LocationStack()
    ed = EnsembleDirectory(libE_specs, ls)
    for i in range(4):  # Should work at any range
        EnsembleDirectory._make_calc_dir(ed, 1, 1, "sim", ls)

    assert "worker1" in os.listdir(ensemble_dir)
    assert "file" in os.listdir(os.path.join(ensemble_dir, "worker1"))

    for dir in [inputdir, ensemble_dir]:
        shutil.rmtree(dir)


def test_loc_stack_FileExists_exceptions():
    inputdir = "./calc"
    copyfile = "./calc/copy"
    symlinkfile = "./calc/symlink"
    ensemble_dir = "./test_ens"

    for dir in [inputdir, copyfile, symlinkfile]:
        os.makedirs(dir, exist_ok=True)

    # Testing loc_stack continuing on FileExistsError when not using sim_dirs
    libE_specs = {
        "sim_dirs_make": False,
        "ensemble_dir_path": ensemble_dir,
        "use_worker_dirs": True,
        "sim_dir_copy_files": [copyfile],
        "sim_dir_symlink_files": [symlinkfile],
    }

    ls = LocationStack()
    ed = EnsembleDirectory(libE_specs=libE_specs, loc_stack=ls)
    for i in range(2):  # Should work at any range
        EnsembleDirectory._make_calc_dir(ed, 1, "1", "sim", ls)

    assert len(os.listdir(ensemble_dir)) == 1, "Should only be a single worker file in ensemble"
    assert "worker1" in os.listdir(ensemble_dir), "Only directory should be worker1"
    assert all(
        [i in os.listdir(os.path.join(ensemble_dir, "worker1")) for i in ["copy", "symlink"]]
    ), "Files to copy and symlink not found in worker directory."

    # Testing loc_stack exception raising when sim_dir re-used - copy
    libE_specs = {
        "sim_dirs_make": True,
        "ensemble_dir_path": ensemble_dir,
        "use_worker_dirs": True,
        "sim_dir_copy_files": [copyfile],
    }

    flag = 1
    ed = EnsembleDirectory(libE_specs=libE_specs, loc_stack=ls)
    EnsembleDirectory._make_calc_dir(ed, 1, "1", "sim", ls)
    try:
        EnsembleDirectory._make_calc_dir(ed, 1, "1", "sim", ls)
    except FileExistsError:
        flag = 0
    assert flag == 0

    # Testing loc_stack exception raising when sim_dir re-used - symlink
    libE_specs = {
        "sim_dirs_make": True,
        "ensemble_dir_path": ensemble_dir,
        "use_worker_dirs": True,
        "sim_dir_symlink_files": [symlinkfile],
    }

    flag = 1
    ed = EnsembleDirectory(libE_specs=libE_specs, loc_stack=ls)
    EnsembleDirectory._make_calc_dir(ed, 1, "2", "sim", ls)
    try:
        EnsembleDirectory._make_calc_dir(ed, 1, "2", "sim", ls)
    except FileExistsError:
        flag = 0
    assert flag == 0

    for dir in [inputdir, ensemble_dir]:
        shutil.rmtree(dir)


if __name__ == "__main__":
    test_range_single_element()
    test_range_two_separate_elements()
    test_range_two_ranges()
    test_range_mixes()
    test_copy_back()
    test_worker_dirs_but_no_sim_dirs()
    test_loc_stack_FileExists_exceptions()
