from libensemble.message_numbers import WORKER_DONE
from libensemble.executors.executor import Executor
import numpy as np


def circ_offset(wid, length):
    offset = length // 2 + (length % 2 > 0)
    circ_index = wid + offset
    if circ_index > length:
        circ_index = circ_index - length
    return circ_index


def exp_nodelist_for_worker(exp_list, workerID, nodes_per_worker, persis_gens):
    """Modify expected node-lists based on workerID"""
    comps = exp_list.split()
    new_line = []
    for comp in comps:
        if comp.startswith("node-"):
            new_node_list = []
            node_list = comp.split(",")
            for node in node_list:
                node_name, node_num = node.split("-")
                offset = workerID - (1 + persis_gens)
                new_num = int(node_num) + int(nodes_per_worker * offset)
                new_node = "-".join([node_name, str(new_num)])
                new_node_list.append(new_node)
            new_list = ",".join(new_node_list)
            new_line.append(new_list)
        else:
            new_line.append(comp)
    return " ".join(new_line)


def runline_check(H, persis_info, sim_specs, libE_info):
    """Check run-lines produced by executor provided by a list of tests"""
    calc_status = 0
    x = H["x"][0][0]
    exctr = Executor.executor
    test_list = sim_specs["user"]["tests"]
    exp_list = sim_specs["user"]["expect"]
    npw = sim_specs["user"]["nodes_per_worker"]
    p_gens = sim_specs["user"].get("persis_gens", 0)

    for i, test in enumerate(test_list):
        task = exctr.submit(
            calc_type="sim",
            num_procs=test.get("nprocs", None),
            num_nodes=test.get("nnodes", None),
            procs_per_node=test.get("ppn", None),
            extra_args=test.get("e_args", None),
            app_args="--testid " + test.get("testid", None),
            stdout="out.txt",
            stderr="err.txt",
            hyperthreads=test.get("ht", None),
            dry_run=True,
        )

        outline = task.runline
        new_exp_list = exp_nodelist_for_worker(exp_list[i], libE_info["workerID"], npw, p_gens)

        if outline != new_exp_list:
            print("outline is: {}\nexp     is: {}".format(outline, new_exp_list), flush=True)

        assert outline == new_exp_list

    calc_status = WORKER_DONE
    output = np.zeros(1, dtype=sim_specs["out"])
    output["f"][0] = np.linalg.norm(x)
    return output, persis_info, calc_status


def runline_check_by_worker(H, persis_info, sim_specs, libE_info):
    """Check run-lines produced by executor provided by a list of lines per worker"""
    offset_for_min_rsets_scheduler = sim_specs["user"].get("offset_for_scheduler", False)
    calc_status = 0
    x = H["x"][0][0]
    exctr = Executor.executor
    test = sim_specs["user"]["tests"][0]
    exp_list = sim_specs["user"]["expect"]
    p_gens = sim_specs["user"].get("persis_gens", 0)

    task = exctr.submit(
        calc_type="sim",
        num_procs=test.get("nprocs", None),
        num_nodes=test.get("nnodes", None),
        procs_per_node=test.get("ppn", None),
        extra_args=test.get("e_args", None),
        app_args="--testid " + test.get("testid", None),
        stdout="out.txt",
        stderr="err.txt",
        hyperthreads=test.get("ht", None),
        dry_run=True,
    )

    outline = task.runline
    wid = libE_info["workerID"]

    # Adjust for minimum slot scheduling in alloc func: e.g. 5 rsets on 2 nodes
    # Node 1: 3 rsets (0,1,2). Node 2: 2 rsets (3,4)
    # The first sim will got to rset 3 as it finds a "smaller slot".
    # Alternative would be for splitter to use opposite splits (e.g. 2,3 rather than 3,2)
    if offset_for_min_rsets_scheduler:
        wid_mod = circ_offset(wid, len(exp_list))
    else:
        wid_mod = wid

    new_exp_list = exp_list[wid_mod - 1 - p_gens]

    if outline != new_exp_list:
        print("Worker {}:\n outline is: {}\n exp     is: {}".format(wid, outline, new_exp_list), flush=True)

    assert outline == new_exp_list

    calc_status = WORKER_DONE
    output = np.zeros(1, dtype=sim_specs["out"])
    output["f"][0] = np.linalg.norm(x)
    return output, persis_info, calc_status
