# import basico
# Lazy import basico
from mobspy.import_manager.lazy_import_class import LazyImporter as ipm_LazyImporter

basico = ipm_LazyImporter("basico")
from joblib import Parallel, delayed
import mobspy.simulation_logging.log_scripts as simlog
import mobspy.sbml_simulator.builder as sbml_builder


# def simulate(list_of_params, models)
def simulate(jobs, list_of_params, models):
    """
    This function coordinates the simulation by calling the necessary jobs
    In the future we hope to implement parallel cluster computing compatibility

    :param list_of_params: (dict) simulation parameters from the text file
    :param models: (list) [{'species_for_sbml':, 'parameters_for_sbml':, 'reactions_for_sbml':,
                        'events_for_sbml':, 'species_not_mapped':, 'mappings':}]
    :return: data (dict) = dictionary containing the resulting data from simulation
    """

    # Run in parallel or sequentially
    # If nothing is specified just run it in parallel
    data = job_execution(list_of_params, models, jobs)

    return data


def job_execution(params, models, jobs):
    """
    This is defined for parallelism purposes
    Uses multiple cores from the processor to execute stochastic simulations

    :param params: (dict) = simulation parameters
    :param models: (list) [{'species_for_sbml':, 'parameters_for_sbml':, 'reactions_for_sbml':,
                        'events_for_sbml':, 'species_not_mapped':, 'mappings':}]
    :param jobs: (int) = number of cores to use, -1 for all available

    :return: parallel_data data of all the individual simulations executed in parallel
    """

    def __single_run(packed):
        i = packed

        added_data = {}
        # This is only need to avoid git errors
        reformatted_data = {}
        for j, (sim_par, model) in enumerate(zip(params, models)):
            # Generate SBML here
            if j > 0:
                sbml_str = __sbml_new_initial_values(
                    reformatted_data, model, sim_par, new_model=True
                )
            else:
                sbml_str = __sbml_new_initial_values({}, model, sim_par)

            end_condition_not_satisfied = True
            if sim_par["_continuous_simulation"]:
                duration = float(sim_par["initial_conditional_duration"])
            else:
                duration = float(sim_par["duration"])

            while end_condition_not_satisfied:
                basico_model = basico.model_io.load_model_from_string(sbml_str)
                data = __run_time_course(basico_model, duration, sim_par, i)

                reformatted_data = reformat_time_series(data)

                if sim_par["_continuous_simulation"]:
                    if reformatted_data["_End_Flag_MetaSpecies"][-1] > 0:
                        reformatted_data = __filter_condition_event_time_data(
                            reformatted_data
                        )
                        end_condition_not_satisfied = False
                    else:
                        sbml_str = __sbml_new_initial_values(
                            reformatted_data, model, sim_par
                        )
                        duration = 2 * duration
                else:
                    end_condition_not_satisfied = False

                reformatted_data = __remap_species(
                    reformatted_data, model["mappings"], model["species_for_sbml"]
                )
                added_data = __add_simulations_data(added_data, reformatted_data)

        return added_data

    parallel_data = Parallel(n_jobs=jobs)(
        delayed(__single_run)(i) for i in range(params[0]["repetitions"])
    )

    if not parallel_data:
        simlog.error(
            "Error: The parallel model has not produced an output."
            + "Try addding ('sequential': True) to parameters"
        )

    return parallel_data


def __run_time_course(basico_model, duration, params, index):
    """
    This function prepares the parameters for the basiCO.run_time_course

    :param basico_model: basico model to be run in the simulation
    :param params: (dict) simulation parameters
    :param index: (int) current index of the run
    :param duration: (int, float) simulation duration in seconds
    """
    params["simulation_method"] = params["simulation_method"].lower()
    if (params["_with_event"] or params["_continuous_simulation"]) and params[
        "simulation_method"
    ] == "stochastic":
        params["simulation_method"] = "directmethod"

    kargs = {
        "model": basico_model,
        "method": params["simulation_method"],
        "start_time": params["start_time"],
        "r_tol": params["r_tol"],
        "a_tol": params["a_tol"],
        "output_event": params["output_event"],
    }

    if "seeds" in params:
        kargs["use_seed"] = True
        kargs["seed"] = params["seeds"][index]

    if "step_size" in params:
        kargs["automatic"] = False
        kargs["step_number"] = int(params["duration"] / params["step_size"])

    return basico.run_time_course(duration, **kargs)


def reformat_time_series(data):
    """
    Transforms the _dot_ in from the results into .

    :param data: (pd.dataframe) simulation data results from BasiCO
    """
    data_dict = {"Time": data.index.tolist()}

    for key in data:
        data_dict[key.replace("_dot_", ".")] = list(data[key])

    return data_dict


def __filter_condition_event_time_data(data):
    """
    This function filters all the data until the event was triggered. It is used for the conditional duration
    simulations

    :param data: (dict) single run data with species as keys and values
    """
    new_data = {}

    for i, e in enumerate(data["_End_Flag_MetaSpecies"]):
        if e == 1:
            stop_index = i
            break

    for key in data:
        new_data[key] = data[key][: stop_index + 1]

    return new_data


def __sbml_new_initial_values(data, model, sim_para, new_model=False):
    """
    This function adjusts the new sbml counts to the end of the previous concatenation simulation or conditional
    simulation

    :param data: (dict) single run data with species as keys and values
    :param model: (dict) model to be used in the next simulation
    :param model: (sim_par) parameters to be used in the next simulation
    :param new_model: (bool) is a new model being used or the model is the same as the old one
    """
    species_for_sbml = model["species_for_sbml"]

    check_list = ["stochastic", "directmethod"]
    for key in data:
        sbml_key = key.replace(".", "_dot_")
        if sbml_key not in species_for_sbml.keys():
            continue

        if key == "Time":
            continue
        try:
            # Case of species set
            if sim_para["simulation_method"].lower() in check_list:
                species_for_sbml[sbml_key] = int(list(data[key])[-1])
            else:
                species_for_sbml[sbml_key] = list(data[key])[-1]
        except KeyError:
            pass

    if new_model:
        try:
            species_for_sbml["_End_Flag_MetaSpecies"] = 0
        except KeyError:
            pass

    return sbml_builder.build(
        species_for_sbml,
        model["parameters_for_sbml"],
        model["reactions_for_sbml"],
        model["events_for_sbml"],
        model["assignments_for_sbml"],
    )


def __add_simulations_data(added_data, reformatted_data):
    """
    Adds the data from the new executed simulation to the data stored so far

    :param added_data: (dict) added data so far in the simulation (key: species string - value: run)
    :param reformatted_data: (dict) data return from the simulation after being formatted using the reformat
    data function
    """
    if added_data != {}:
        time_to_add = added_data["Time"][-1]
    else:
        time_to_add = 0
    new_data = {}
    already_added_keys = set()

    for key in added_data:
        new_data[key] = added_data[key]

    for i, time in enumerate(reformatted_data["Time"]):
        # Remove the repeated initial value from following simulations.
        # The final value of summed simulations is repeated
        if time == 0 and added_data != {}:
            for key in reformatted_data:
                reformatted_data[key].pop(0)

        # This is for simulations with only one value. The system did not change at all
        try:
            reformatted_data["Time"][i] = reformatted_data["Time"][i] + time_to_add
        except IndexError:
            pass

    for key in added_data:
        if key == "Time":
            continue
        try:
            new_data[key] = added_data[key] + reformatted_data[key]
            already_added_keys.add(key)
        except KeyError:
            dummy = added_data[key][-1]
            new_data[key] += [dummy for _ in reformatted_data["Time"]]

    for key in reformatted_data:
        if key == "Time":
            continue

        if key in already_added_keys:
            continue
        else:
            if time_to_add != 0:
                new_data[key] = [0 for _ in added_data["Time"]]
                new_data[key] = new_data[key] + reformatted_data[key]
            else:
                new_data[key] = reformatted_data[key]

    if time_to_add != 0:
        new_data["Time"] = added_data["Time"] + reformatted_data["Time"]
    else:
        new_data["Time"] = reformatted_data["Time"]

    return new_data


def __remap_species(data, mapping, species_not_mapped):
    """
    Takes the simulated species (data) and add ones defined by
    mapping (mapping).
    By default, a mapping is a list of species in which case the
    sum is taken: ['a', 'b', ...]
    It also adds species that have been removed from the basico simulation

    Parameters:
    :param data: (dict) Data in MobsPy format
    :param mapping: (dict) = Mappings between species and meta-species
    :param species_not_mapped: (list of str) Species not mapped by BasiCO
    """

    mapped_data = {"Time": data["Time"]}
    T = range(len(data["Time"]))

    # copy over all unmapped ones
    for k in data.keys():
        mapped_data[k] = data[k]

    dot_species_not_mapped = {}
    for key in species_not_mapped:
        dot_species_not_mapped[key.replace("_dot_", ".")] = species_not_mapped[key]

    # 1st pass with sum mappings
    for group in mapping.keys():
        the_mapping = mapping[group]
        mapped_data[group] = {"runs": []}

        try:
            # check if is a list -> sum
            if type(the_mapping) is list:
                this_run = []
                runs_not_returned_by_basico = {}
                for t in T:
                    mapping_sum = 0
                    for spe in the_mapping:
                        try:
                            mapping_sum = mapping_sum + data[spe][t]
                        except KeyError:
                            mapping_sum = mapping_sum + dot_species_not_mapped[spe]
                            try:
                                runs_not_returned_by_basico[spe] += [
                                    dot_species_not_mapped[spe]
                                ]
                            except KeyError:
                                runs_not_returned_by_basico[spe] = [
                                    dot_species_not_mapped[spe]
                                ]
                    this_run.append(mapping_sum)
                for spe in runs_not_returned_by_basico:
                    try:
                        mapped_data[spe] = runs_not_returned_by_basico[spe]
                    except KeyError:
                        mapped_data[spe] = [runs_not_returned_by_basico[spe]]
                mapped_data[group] = this_run

        except IndexError:
            simlog.error(
                f'run: remap_species: error when remapping "{the_mapping}".'
                + "Possible fix: All runs must have the same time"
            )

        except TypeError:
            simlog.warning(
                "Copasi removes A >> A species from reaction calculations and does not provide an output"
            )
            simlog.warning(
                "Please check the output data to see if this is the problem"
            )
            for key in data:
                print(key, data[key])
            exit(1)

    return mapped_data
