# Path: flare/kernels/utils.py
# def from_mask_to_args(hyps, cutoffs, hyps_mask=None):
#     """Return the tuple of arguments needed for kernel function.
#     The order of the tuple has to be exactly the same as the one taken by
#         the kernel function.
# 
#     :param hyps: list of hyperparmeter values
#     :type hyps: nd.array
#     :param hyps_mask: all the remaining parameters needed
#     :type hyps_mask: dictionary
#     :param cutoffs: cutoffs used
# 
#     :return: args
#     """
# 
#     # no special setting
#     multihyps = True
#     if hyps_mask is None:
#         multihyps = False
#     elif hyps_mask["nspecie"] == 1:
#         multihyps = False
# 
#     if not multihyps:
# 
#         cutoffs_array = [0, 0, 0]
#         cutoffs_array[0] = cutoffs.get("twobody", 0)
#         cutoffs_array[1] = cutoffs.get("threebody", 0)
#         cutoffs_array[2] = cutoffs.get("manybody", 0)
#         return (hyps, cutoffs_array)
# 
#     # setting for mc_sephyps
#     nspecie = hyps_mask["nspecie"]
#     n2b = hyps_mask.get("ntwobody", 0)
# 
#     n3b = hyps_mask.get("nthreebody", 0)
#     nmanybody = hyps_mask.get("nmanybody", 0)
#     ncut3b = hyps_mask.get("ncut3b", 0)
# 
#     twobody_mask = hyps_mask.get("twobody_mask", None)
#     threebody_mask = hyps_mask.get("threebody_mask", None)
#     manybody_mask = hyps_mask.get("manybody_mask", None)
#     cut3b_mask = hyps_mask.get("cut3b_mask", None)
# 
#     # TO DO , should instead use the non-sephyps kernel
#     if n2b == 1:
#         twobody_mask = np.zeros(nspecie ** 2, dtype=int)
#     if n3b == 1:
#         threebody_mask = np.zeros(nspecie ** 3, dtype=int)
#     if nmanybody == 1:
#         manybody_mask = np.zeros(nspecie ** 2, dtype=int)
# 
#     cutoff_2b = cutoffs.get("twobody", 0)
#     cutoff_3b = cutoffs.get("threebody", 0)
#     cutoff_mb = cutoffs.get("manybody", 0)
# 
#     if "bond_cutoff_list" in hyps_mask:
#         cutoff_2b = hyps_mask["bond_cutoff_list"]
#     else:
#         cutoff_2b = np.ones(nspecie ** 2, dtype=float) * cutoff_2b
# 
#     if "threebody_cutoff_list" in hyps_mask:
#         cutoff_3b = hyps_mask["threebody_cutoff_list"]
#     if "manybody_cutoff_list" in hyps_mask:
#         cutoff_mb = hyps_mask["manybody_cutoff_list"]
# 
#     (sig2, ls2) = Parameters.get_component_hyps(hyps_mask, "twobody", hyps=hyps)
#     (sig3, ls3) = Parameters.get_component_hyps(hyps_mask, "threebody", hyps=hyps)
#     (sigm, lsm) = Parameters.get_component_hyps(hyps_mask, "manybody", hyps=hyps)
# 
#     return (
#         cutoff_2b,
#         cutoff_3b,
#         cutoff_mb,
#         nspecie,
#         np.array(hyps_mask["species_mask"]),
#         n2b,
#         twobody_mask,
#         n3b,
#         threebody_mask,
#         ncut3b,
#         cut3b_mask,
#         nmanybody,
#         manybody_mask,
#         sig2,
#         ls2,
#         sig3,
#         ls3,
#         sigm,
#         lsm,
#     )
# 
# def from_grad_to_mask(grad, hyps_mask=None):
#     """
#     Return gradient which only includes hyperparameters
#     which are meant to vary
# 
#     :param grad: original gradient vector
#     :param hyps_mask: dictionary for hyper-parameters
# 
#     :return: newgrad
#     """
# 
#     constrain = True
#     if hyps_mask is None:
#         constrain = False
#     elif "map" not in hyps_mask:
#         constrain = False
#     if not constrain:
#         return grad
# 
#     hyp_index = hyps_mask["map"]
# 
#     # setting for mc_sephyps
#     # if the last element is not sigma_noise
#     if hyp_index[-1] == len(grad):
#         hm = hyp_index[:-1]
#     else:
#         hm = hyp_index
# 
#     newgrad = np.zeros(len(hm), dtype=np.float64)
#     for i, mapid in enumerate(hm):
#         newgrad[i] = grad[mapid]
# 
#     return newgrad


import logging
import math
import multiprocessing as mp
import numpy as np
import time
from typing import List, Callable
from flare.kernels.utils import from_mask_to_args, from_grad_to_mask
def get_distance_mat_pack(
    hyps: np.ndarray,
    name: str,
    s1: int,
    e1: int,
    s2: int,
    e2: int,
    same: bool,
    kernel,
    cutoffs,
    hyps_mask,
):
    """Compute covariance matrix element between set1 and set2
    :param hyps: list of hyper-parameters
    :param name: name of the gp instance.
    :param same: whether the row and column are the same
    :param kernel: function object of the kernel
    :param cutoffs: The cutoff values used for the atomic environments
    :type cutoffs: list of 2 float numbers
    :param hyps_mask: dictionary used for multi-group hyperparmeters
    :return: covariance matrix
    """

    # initialize matrices
    training_data = _global_training_data[name]
    size1 = e1 - s1
    size2 = e2 - s2
    dist_mat = np.zeros([size1, size2])

    # calculate elements

# Path: flare/kernels/utils.py
# def from_mask_to_args(hyps, cutoffs, hyps_mask=None):
#     """Return the tuple of arguments needed for kernel function.
#     The order of the tuple has to be exactly the same as the one taken by
#         the kernel function.
# 
#     :param hyps: list of hyperparmeter values
#     :type hyps: nd.array
#     :param hyps_mask: all the remaining parameters needed
#     :type hyps_mask: dictionary
#     :param cutoffs: cutoffs used
# 
#     :return: args
#     """
# 
#     # no special setting
#     multihyps = True
#     if hyps_mask is None:
#         multihyps = False
#     elif hyps_mask["nspecie"] == 1:
#         multihyps = False
# 
#     if not multihyps:
# 
#         cutoffs_array = [0, 0, 0]
#         cutoffs_array[0] = cutoffs.get("twobody", 0)
#         cutoffs_array[1] = cutoffs.get("threebody", 0)
#         cutoffs_array[2] = cutoffs.get("manybody", 0)
#         return (hyps, cutoffs_array)
# 
#     # setting for mc_sephyps
#     nspecie = hyps_mask["nspecie"]
#     n2b = hyps_mask.get("ntwobody", 0)
# 
#     n3b = hyps_mask.get("nthreebody", 0)
#     nmanybody = hyps_mask.get("nmanybody", 0)
#     ncut3b = hyps_mask.get("ncut3b", 0)
# 
#     twobody_mask = hyps_mask.get("twobody_mask", None)
#     threebody_mask = hyps_mask.get("threebody_mask", None)
#     manybody_mask = hyps_mask.get("manybody_mask", None)
#     cut3b_mask = hyps_mask.get("cut3b_mask", None)
# 
#     # TO DO , should instead use the non-sephyps kernel
#     if n2b == 1:
#         twobody_mask = np.zeros(nspecie ** 2, dtype=int)
#     if n3b == 1:
#         threebody_mask = np.zeros(nspecie ** 3, dtype=int)
#     if nmanybody == 1:
#         manybody_mask = np.zeros(nspecie ** 2, dtype=int)
# 
#     cutoff_2b = cutoffs.get("twobody", 0)
#     cutoff_3b = cutoffs.get("threebody", 0)
#     cutoff_mb = cutoffs.get("manybody", 0)
# 
#     if "bond_cutoff_list" in hyps_mask:
#         cutoff_2b = hyps_mask["bond_cutoff_list"]
#     else:
#         cutoff_2b = np.ones(nspecie ** 2, dtype=float) * cutoff_2b
# 
#     if "threebody_cutoff_list" in hyps_mask:
#         cutoff_3b = hyps_mask["threebody_cutoff_list"]
#     if "manybody_cutoff_list" in hyps_mask:
#         cutoff_mb = hyps_mask["manybody_cutoff_list"]
# 
#     (sig2, ls2) = Parameters.get_component_hyps(hyps_mask, "twobody", hyps=hyps)
#     (sig3, ls3) = Parameters.get_component_hyps(hyps_mask, "threebody", hyps=hyps)
#     (sigm, lsm) = Parameters.get_component_hyps(hyps_mask, "manybody", hyps=hyps)
# 
#     return (
#         cutoff_2b,
#         cutoff_3b,
#         cutoff_mb,
#         nspecie,
#         np.array(hyps_mask["species_mask"]),
#         n2b,
#         twobody_mask,
#         n3b,
#         threebody_mask,
#         ncut3b,
#         cut3b_mask,
#         nmanybody,
#         manybody_mask,
#         sig2,
#         ls2,
#         sig3,
#         ls3,
#         sigm,
#         lsm,
#     )
# 
# def from_grad_to_mask(grad, hyps_mask=None):
#     """
#     Return gradient which only includes hyperparameters
#     which are meant to vary
# 
#     :param grad: original gradient vector
#     :param hyps_mask: dictionary for hyper-parameters
# 
#     :return: newgrad
#     """
# 
#     constrain = True
#     if hyps_mask is None:
#         constrain = False
#     elif "map" not in hyps_mask:
#         constrain = False
#     if not constrain:
#         return grad
# 
#     hyp_index = hyps_mask["map"]
# 
#     # setting for mc_sephyps
#     # if the last element is not sigma_noise
#     if hyp_index[-1] == len(grad):
#         hm = hyp_index[:-1]
#     else:
#         hm = hyp_index
# 
#     newgrad = np.zeros(len(hm), dtype=np.float64)
#     for i, mapid in enumerate(hm):
#         newgrad[i] = grad[mapid]
# 
#     return newgrad


# Path: flare/gp_algebra.py
import logging
import math
import multiprocessing as mp
import numpy as np
import time
from typing import List, Callable
from flare.kernels.utils import from_mask_to_args, from_grad_to_mask

def get_distance_mat_pack(
    hyps: np.ndarray,
    name: str,
    s1: int,
    e1: int,
    s2: int,
    e2: int,
    same: bool,
    kernel,
    cutoffs,
    hyps_mask,
):
    """Compute covariance matrix element between set1 and set2
    :param hyps: list of hyper-parameters
    :param name: name of the gp instance.
    :param same: whether the row and column are the same
    :param kernel: function object of the kernel
    :param cutoffs: The cutoff values used for the atomic environments
    :type cutoffs: list of 2 float numbers
    :param hyps_mask: dictionary used for multi-group hyperparmeters
    :return: covariance matrix
    """

    # initialize matrices
    training_data = _global_training_data[name]
    size1 = e1 - s1
    size2 = e2 - s2
    dist_mat = np.zeros([size1, size2])

    # calculate elements

    args = from_mask_to_args(hyps, cutoffs, hyps_mask)