Utils

def compute_dot_product(model: dict, params: dict, device: Optional[str] = None) -> torch.tensor:
    """Compute the dot product between model and input parameters.

    Args:
        model: OrderedDict representing model state
        params: OrderedDict containing correction parameters

    Returns:
        A tensor containing a single numerical value which is the dot product.
    """
    model_p = model.values()
    correction_state = params.values()
    if device is None:
        if model_p:
            device = list(model_p)[0].device
        else:
            # if device is not found, set it to `cpu`
            device = 'cpu'
    dot_prod = sum([torch.sum(m * torch.tensor(p).float().to(device)) for m, p in zip(model_p, correction_state)])
    return dot_prod

def convert_iterator_to_list_of_python_floats(iterator: Iterator) -> List[float]:
    """Converts numerical values of array-like object to float

    Args:
        iterator: Array-like numeric object to convert numerics to float

    Returns:
        Numerical elements as converted to List of floats
    """

    if not isinstance(iterator, Iterable):
        raise FedbiomedError(f"{ErrorNumbers.FB627.value}: object {type(iterator)} is not iterable")

    list_of_floats = []
    if isinstance(iterator, dict):
        # specific processing for dictionaries
        for val in iterator.values():
            list_of_floats.append(convert_to_python_float(val))
    else:
        for it in iterator:
            list_of_floats.append(convert_to_python_float(it))
    return list_of_floats

def convert_to_python_float(value: Union[torch.Tensor, np.integer, np.floating, float, int]) -> float:
    """ Convert numeric types to float

    Args:
        value: value to convert python type float

    Returns:
        Python float
    """

    if not isinstance(value, (torch.Tensor, np.integer, np.floating, float, int)):
        raise FedbiomedError(
            f"{ErrorNumbers.FB627.value}: Converting {type(value)} to python to float is not supported.")

    # if the result is a tensor, convert it back to numpy
    if isinstance(value, torch.Tensor):
        value = value.numpy()

    if isinstance(value, Iterable) and value.size > 1:
        raise FedbiomedError(f"{ErrorNumbers.FB627.value}: Can not convert array-type objects to float.")

    return float(value)

def create_fedbiomed_setup_folders(root: str):
    """Creates folders reequired by Fed-BioMed component setup

    Args:
        root: Root directory of Fed-BioMed component setup
    """

    etc_config_dir = os.path.join(root, CONFIG_FOLDER_NAME)
    var_dir = os.path.join(root, VAR_FOLDER_NAME)
    cache_dir = os.path.join(var_dir, CACHE_FOLDER_NAME)
    tmp_dir  = os.path.join(var_dir, TMP_FOLDER_NAME)

    for dir_ in [etc_config_dir, var_dir, cache_dir, tmp_dir]:
        if not os.path.isdir(dir_):
            os.makedirs(dir_)

def divide(xs: List[int], k: int) -> List[int]:
    """
    Divides a list of integers by a constant

    Args:
        xs: List of integers
        k: Constant to divide by

    Returns:
        List of divided integers
    """
    # Quicker than converting to/from numpy
    return [e / k for e in xs]

def get_all_existing_certificates() -> List[Dict[str, str]]:
    """Gets all existing certificates from Fed-BioMed `etc` directory.

    This method parse all available configs in `etc` directory.

    Returns:
        List of certificate objects that contain  component type as `component`,
            party id `id`, public key content (not path)  as `certificate`.
    """

    config_files = get_all_existing_config_files()

    certificates = []
    for config in config_files:
        certificates.append(get_component_certificate_from_config(config))

    return certificates

def get_all_existing_config_files():
    """Gets all existing config files from Fed-BioMed `etc` directory"""
    etc = os.path.join(ROOT_DIR, CONFIG_FOLDER_NAME, '')
    return [file for file in glob.glob(f"{etc}*.ini")]

def get_class_source(cls: Callable) -> str:
    """Get source of the class.

    It uses different methods for getting the class source based on shell type; IPython,Notebook
    shells or Python shell

    Args:
        cls: The class to extract the source code from

    Returns:
        str: Source code of the given class

    Raises:
        FedbiomedError: if argument is not a class
    """

    if not inspect.isclass(cls):
        raise FedbiomedError(f'{ErrorNumbers.FB627.value}: The argument `cls` must be a python class')

    # Check ipython status
    status = is_ipython()

    if status:
        file = get_ipython_class_file(cls)
        codes = "".join(inspect.linecache.getlines(file))
        class_code = extract_symbols(codes, cls.__name__)[0][0]
        return class_code

    return inspect.getsource(cls)

def get_component_certificate_from_config(
        config_path: str
) -> Dict[str, str]:
    """Gets component certificate, id and component type by given config file path.

    Args:
        config_path: Path where config file is located.

    Returns:
        Certificate object that contains  component type as `component`, party id `id`, public key content
            (not path)  as `certificate`

    Raises:
        FedbiomedError:
            - If config file does not contain `node_id` or `researcher_id` under `default` section.
            - If config file does not contain `public_key` under `ssl` section.
            - If certificate file is not found or not readable
    """

    config = get_component_config(config_path)
    component_id = config.get("default", "id")
    component_type = config.get("default", "component")

    certificate = config.get("certificate", "public_key")
    certificate_path = os.path.join(
        os.path.dirname(config_path), certificate)

    if not os.path.isfile(certificate_path):
        raise FedbiomedError(
            f"The certificate for component '{component_id}' not found in {certificate_path}"
        )

    certificate = read_file(certificate_path)

    return {
        "party_id": component_id,
        "certificate": certificate,
        "component": component_type
    }

def get_component_config(
        config_path: str
) -> configparser.ConfigParser:
    """Gets config object from given config path.

    Args:
        config_path: The path where config file is stored.

    Returns:
        Configuration object.

    Raises:
        FedbiomedError: If config file is not readable or not existing.
    """
    config = configparser.ConfigParser()

    try:
        config.read(config_path)
    except Exception:
        raise FedbiomedError(f"Can not read config file. Please make sure it is existing or it has valid format. "
                             f"{config_path}")

    return config

def get_existing_component_db_names():
    """Gets DB_PATHs of all existing components in Fed-BioMed root"""

    config_files = get_all_existing_config_files()
    db_names = {}

    for _config in config_files:
        config = get_component_config(_config)
        component_id = config['default']['id']

        db_name = f"{DB_PREFIX}{component_id}"
        db_names = {**db_names, component_id: db_name}

    return db_names

def get_ipython_class_file(cls: Callable) -> str:
    """Get source file/cell-id of the class which is defined in ZMQInteractiveShell or TerminalInteractiveShell

    Args:
        cls: Python class defined on the IPython kernel

    Returns:
        File path or id of Jupyter cell. On IPython's interactive shell, it returns cell ID
    """

    # Lookup by parent module
    if hasattr(cls, '__module__'):
        object_ = sys.modules.get(cls.__module__)
        # If module has `__file__` attribute
        if hasattr(object_, '__file__'):
            return object_.__file__

        # If parent module is __main__
        for name, member in inspect.getmembers(cls):
            if inspect.isfunction(member) and cls.__qualname__ + '.' + member.__name__ == member.__qualname__:
                return inspect.getfile(member)
    else:
        raise FedbiomedError(f'{ErrorNumbers.FB627.value}: {cls} has no attribute `__module__`, source is not found.')

def get_method_spec(method: Callable) -> dict:
    """ Helper to get argument specification

    Args:
        method: The function/method to extract argument specification from

    Returns:
         Specification of the method
    """

    method_spec = {}
    parameters = inspect.signature(method).parameters
    for (key, val) in parameters.items():
        method_spec[key] = {
            'name': val.name,
            'default': None if val.default is inspect._empty else val.default,
            'annotation': None if val.default is inspect._empty else val.default
        }

    return method_spec

def import_class_from_file(module_path: str, class_name: str) -> Tuple[Any, Any]:
    """Import a module from a file and return a specified class of the module.

    Args:
        module_path: path to python module file
        class_name: name of the class

    Returns:
        Tuple of the module created and the training plan class loaded

    Raises:
        FedbiomedError: bad argument type
        FedbiomedError: cannot load module or class
    """
    if not os.path.isfile(module_path):
        raise FedbiomedError(f"{ErrorNumbers.FB627}: Given path for importing {class_name} is not existing")

    module_base_name = os.path.basename(module_path)
    pattern = re.compile("(.*).py$")

    match = pattern.match(module_base_name)
    if not match:
        raise FedbiomedError(f"{ErrorNumbers.FB627}: File is not a python file.")

    module = match.group(1)
    sys.path.insert(0, os.path.dirname(module_path))


    class_ = import_object(module, class_name)
    sys.path.pop(0)

    return module, class_

def import_class_from_spec(code: str, class_name: str) -> Tuple[Any, Any] :
    """Import a module from a code and extract the code of a specified class of the module.

    Args:
        code: code of the module
        class_name: name of the class

    Returns:
        Tuple of the module created and the extracted class

    Raises:
        FedbiomedError: bad argument type
        FedbiomedError: cannot load module or extract clas
    """

    for arg in [code, class_name]:
        if not isinstance(arg, str):
            raise FedbiomedError(f"{ErrorNumbers.FB627.value}: Expected argument type is string but got '{type(arg)}'")

    try:
        spec = importlib.util.spec_from_loader("module_", loader=None)
        module = importlib.util.module_from_spec(spec)
        exec(code, module.__dict__)
    except Exception as e:
        raise FedbiomedError(f"{ErrorNumbers.FB627.value}: Can not load module from given code: {e}")

    try:
        class_ = getattr(module, class_name)
    except AttributeError:
        raise FedbiomedError(f"{ErrorNumbers.FB627.value}: Can not import {class_name} from given code")

    return module, class_

def import_class_object_from_file(module_path: str, class_name: str) -> Tuple[Any, Any]:
    """Import a module from a file and create an instance of a specified class of the module.

    Args:
        module_path: path to python module file
        class_name: name of the class

    Returns:
        Tuple of the module created and the training plan object created

    Raises:
        FedbiomedError: bad argument type
        FedbiomedError: cannot instantiate object
    """
    for arg in [module_path, class_name]:
        if not isinstance(arg, str):
            raise FedbiomedError(f"{ErrorNumbers.FB627.value}: Expected argument type is string but got '{type(arg)}'")

    module, train_class = import_class_from_file(module_path, class_name)

    try:
        train_class_instance = train_class()
    except Exception as e:
        raise FedbiomedError(f"{ErrorNumbers.FB627.value}: Cannot instantiate training plan object: {e}")

    return module, train_class_instance

def import_object(module: str, obj_name: str) -> Any:
    """Imports given object/class from given module


    Args:
        module: Module that the object will be imported from
        obj_name: Name of the object or class defined in the module

    Returns:
        Python object
    """
    try:
        module = importlib.import_module(module)
    except ModuleNotFoundError as exp:
        raise FedbiomedError(f"Specified module is not existing. {exp}") from exp

    try:
        obj_ = getattr(module, obj_name)
    except AttributeError as exp:
        raise FedbiomedError(f"{ErrorNumbers.FB627}, Attribute error while loading the class "
                             f"{obj_name} from {module}. Error: {exp}") from exp


    return obj_

def matching_parties_dh(context: dict, parties: list) -> bool:
    """Check if parties of given context are compatible with the parties
        of a secagg Diffie Hellman element.

        Args:
            context: context to be compared with the secagg servkey element parties
            parties: the secagg servkey element parties

        Returns:
            True if this context can be used with this element, False if not.
        """
    # Need to ensure that:
    # - no check on first party (no cryptographic material attached to the researcher).
    #   The context is established for a given experiment, thus a given researcher but this should
    #   be tested prior to this call.
    # - existing element was established for the same node parties or a superset of the node parties
    #   (order can differ, as nodes are ordered by the cipher code)
    #
    # eg: [ 'un', 'deux', 'trois' ] parties compatible with [ 'un', 'trois', 'deux' ] context
    # but not with [ 'deux', 'un', 'trois' ]
    # eg: [ 'un', 'deux', 'trois' ] parties compatible with [ 'un', 'trois', 'quatre', 'deux' ] context
    # but not with [ 'un', 'deux', 'quatre' ]
    return (
        # Commented tests can be assumed from calling functions
        #
        # isinstance(context, dict) and
        # 'parties' in context and
        # isinstance(context['parties'], list) and
        # len(context['parties']) >= 1 and
        # isinstance(parties, list) and
        set(parties[1:]).issubset(set(context['parties'][1:])))

def matching_parties_servkey(context: dict, parties: list) -> bool:
    """Check if parties of given context are compatible with the parties
        of a secagg servkey element.

        Args:
            context: context to be compared with the secagg servkey element parties
            parties: the secagg servkey element parties

        Returns:
            True if this context can be used with this element, False if not.
        """
    # Need to ensure that:
    # - existing element was established for the same parties
    # - first party needs to be the same for both
    # - set of other parties needs to be the same for both (order can differ)
    #
    # eg: [ 'un', 'deux', 'trois' ] compatible with [ 'un', 'trois', 'deux' ]
    # but not with [ 'deux', 'un', 'trois' ]
    return (
        # Commented tests can be assumed from calling functions
        #
        # isinstance(context, dict) and
        # 'parties' in context and
        # isinstance(context['parties'], list) and
        # len(context['parties']) >= 1 and
        # isinstance(parties, list) and
        parties[0] == context['parties'][0] and
        set(parties[1:]) == set(context['parties'][1:]))

def multiply(xs: List[int], k: int) -> List[int]:
    """
    Multiplies a list of integers by a constant

    Args:
        xs: List of integers
        k: Constant to multiply by

    Returns:
        List of multiplied integers
    """
    # Quicker than converting to/from numpy
    return [e * k for e in xs]

def quantize(
    weights: List[float],
    clipping_range: Union[int, None] = None,
    target_range: int = SAParameters.TARGET_RANGE,
) -> List[int]:
    """Quantization step implemented by: https://dl.acm.org/doi/pdf/10.1145/3488659.3493776

    This function returns a vector in the range [0, target_range-1].

    Args:
        weights: List of model weight values
        clipping_range: Clipping range
        target_range: Target range

    Returns:
        Quantized model weights as numpy array.

    """
    if clipping_range is None:
        clipping_range = SAParameters.CLIPPING_RANGE

    _check_clipping_range(weights, clipping_range)

    # CAVEAT: ensure to be converting from `float` to `uint64` (no intermediate `int64`)
    # Process ensures to compute an `int`` in the range [0, target_range -1]
    # This enables to use at most 2**64 as target_range (max value of `uint` - 1)
    f = np.vectorize(
        lambda x: min(
            target_range - 1,
            (sorted((-clipping_range, x, clipping_range))[1] +
            clipping_range) *
            target_range / (2 * clipping_range),
        ),
        otypes=[np.uint64]
    )
    quantized_list = f(weights)

    return quantized_list.tolist()

def raise_for_version_compatibility(their_version: Union[FBM_Component_Version, str],
                                    our_version: Union[FBM_Component_Version, str],
                                    error_msg: Optional[str] = None) -> None:
    """Check version compatibility and behave accordingly.

    Raises an exception if the versions are incompatible, otherwise outputs a warning or info message.

    Args:
        their_version: the version that we detected in the component
        our_version: the version of the component within the current runtime
        error_msg: an optional error message. It may contain two %s placeholders which will be substituted with
            the values of their_version and our_version.

    Raises:
        FedbiomedVersionError: if the versions are incompatible
    """
    if isinstance(our_version, str):
        our_version = FBM_Component_Version(our_version)
    if isinstance(their_version, str):
        their_version = FBM_Component_Version(their_version)
    if not isinstance(our_version, FBM_Component_Version):
        msg = f"{ErrorNumbers.FB625.value}: Component version has incorrect type `our_version` type={str(type(our_version))} value={our_version}"
        logger.critical(msg)
        raise FedbiomedVersionError(msg)
    if not isinstance(their_version, FBM_Component_Version):
        msg = f"{ErrorNumbers.FB625.value}: Component version has incorrect type `their_version` type={str(type(their_version))} value={their_version}"
        logger.critical(msg)
        raise FedbiomedVersionError(msg)
    if our_version != their_version:
        # note: the checks below rely on the short-circuiting behaviour of the or operator
        # (e.g. when checking our_version.minor < their_version.minor we have the guarantee that
        # our_version.major == their_version.major
        if our_version.major != their_version.major or \
                our_version.minor < their_version.minor or \
                (our_version.minor == their_version.minor and our_version.micro < their_version.micro):
            msg = _create_msg_for_version_check(
                f"{ErrorNumbers.FB625.value}: Found incompatible version %s, expected version %s" if error_msg is None else error_msg,
                their_version,
                our_version
            )
            logger.critical(msg)
            raise FedbiomedVersionError(msg)
        else:
            msg = _create_msg_for_version_check(
                "Found version %s, expected version %s",
                their_version,
                our_version
            )
            logger.warning(msg)

def read_file(path):
    """Read given file

    Args:
        path: Path to file to be read

    Raises:
        FedbiomedError: If the file is not existing or readable.
    """
    try:
        with open(path, "r", encoding="UTF-8") as file:
            content = file.read()
            file.close()
    except Exception as e:
        raise FedbiomedError(
            f"{ErrorNumbers.FB627.value}: Can not read file {path}. Error: {e}"
        ) from e

    return content

def reverse_quantize(
    weights: List[int],
    clipping_range: Union[int, None] = None,
    target_range: int = SAParameters.TARGET_RANGE,
) -> List[float]:
    """Reverse quantization step implemented by: https://dl.acm.org/doi/pdf/10.1145/3488659.3493776

     Args:
        weights: List of quantized model weights
        clipping_range: Clipping range used for quantization
        target_range: Target range used for quantization

    Returns:
        Reversed quantized model weights as numpy array.
    """

    if clipping_range is None:
        clipping_range = SAParameters.CLIPPING_RANGE

    # CAVEAT: there should not be any weight received that does not fit in `uint64`
    max_val = np.iinfo(np.uint64).max
    if any([v > max_val or v < 0 for v in weights]):
        raise FedbiomedSecaggCrypterError(
            f"{ErrorNumbers.FB624.value}: Cannot reverse quantize, received values exceed maximum number"
        )

    max_range = clipping_range
    min_range = -clipping_range
    step_size = (max_range - min_range) / (target_range - 1)
    # Compute as input type (`np.uint64` then convert to `np.float64`)
    f = np.vectorize(
        lambda x: (min_range + step_size * x),
        otypes=[np.float64]
    )

    # TODO: we could check that received values are in the range
    weights = np.array(weights, dtype=np.uint64)
    reverse_quantized_list = f(weights)

    return reverse_quantized_list.tolist()