from mobspy.simulation_logging.log_scripts import error as simlog_error
from inspect import stack as inspect_stack
from mobspy.modules.meta_class import List_Species, Species, Reacting_Species
from mobspy.modules.mobspy_parameters import (
    Internal_Parameter_Constructor as mp_Mobspy_Parameter,
)
from pint import Quantity
from numpy import integer as np_int_, floating as np_float_


def set_counts(count_dic):
    """
    Adds counts to meta-species using a given dictionary. Keys from this dictionary can be either meta-species
    objects or strings. Items must be the assign counts to species

    :param count_dic: Dictionary where keys can be either meta-species or strings
    :raise simlog.error:
        - If a count is assigned to a Reacting_Species with more than one meta-species.
        - If there are two species with the same name and a string assignment is performed.
        - If the keys are not strings or meta-species.
        - If the counts are not Quantities, Floats or Ints.
        - If the species was not found in the stack.
    :return: List_Species object. All meta-species that had a count assigned in this dictionary will be returned
        as a List_Species which can be passed as a model to the simulation object
    """
    new_count_dict = {}
    for key, item in count_dic.items():
        if (
            type(item) == int
            or type(item) == float
            or isinstance(item, Quantity)
            or isinstance(item, mp_Mobspy_Parameter)
        ):
            new_count_dict[key] = item
        elif isinstance(item, (np_int_, np_float_)):
            new_count_dict[key] = float(item)
        else:
            simlog_error(
                f"Reactant_species count assignment does not support the type {type(item)}",
            )
    count_dic = new_count_dict

    def find_species():
        found_species = set()

        for i in range(len(inspect_stack())):
            local_names = inspect_stack()[i][0].f_locals
            global_names = inspect_stack()[i][0].f_globals
            for key, item in global_names.items():
                try:
                    if isinstance(item, Species) and type(item) != type:
                        found_species.add(item)
                except AttributeError:
                    pass
            for key, item in local_names.items():
                try:
                    if isinstance(item, Species) and type(item) != type:
                        found_species.add(item)
                except AttributeError:
                    pass

        return found_species

    for key in count_dic:
        if type(key) == str:
            all_found_species = find_species()

    model = set()
    for key, item in count_dic.items():
        if type(key) == str:
            already_found = False
            str_name = key.split(".")[0]
            str_characteristics = set(key.split(".")[1:])
            for spe in all_found_species:
                if spe.get_name() == str_name and not already_found:
                    already_found = True
                    temp_set = set(str_characteristics)
                    temp_set.discard("all$")
                    if temp_set.issubset(spe.get_all_characteristics()):
                        spe.add_quantities(str_characteristics, item)
                    else:
                        simlog_error(
                            "Characteristics not found in species with equal name",
                        )
                    model.add(spe)
                elif spe.get_name() == str_name and already_found:
                    simlog_error(
                        "There are two different meta-species with the same name. Set_counts cannot resolve",
                    )
            if not already_found:
                simlog_error(
                    f"Meta-species with the following name {key} not found",
                )
        else:
            try:
                if isinstance(key, Species) or isinstance(key, Reacting_Species):
                    if not isinstance(key, Species):
                        if len(key.list_of_reactants) != 1:
                            simlog_error(
                                "Assignment used incorrectly. Only one species at a time",
                            )
                        model.add(key.list_of_reactants[0]["object"])
                    if isinstance(key, Species):
                        model.add(key)
                    key(item)
            except AttributeError:
                simlog_error(
                    "Keys must be either meta-species or strings"
                )

    return List_Species(model)
