try:
    import cupy
except:
    pass
import numpy as np
import numba as nb
from .distutils import *

# Base functions: Distance and energy computation --------------------------


def PartitionData_cp(
    Xcp,
    NodePositions,
    MaxBlockSize,
    SquaredXcp,
    TrimmingRadius=float("inf"),
    precomp=False,
):
    """
    # Partition the data by proximity to graph nodes
    # (same step as in K-means EM procedure)
    #
    # Inputs:
    #   X is n-by-m matrix of datapoints with one data point per row. n is
    #       number of data points and m is dimension of data space.
    #   NodePositions is k-by-m matrix of embedded coordinates of graph nodes,
    #       where k is number of nodes and m is dimension of data space.
    #   MaxBlockSize integer number which defines maximal number of
    #       simultaneously calculated distances. Maximal size of created matrix
    #       is MaxBlockSize-by-k, where k is number of nodes.
    #   SquaredX is n-by-1 vector of data vectors length: SquaredX = sum(X.^2,2);
    #   TrimmingRadius (optional) is squared trimming radius.
    #
    # Outputs
    #   partition is n-by-1 vector. partition[i] is number of the node which is
    #       associated with data point X[i, ].
    #   dists is n-by-1 vector. dists[i] is squared distance between the node with
    #       number partition[i] and data point X[i, ].
    """
    NodePositionscp = cupy.asarray(NodePositions)
    n = Xcp.shape[0]
    partition = cupy.zeros((n, 1), dtype=int)
    dists = cupy.zeros((n, 1))
    d = cupy.zeros((n, len(NodePositions)))

    # Calculate squared length of centroids
    cent = NodePositionscp.T
    centrLength = (cent ** 2).sum(axis=0)
    # Process partitioning without trimming

    for i in range(0, n, MaxBlockSize):
        # Define last element for calculation
        last = i + MaxBlockSize
        if last > n:
            last = n
        # Calculate distances
        _d = (
            SquaredXcp[i:last]
            + centrLength
            - 2
            * cupy.dot(
                Xcp[
                    i:last,
                ],
                cent,
            )
        )
        tmp = _d.argmin(axis=1)
        partition[i:last] = tmp[:, cupy.newaxis]
        dists[i:last] = _d[cupy.arange(_d.shape[0]), tmp][:, cupy.newaxis]
        d[i:last] = _d

    # Apply trimming
    if not cupy.isinf(TrimmingRadius):
        ind = dists > (TrimmingRadius ** 2)
        partition[ind] = -1
        dists[ind] = TrimmingRadius ** 2
    if precomp:
        return cupy.asnumpy(partition), cupy.asnumpy(dists), cupy.asnumpy(d)
    else:
        return cupy.asnumpy(partition), cupy.asnumpy(dists)


def PartitionData(
    X, NodePositions, MaxBlockSize, SquaredX, TrimmingRadius=float("inf"), precomp=False
):
    """
    # Partition the data by proximity to graph nodes
    # (same step as in K-means EM procedure)
    #
    # Inputs:
    #   X is n-by-m matrix of datapoints with one data point per row. n is
    #       number of data points and m is dimension of data space.
    #   NodePositions is k-by-m matrix of embedded coordinates of graph nodes,
    #       where k is number of nodes and m is dimension of data space.
    #   MaxBlockSize integer number which defines maximal number of
    #       simultaneously calculated distances. Maximal size of created matrix
    #       is MaxBlockSize-by-k, where k is number of nodes.
    #   SquaredX is n-by-1 vector of data vectors length: SquaredX = sum(X.^2,2);
    #   TrimmingRadius (optional) is squared trimming radius.
    #
    # Outputs
    #   partition is n-by-1 vector. partition[i] is number of the node which is
    #       associated with data point X[i, ].
    #   dists is n-by-1 vector. dists[i] is squared distance between the node with
    #       number partition[i] and data point X[i, ].
    """
    n = X.shape[0]
    partition = np.zeros((n, 1), dtype=int)
    dists = np.zeros((n, 1))
    d = np.zeros((n, len(NodePositions)))
    # Calculate squared length of centroids
    cent = NodePositions.T
    centrLength = (cent ** 2).sum(axis=0)
    # Process partitioning without trimming
    for i in range(0, n, MaxBlockSize):
        # Define last element for calculation
        last = i + MaxBlockSize
        if last > n:
            last = n
        # Calculate distances
        _d = (
            SquaredX[i:last]
            + centrLength
            - 2
            * np.dot(
                X[
                    i:last,
                ],
                cent,
            )
        )
        tmp = _d.argmin(axis=1)
        partition[i:last] = tmp[:, np.newaxis]
        dists[i:last] = _d[np.arange(_d.shape[0]), tmp][:, np.newaxis]
        d[i:last] = _d

    # Apply trimming
    if not np.isinf(TrimmingRadius):
        ind = dists > (TrimmingRadius ** 2)
        partition[ind] = -1
        dists[ind] = TrimmingRadius ** 2
    if precomp:
        return partition, dists, d
    else:
        return partition, dists


def RePartitionData(
    X,
    NodePositions,
    NewNodePositions,
    NodeIndicesArray,
    opType,
    precomp_d,
    SquaredX,
    TrimmingRadius=np.inf,
):

    if "remove" in opType:
        d = precomp_d[:, NodeIndicesArray]
    elif "shrink" in opType:
        repositioned_node_idx = np.all(
            NodePositions[NodeIndicesArray] != NewNodePositions, axis=1
        )
        newcent = NewNodePositions[repositioned_node_idx].T
        newcentrLength = (newcent ** 2).sum(axis=0)
        new_d = SquaredX + newcentrLength - 2 * np.dot(X, newcent)
        d = precomp_d[:, NodeIndicesArray]
        d[:, repositioned_node_idx] = new_d

    else:  # grow grammar:
        newcent = NewNodePositions[[-1]].T
        newcentrLength = (newcent ** 2).sum(axis=0)
        new_d = SquaredX + newcentrLength - 2 * np.dot(X, newcent)
        d = np.hstack((precomp_d, new_d))

    tmp = d.argmin(axis=1)
    partition = tmp[:, np.newaxis]
    dists = d[np.arange(d.shape[0]), tmp][:, np.newaxis]
    # Apply trimming
    if not np.isinf(TrimmingRadius):
        ind = dists > (TrimmingRadius ** 2)
        partition[ind] = -1
        dists[ind] = TrimmingRadius ** 2

    return partition, dists, d


def PartitionDataBarycenter(
    X,
    NodePositions,
    MaxBlockSize,
    SquaredX,
    TrimmingRadius=float("inf"),
    previous_partition=None,
):
    """
    # Partition the data by barycentric proximity to graph nodes
    # https://arxiv.org/pdf/1902.10288.pdf
    # (same step as in barycentric K-means procedure)
    #
    # Inputs:
    #   X is n-by-m matrix of datapoints with one data point per row. n is
    #       number of data points and m is dimension of data space.
    #   NodePositions is k-by-m matrix of embedded coordinates of graph nodes,
    #       where k is number of nodes and m is dimension of data space.
    #   MaxBlockSize integer number which defines maximal number of
    #       simultaneously calculated distances. Maximal size of created matrix
    #       is MaxBlockSize-by-k, where k is number of nodes.
    #   SquaredX is n-by-1 vector of data vectors length: SquaredX = sum(X.^2,2);
    #   TrimmingRadius (optional) is squared trimming radius.
    #
    # Outputs
    #   partition is n-by-1 vector. partition[i] is number of the node which is
    #       associated with data point X[i, ].
    #   dists is n-by-1 vector. dists[i] is squared distance between the node with
    #       number partition[i] and data point X[i, ].
    """
    # ------clusters standard deviations
    nodes = np.arange(len(NodePositions))
    sds = np.zeros(len(nodes))

    if previous_partition is not None:
        previous_partition = previous_partition.ravel()
        for i, node in enumerate(nodes):
            node_data = X[previous_partition == node]
            if len(node_data) < 2:
                # adding epsilon for clusters with single point or no point
                sds[i] = 1e-100
            else:
                sds[i] = np.std(node_data)

    n = X.shape[0]
    partition = np.zeros((n, 1), dtype=int)
    dists = np.zeros((n, 1))
    # Calculate squared length of centroids
    cent = NodePositions.T
    centrLength = (cent ** 2).sum(axis=0)
    # Process partitioning without trimming
    for i in range(0, n, MaxBlockSize):
        # Define last element for calculation
        last = i + MaxBlockSize
        if last > n:
            last = n
        # Calculate distances
        d = (
            SquaredX[i:last]
            + centrLength
            - 2
            * np.dot(
                X[
                    i:last,
                ],
                cent,
            )
        )
        if previous_partition is not None:
            d = d / sds + sds
            tmp = d.argmin(axis=1)
        else:
            tmp = d.argmin(axis=1)
            for j, node in enumerate(nodes):
                node_data = X[tmp == node]
                if len(node_data) < 2:
                    # adding epsilon for clusters with single point or no point
                    sds[j] = 1e-100
                else:
                    sds[j] = np.std(node_data)
            d = d / sds + sds
            tmp = d.argmin(axis=1)

        partition[i:last] = tmp[:, np.newaxis]
        dists[i:last] = d[np.arange(d.shape[0]), tmp][:, np.newaxis]
    # Apply trimming
    if not np.isinf(TrimmingRadius):
        ind = dists > (TrimmingRadius ** 2)
        partition[ind] = -1
        dists[ind] = TrimmingRadius ** 2
    return partition, dists


def MakeUniformElasticMatrix(Edges, Lambda, Mu):
    """
    # Base function: Function to deal with elastic matrices --------------------------
    #' Create a uniform elastic matrix from a set of edges
    #'
    #' @param Edges an e-by-2 matrix containing the index of the edges connecting the nodes
    #' @param Lambda the lambda parameter. It can be a real value or a vector of lenght e
    #' @param Mu the mu parameter. It can be a real value or a vector with a length equal to the number of nodes
    #'
    #' @return the elastic matrix
    #'
    #' @export
    #'
    #' @examples
    """
    NumberOfNodes = Edges.max() + 1
    ElasticMatrix = np.zeros((NumberOfNodes, NumberOfNodes))
    for i in range(Edges.shape[0]):
        ElasticMatrix[Edges[i][0], Edges[i][1]] = Lambda
        ElasticMatrix[Edges[i][1], Edges[i][0]] = Lambda
    Connect = (ElasticMatrix > 0).sum(axis=0)
    ind = Connect > 1
    Mus = np.zeros((NumberOfNodes, 1))
    Mus[ind] = Mu
    ElasticMatrix = ElasticMatrix + np.diag(Mus.ravel())
    return ElasticMatrix


def Encode2ElasticMatrix(Edges, Lambdas, Mus):
    """
    #' Create an Elastic matrix from a set of edges
    #'
    #' @param Lambdas the lambda parameters. Either a single value (which will be used for all the edges),
    #' or a vector containing the values for each edge
    #' @param Mus the mu parameters. Either a single value (which will be used for all the nodes),
    #' or a vector containing the values for each node
    #' @param Edges an e-by-2 matrix containing the index of the edges connecting the nodes
    #'
    #' @return the elastic matrix
    #'
    #' @export
    #'
    #' @examples"""

    NumberOfNodes = np.max(Edges) + 1
    NumberOfEdges = Edges.shape[0]

    EM = np.zeros((NumberOfNodes, NumberOfNodes))

    if isinstance(Lambdas, int) or isinstance(Lambdas, float):
        Lambdas = np.array([Lambdas] * NumberOfEdges)

    if isinstance(Mus, int) or isinstance(Mus, float):
        Mus = np.array([Mus] * NumberOfNodes)

    for i in range(NumberOfEdges):
        EM[Edges[i, 0], Edges[i, 1]] = Lambdas[i]
        EM[Edges[i, 1], Edges[i, 0]] = Lambdas[i]

    return EM + np.diag(Mus)


# def ComputeSpringLaplacianMatrix(ElasticMatrix):
#     '''
#     #' Compute the Laplacian matrix
#     #'
#     #' @param ElasticMatrix an e-by-e elastic matrix
#     #'
#     #' @return the Laplacian matrix
#     #'
#     #' @export
#     #'
#     #' @examples'''
#     NumberOfNodes = ElasticMatrix.shape[0]
#     # first, make the vector of mu coefficients
#     Mu = ElasticMatrix.diagonal()
#     # create the matrix with edge elasticity moduli at non-diagonal elements
#     Lambda = ElasticMatrix - np.diag(Mu)
#     # Diagonal matrix of edge elasticities
#     LambdaSums = Lambda.sum(axis=0)
#     # E matrix (contribution from edges) is simply weighted Laplacian
#     E = np.diag(LambdaSums) - Lambda
#     # matrix S (contribution from stars) is composed of Laplacian for
#     # positive strings (star edges) with elasticities mu/k, where k is the
#     # order of the star, and Laplacian for negative strings with
#     # elasticities -mu/k^2. Negative springs connect all star leafs in a
#     # clique.
#     StarCenterIndices = np.nonzero(Mu > 0)[0]
#     S = np.zeros((NumberOfNodes, NumberOfNodes))
#     for i in range(np.size(StarCenterIndices)):
#         Spart = np.zeros((NumberOfNodes, NumberOfNodes))
#         # leaves indices
#         leaves = Lambda.take(StarCenterIndices[i], axis=1) > 0
#         # order of the star
#         K = leaves.sum()
#         Spart[StarCenterIndices[i], StarCenterIndices[i]] = (
#                 Mu[StarCenterIndices[i]])
#         Spart[StarCenterIndices[i], leaves] = -Mu[StarCenterIndices[i]]/K
#         Spart[leaves, StarCenterIndices[i]] = -Mu[StarCenterIndices[i]]/K
#         tmp = np.repeat(leaves[np.newaxis],
#                         ElasticMatrix.shape[0], axis=0)
#         tmp = np.logical_and(tmp, tmp.transpose())
#         Spart[tmp] = Mu[StarCenterIndices[i]]/(K**2)
#         S = S + Spart
#     return E + S


@nb.njit(cache=True)
def ComputeSpringLaplacianMatrix(ElasticMatrix):
    NumberOfNodes = ElasticMatrix.shape[0]
    # first, make the vector of mu coefficients
    Mu = np.diag(ElasticMatrix)
    # create the matrix with edge elasticity moduli at non-diagonal elements
    Lambda = ElasticMatrix - np.diag(Mu)
    # set mu of endpoints to 0.
    # endpoints = np.where((np.bincount(np.concatenate(Lambda.nonzero())) // 2) == 1)[0]
    # Mu[endpoints] = 0.0
    # Diagonal matrix of edge elasticities
    LambdaSums = Lambda.sum(axis=0)
    # E matrix (contribution from edges) is simply weighted Laplacian
    E = np.diag(LambdaSums) - Lambda
    # matrix S (contribution from stars) is composed of Laplacian for
    # positive strings (star edges) with elasticities mu/k, where k is the
    # order of the star, and Laplacian for negative strings with
    # elasticities -mu/k^2. Negative springs connect all star leafs in a
    # clique.
    StarCenterIndices = (Mu > 0).nonzero()[0]
    S = np.zeros((NumberOfNodes, NumberOfNodes))
    for i in range(StarCenterIndices.size):
        Spart = np.zeros((NumberOfNodes, NumberOfNodes))
        # leaves indices
        leaves = Lambda[StarCenterIndices[i]] > 0
        ind_leaves = (leaves).nonzero()[0]
        # order of the star
        K = leaves.sum()
        Spart[StarCenterIndices[i], StarCenterIndices[i]] = Mu[StarCenterIndices[i]]
        for j in ind_leaves:
            Spart[StarCenterIndices[i], j] = -Mu[StarCenterIndices[i]] / K
            Spart[j, StarCenterIndices[i]] = -Mu[StarCenterIndices[i]] / K
        tmp = np.repeat(leaves, ElasticMatrix.shape[0]).reshape(
            (-1, ElasticMatrix.shape[0])
        )
        tmp_r, tmp_c = np.where(np.logical_and(tmp, tmp.transpose()))
        K2 = K ** 2
        for it in range(len(tmp_r)):
            Spart[tmp_r[it], tmp_c[it]] = Mu[StarCenterIndices[i]] / K2
        S = S + Spart
    return E + S


def DecodeElasticMatrix(ElasticMatrix):
    """
    #' Converts ElasticMatrix into a set of edges and vectors of elasticities for edges and stars
    #'
    #' @param ElasticMatrix an e-by-e elastic matrix
    #'
    #' @return a list with three elements: a matrix with the edges (Edges), a vector of lambdas (Lambdas), and a vector of Mus (Mus)
    #'
    #' @export
    #'
    #' @examples
    """
    Mus = ElasticMatrix.diagonal()
    Lambda = ElasticMatrix.copy()
    np.fill_diagonal(Lambda, 0)

    Edges_i, Edges_j = np.where(Lambda > 0)
    Edges = np.concatenate((Edges_j[:, None], Edges_i[:, None]), axis=1)
    inds = np.where(Edges_i > Edges_j)

    Edges = Edges[inds]

    Lambdas = Lambda[Edges[:, 0], Edges[:, 1]]

    return Edges, Lambdas, Mus


def DecodeElasticMatrix2(ElasticMatrix):
    Mus = ElasticMatrix.diagonal()
    Lambda = ElasticMatrix.copy()
    np.fill_diagonal(Lambda, 0)

    Edges_i, Edges_j = np.where(Lambda > 0)
    Edges = np.concatenate((Edges_i[:, None], Edges_j[:, None]), axis=1)
    inds = np.where(Edges_i < Edges_j)

    Edges = Edges[inds]

    Lambdas = Lambda[Edges[:, 0], Edges[:, 1]]

    return Edges, Lambdas, Mus


@nb.njit(cache=True)
def ComputeRelativeChangeOfNodePositions(NodePositions, NewNodePositions):
    """
    #' Estimates the relative difference between two node configurations
    #'
    #' @param NodePositions a k-by-m numeric matrix with the coordiantes of the nodes in the old configuration
    #' @param NewNodePositions a k-by-m numeric matrix with the coordiantes of the nodes in the new configuration
    #' @param Mode an integer indicating the modality used to compute the difference (currently only 1 is an accepted value)
    #' @param X an n-by-m numeric matrix with the coordinates of the data points
    #'
    #' @return
    #' @export
    #'
    #' @examples
    """
    return (
        ((NodePositions - NewNodePositions) ** 2).sum(axis=1)
        / (NewNodePositions ** 2).sum(axis=1)
    ).max()


def PrimitiveElasticGraphBarycentricEmbedment(
    X,
    NodePositions,
    ElasticMatrix,
    MaxNumberOfIterations=10,
    eps=0.01,
    Mode=1,
    FinalEnergy="Base",
    alpha=0,
    beta=0,
    prob=1,
    DisplayWarnings=True,
    PointWeights=None,
    MaxBlockSize=100000000,
    verbose=False,
    TrimmingRadius=float("inf"),
    SquaredX=None,
    FixNodesAtPoints=[],
):

    """
    #' Function fitting a primitive elastic graph to the data
    #'
    #' @param X is n-by-m matrix containing the positions of the n points in the m-dimensional space
    #' @param NodePositions is k-by-m matrix of positions of the graph nodes in the same space as X
    #' @param ElasticMatrix is a k-by-k symmetric matrix describing the connectivity and the elastic
    #' properties of the graph. Star elasticities (mu coefficients) are along the main diagonal
    #' (non-zero entries only for star centers), and the edge elasticity moduli are at non-diagonal elements.
    #' @param MaxNumberOfIterations is an integer number indicating the maximum number of iterations for the EM algorithm
    #' @param TrimmingRadius is a real value indicating the trimming radius, a parameter required for robust principal graphs
    #' (see https://github.com/auranic/Elastic-principal-graphs/wiki/Robust-principal-graphs)
    #' @param eps a real number indicating the minimal relative change in the nodenpositions
    #' to be considered the graph embedded (convergence criteria)
    #' @param verbose is a boolean indicating whether diagnostig informations should be plotted
    #' @param Mode integer, the energy mode. It can be 1 (difference is computed using the position of the nodes) and
    #' 2 (difference is computed using the changes in elestic energy of the configuraztions)
    #' @param SquaredX the sum (by node) of X squared. It not specified, it will be calculated by the fucntion
    #' @param FastSolve boolean, shuold the Fastsolve of Armadillo by enabled?
    #' @param DisplayWarnings boolean, should warning about convergence be displayed?
    #' @param FinalEnergy string indicating the final elastic emergy associated with the configuration. Currently it can be "Base" or "Penalized"
    #' @param alpha positive numeric, the value of the alpha parameter of the penalized elastic energy
    #' @param beta positive numeric, the value of the beta parameter of the penalized elastic energy
    #' @param prob numeric between 0 and 1. If less than 1 point will be sampled at each iteration. Prob indicate the probability of
    #' using each points. This is an *experimental* feature, which may helps speeding up the computation if a large number of points is present.
    #'
    #' @return
    #' @export
    #'
    #' @examples
    """

    if prob < 1:
        raise ValueError("probPoint < 1 option not implemented yet")
    N = X.shape[0]

    # Auxiliary computations
    if PointWeights is None:
        PointWeights = np.ones((N, 1))
    if SquaredX is None:
        SquaredX = (X ** 2).sum(axis=1).reshape((N, 1))
    SpringLaplacianMatrix = ComputeSpringLaplacianMatrix(ElasticMatrix)

    if FixNodesAtPoints != []:
        # -----Index data and graph into moving nodes and fixed nodes------#
        flat_FixNodesAtPoints = [
            item for sublist in FixNodesAtPoints for item in sublist
        ]  # fixed datapoints

        # partition, dists = PartitionData(
        #    X, NodePositions, MaxBlockSize, SquaredX, TrimmingRadius,
        # )
        # move_data_idx = np.where(~np.isin(partition, range(len(FixNodesAtPoints))))[0]

        move_data_idx = [i for i in range(len(X)) if i not in flat_FixNodesAtPoints]
        move_nodes_idx = np.arange(len(FixNodesAtPoints), len(NodePositions))
    else:
        move_data_idx = np.arange(len(X))
        move_nodes_idx = np.arange(len(NodePositions))

    # fitted subset
    move_X, move_PointWeights, move_SquaredX, move_NodePositions = (
        X[move_data_idx, :],
        PointWeights[move_data_idx],
        SquaredX[move_data_idx],
        NodePositions[move_nodes_idx],
    )

    # Main iterative EM cycle: partition, fit given the partition, repeat
    move_partition, move_dists = PartitionDataBarycenter(
        move_X,
        move_NodePositions,
        MaxBlockSize,
        move_SquaredX,
        TrimmingRadius,
        previous_partition=None,
    )

    if verbose or Mode == 2:
        # if FixNodesAtPoints != []:
        #    dists = np.concatenate((FixedNodesDists, move_dists))
        OldElasticEnergy, MSE, EP, RP = ComputePrimitiveGraphElasticEnergy(
            NodePositions, ElasticMatrix, move_dists
        )

    ElasticEnergy = 0
    for i in range(MaxNumberOfIterations):
        # Updated positions
        if FixNodesAtPoints != []:
            NewNodePositions = FitSubGraph2DataGivenPartition(
                move_X,
                move_PointWeights,
                SpringLaplacianMatrix,
                NodePositions,
                move_partition,
                move_nodes_idx,
            )
        else:
            NewNodePositions = FitGraph2DataGivenPartition(
                move_X,
                move_PointWeights,
                SpringLaplacianMatrix,
                move_partition,
            )

        # Look at differences
        if verbose or Mode == 2:
            # if FixNodesAtPoints != []:
            #    dists = np.concatenate((FixedNodesDists, move_dists))
            ElasticEnergy, MSE, EP, RP = ComputePrimitiveGraphElasticEnergy(
                NewNodePositions, ElasticMatrix, move_dists
            )

        if Mode == 1:
            diff = ComputeRelativeChangeOfNodePositions(NodePositions, NewNodePositions)
        elif Mode == 2:
            diff = (OldElasticEnergy - ElasticEnergy) / ElasticEnergy

        # Print Info
        if verbose:
            print(
                "Iteration ",
                (i + 1),
                " difference of node position=",
                np.around(diff, 5),
                ", Energy=",
                np.around(ElasticEnergy, 5),
                ", MSE=",
                np.around(MSE, 5),
                ", EP=",
                np.around(EP, 5),
                ", RP=",
                np.around(RP, 5),
            )

        # Have we converged?
        if not np.isfinite(diff):
            diff = 0
        if diff < eps:
            break

        elif i < MaxNumberOfIterations - 1:
            move_NodePositions = NewNodePositions[move_nodes_idx]
            move_partition, move_dists = PartitionDataBarycenter(
                move_X,
                move_NodePositions,
                MaxBlockSize,
                move_SquaredX,
                TrimmingRadius,
                previous_partition=move_partition,
            )

            NodePositions = NewNodePositions
            OldElasticEnergy = ElasticEnergy

    if DisplayWarnings and not (diff < eps):
        print(
            "Maximum number of iterations (",
            MaxNumberOfIterations,
            ") has been reached. diff = ",
            diff,
        )

    # If we
    # 1) Didn't use use energy during the embedment
    # 2) Didn't compute energy step by step due to verbose being false
    # or
    # 3) FinalEnergy != "Penalized"

    if FixNodesAtPoints != []:
        partition, dists = PartitionData(
            X,
            NewNodePositions,
            MaxBlockSize,
            SquaredX,
            TrimmingRadius,
        )
    else:
        partition, dists = move_partition, move_dists
    # if FixNodesAtPoints != []:
    #    move_partition[move_partition >= 0] += len(FixNodesAtPoints)
    #    partition = np.concatenate((FixedNodesPart, move_partition))
    #    dists = np.concatenate((FixedNodesDists, move_dists))

    if (FinalEnergy != "Base") or (not (verbose) and (Mode != 2)):
        if FinalEnergy == "Base":
            ElasticEnergy, MSE, EP, RP = ComputePrimitiveGraphElasticEnergy(
                NewNodePositions, ElasticMatrix, dists
            )

        elif FinalEnergy == "Penalized":
            ElasticEnergy, MSE, EP, RP = ComputePenalizedPrimitiveGraphElasticEnergy(
                NewNodePositions, ElasticMatrix, dists, alpha, beta
            )

    return NewNodePositions, ElasticEnergy, partition, dists, MSE, EP, RP


def PrimitiveElasticGraphEmbedment(
    X,
    NodePositions,
    ElasticMatrix,
    MaxNumberOfIterations=10,
    eps=0.01,
    Mode=1,
    FinalEnergy="Base",
    alpha=0,
    beta=0,
    prob=1,
    DisplayWarnings=True,
    PointWeights=None,
    MaxBlockSize=100000000,
    verbose=False,
    TrimmingRadius=float("inf"),
    SquaredX=None,
    FixNodesAtPoints=[],
):

    """
    #' Function fitting a primitive elastic graph to the data
    #'
    #' @param X is n-by-m matrix containing the positions of the n points in the m-dimensional space
    #' @param NodePositions is k-by-m matrix of positions of the graph nodes in the same space as X
    #' @param ElasticMatrix is a k-by-k symmetric matrix describing the connectivity and the elastic
    #' properties of the graph. Star elasticities (mu coefficients) are along the main diagonal
    #' (non-zero entries only for star centers), and the edge elasticity moduli are at non-diagonal elements.
    #' @param MaxNumberOfIterations is an integer number indicating the maximum number of iterations for the EM algorithm
    #' @param TrimmingRadius is a real value indicating the trimming radius, a parameter required for robust principal graphs
    #' (see https://github.com/auranic/Elastic-principal-graphs/wiki/Robust-principal-graphs)
    #' @param eps a real number indicating the minimal relative change in the nodenpositions
    #' to be considered the graph embedded (convergence criteria)
    #' @param verbose is a boolean indicating whether diagnostig informations should be plotted
    #' @param Mode integer, the energy mode. It can be 1 (difference is computed using the position of the nodes) and
    #' 2 (difference is computed using the changes in elestic energy of the configuraztions)
    #' @param SquaredX the sum (by node) of X squared. It not specified, it will be calculated by the fucntion
    #' @param FastSolve boolean, shuold the Fastsolve of Armadillo by enabled?
    #' @param DisplayWarnings boolean, should warning about convergence be displayed?
    #' @param FinalEnergy string indicating the final elastic emergy associated with the configuration. Currently it can be "Base" or "Penalized"
    #' @param alpha positive numeric, the value of the alpha parameter of the penalized elastic energy
    #' @param beta positive numeric, the value of the beta parameter of the penalized elastic energy
    #' @param prob numeric between 0 and 1. If less than 1 point will be sampled at each iteration. Prob indicate the probability of
    #' using each points. This is an *experimental* feature, which may helps speeding up the computation if a large number of points is present.
    #'
    #' @return
    #' @export
    #'
    #' @examples
    """

    if prob < 1:
        raise ValueError("probPoint < 1 option not implemented yet")
    N = X.shape[0]

    # Auxiliary computations
    if PointWeights is None:
        PointWeights = np.ones((N, 1))
    if SquaredX is None:
        SquaredX = (X ** 2).sum(axis=1).reshape((N, 1))
    SpringLaplacianMatrix = ComputeSpringLaplacianMatrix(ElasticMatrix)

    if FixNodesAtPoints != []:
        # -----Index data and graph into moving nodes and fixed nodes------#
        flat_FixNodesAtPoints = [
            item for sublist in FixNodesAtPoints for item in sublist
        ]  # fixed datapoints

        # partition, dists = PartitionData(
        #    X, NodePositions, MaxBlockSize, SquaredX, TrimmingRadius,
        # )
        # move_data_idx = np.where(~np.isin(partition, range(len(FixNodesAtPoints))))[0]

        move_data_idx = [i for i in range(len(X)) if i not in flat_FixNodesAtPoints]
        move_nodes_idx = np.arange(len(FixNodesAtPoints), len(NodePositions))
    else:
        move_data_idx = np.arange(len(X))
        move_nodes_idx = np.arange(len(NodePositions))

    # fitted subset
    move_X, move_PointWeights, move_SquaredX, move_NodePositions = (
        X[move_data_idx, :],
        PointWeights[move_data_idx],
        SquaredX[move_data_idx],
        NodePositions[move_nodes_idx],
    )

    # Main iterative EM cycle: partition, fit given the partition, repeat
    move_partition, move_dists = PartitionData(
        move_X,
        move_NodePositions,
        MaxBlockSize,
        move_SquaredX,
        TrimmingRadius,
    )

    if verbose or Mode == 2:
        # if FixNodesAtPoints != []:
        #    dists = np.concatenate((FixedNodesDists, move_dists))
        OldElasticEnergy, MSE, EP, RP = ComputePrimitiveGraphElasticEnergy(
            NodePositions, ElasticMatrix, move_dists
        )

    ElasticEnergy = 0
    for i in range(MaxNumberOfIterations):
        # Updated positions
        if FixNodesAtPoints != []:
            NewNodePositions = FitSubGraph2DataGivenPartition(
                move_X,
                move_PointWeights,
                SpringLaplacianMatrix,
                NodePositions,
                move_partition,
                move_nodes_idx,
            )
        else:
            NewNodePositions = FitGraph2DataGivenPartition(
                move_X,
                move_PointWeights,
                SpringLaplacianMatrix,
                move_partition,
            )

        # Look at differences
        if verbose or Mode == 2:
            # if FixNodesAtPoints != []:
            #    dists = np.concatenate((FixedNodesDists, move_dists))
            ElasticEnergy, MSE, EP, RP = ComputePrimitiveGraphElasticEnergy(
                NewNodePositions, ElasticMatrix, move_dists, move_PointWeights
            )

        if Mode == 1:
            diff = ComputeRelativeChangeOfNodePositions(NodePositions, NewNodePositions)
        elif Mode == 2:
            diff = (OldElasticEnergy - ElasticEnergy) / ElasticEnergy

        # Print Info
        if verbose:
            print(
                "Iteration ",
                (i + 1),
                " difference of node position=",
                np.around(diff, 5),
                ", Energy=",
                np.around(ElasticEnergy, 5),
                ", MSE=",
                np.around(MSE, 5),
                ", EP=",
                np.around(EP, 5),
                ", RP=",
                np.around(RP, 5),
            )

        # Have we converged?
        if not np.isfinite(diff):
            diff = 0
        if diff < eps:
            break

        elif i < MaxNumberOfIterations - 1:
            move_NodePositions = NewNodePositions[move_nodes_idx]
            move_partition, move_dists = PartitionData(
                move_X,
                move_NodePositions,
                MaxBlockSize,
                move_SquaredX,
                TrimmingRadius,
            )

            NodePositions = NewNodePositions
            OldElasticEnergy = ElasticEnergy

    if DisplayWarnings and not (diff < eps):
        print(
            "Maximum number of iterations (",
            MaxNumberOfIterations,
            ") has been reached. diff = ",
            diff,
        )

    # If we
    # 1) Didn't use use energy during the embedment
    # 2) Didn't compute energy step by step due to verbose being false
    # or
    # 3) FinalEnergy != "Penalized"

    if FixNodesAtPoints != []:
        partition, dists = PartitionData(
            X,
            NewNodePositions,
            MaxBlockSize,
            SquaredX,
            TrimmingRadius,
        )
    else:
        partition, dists = move_partition, move_dists
    # if FixNodesAtPoints != []:
    #    move_partition[move_partition >= 0] += len(FixNodesAtPoints)
    #    partition = np.concatenate((FixedNodesPart, move_partition))
    #    dists = np.concatenate((FixedNodesDists, move_dists))

    if (FinalEnergy != "Base") or (not (verbose) and (Mode != 2)):
        if FinalEnergy == "Base":
            ElasticEnergy, MSE, EP, RP = ComputePrimitiveGraphElasticEnergy(
                NewNodePositions, ElasticMatrix, dists, PointWeights
            )

        elif FinalEnergy == "Penalized":
            ElasticEnergy, MSE, EP, RP = ComputePenalizedPrimitiveGraphElasticEnergy_v2(
                NewNodePositions,
                ElasticMatrix,
                dists,
                alpha,
                beta,
                PointWeights=PointWeights,
            )

    return NewNodePositions, ElasticEnergy, partition, dists, MSE, EP, RP


def PrimitiveElasticGraphEmbedment_cp(
    X,
    NodePositions,
    ElasticMatrix,
    MaxNumberOfIterations=10,
    eps=0.01,
    Mode=1,
    FinalEnergy="Base",
    alpha=0,
    beta=0,
    prob=1,
    DisplayWarnings=True,
    PointWeights=None,
    MaxBlockSize=100000000,
    verbose=False,
    TrimmingRadius=float("inf"),
    SquaredX=None,
    Xcp=None,
    SquaredXcp=None,
    FixNodesAtPoints=[],
):

    """
    #' Function fitting a primitive elastic graph to the data
    #'
    #' @param X is n-by-m matrix containing the positions of the n points in the m-dimensional space
    #' @param NodePositions is k-by-m matrix of positions of the graph nodes in the same space as X
    #' @param ElasticMatrix is a k-by-k symmetric matrix describing the connectivity and the elastic
    #' properties of the graph. Star elasticities (mu coefficients) are along the main diagonal
    #' (non-zero entries only for star centers), and the edge elasticity moduli are at non-diagonal elements.
    #' @param MaxNumberOfIterations is an integer number indicating the maximum number of iterations for the EM algorithm
    #' @param TrimmingRadius is a real value indicating the trimming radius, a parameter required for robust principal graphs
    #' (see https://github.com/auranic/Elastic-principal-graphs/wiki/Robust-principal-graphs)
    #' @param eps a real number indicating the minimal relative change in the nodenpositions
    #' to be considered the graph embedded (convergence criteria)
    #' @param verbose is a boolean indicating whether diagnostig informations should be plotted
    #' @param Mode integer, the energy mode. It can be 1 (difference is computed using the position of the nodes) and
    #' 2 (difference is computed using the changes in elestic energy of the configuraztions)
    #' @param SquaredX the sum (by node) of X squared. It not specified, it will be calculated by the fucntion
    #' @param FastSolve boolean, shuold the Fastsolve of Armadillo by enabled?
    #' @param DisplayWarnings boolean, should warning about convergence be displayed?
    #' @param FinalEnergy string indicating the final elastic emergy associated with the configuration. Currently it can be "Base" or "Penalized"
    #' @param alpha positive numeric, the value of the alpha parameter of the penalized elastic energy
    #' @param beta positive numeric, the value of the beta parameter of the penalized elastic energy
    #' @param prob numeric between 0 and 1. If less than 1 point will be sampled at each iteration. Prob indicate the probability of
    #' using each points. This is an *experimental* feature, which may helps speeding up the computation if a large number of points is present.
    #'
    #' @return
    #' @export
    #'
    #' @examples
    """

    if prob < 1:
        raise ValueError("probPoint < 1 option not implemented yet")
    N = X.shape[0]

    # Auxiliary computations
    if PointWeights is None:
        PointWeights = np.ones((N, 1))

    SpringLaplacianMatrix = ComputeSpringLaplacianMatrix(ElasticMatrix)

    if FixNodesAtPoints != []:
        # -----Index data and graph into moving nodes and fixed nodes------#
        flat_FixNodesAtPoints = [
            item for sublist in FixNodesAtPoints for item in sublist
        ]  # fixed datapoints

        # partition, dists = PartitionData(
        #    X, NodePositions, MaxBlockSize, SquaredX, TrimmingRadius,
        # )
        # move_data_idx = np.where(~np.isin(partition, range(len(FixNodesAtPoints))))[0]

        move_data_idx = np.array(
            [i for i in range(len(X)) if i not in flat_FixNodesAtPoints]
        )
        move_nodes_idx = np.arange(len(FixNodesAtPoints), len(NodePositions))
    else:
        move_data_idx = np.arange(len(X))
        move_nodes_idx = np.arange(len(NodePositions))

    # fitted subset
    (move_X, move_Xcp, move_PointWeights, move_SquaredXcp, move_NodePositions,) = (
        X[move_data_idx, :],
        Xcp[move_data_idx, :],
        PointWeights[move_data_idx],
        SquaredXcp[move_data_idx],
        NodePositions[move_nodes_idx],
    )

    # Main iterative EM cycle: partition, fit given the partition, repeat
    move_partition, move_dists = PartitionData_cp(
        move_Xcp,
        move_NodePositions,
        MaxBlockSize,
        move_SquaredXcp,
        TrimmingRadius,
    )

    if verbose or Mode == 2:
        # if FixNodesAtPoints != []:
        #    dists = np.concatenate((FixedNodesDists, move_dists))
        OldElasticEnergy, MSE, EP, RP = ComputePrimitiveGraphElasticEnergy(
            NodePositions, ElasticMatrix, move_dists
        )

    ElasticEnergy = 0
    for i in range(MaxNumberOfIterations):
        # Updated positions
        if FixNodesAtPoints != []:
            NewNodePositions = FitSubGraph2DataGivenPartition(
                move_X,
                move_PointWeights,
                SpringLaplacianMatrix,
                NodePositions,
                move_partition,
                move_nodes_idx,
            )
        else:
            NewNodePositions = FitGraph2DataGivenPartition(
                move_X,
                move_PointWeights,
                SpringLaplacianMatrix,
                move_partition,
            )

        # Look at differences
        if verbose or Mode == 2:
            # if FixNodesAtPoints != []:
            #    dists = np.concatenate((FixedNodesDists, move_dists))
            ElasticEnergy, MSE, EP, RP = ComputePrimitiveGraphElasticEnergy(
                NewNodePositions, ElasticMatrix, move_dists, move_PointWeights
            )

        if Mode == 1:
            diff = ComputeRelativeChangeOfNodePositions(NodePositions, NewNodePositions)
        elif Mode == 2:
            diff = (OldElasticEnergy - ElasticEnergy) / ElasticEnergy

        # Print Info
        if verbose:
            print(
                "Iteration ",
                (i + 1),
                " difference of node position=",
                np.around(diff, 5),
                ", Energy=",
                np.around(ElasticEnergy, 5),
                ", MSE=",
                np.around(MSE, 5),
                ", EP=",
                np.around(EP, 5),
                ", RP=",
                np.around(RP, 5),
            )

        # Have we converged?
        if not np.isfinite(diff):
            diff = 0
        if diff < eps:
            break

        elif i < MaxNumberOfIterations - 1:
            move_NodePositions = NewNodePositions[move_nodes_idx]
            move_partition, move_dists = PartitionData_cp(
                move_Xcp,
                move_NodePositions,
                MaxBlockSize,
                move_SquaredXcp,
                TrimmingRadius,
            )

            NodePositions = NewNodePositions
            OldElasticEnergy = ElasticEnergy

    if DisplayWarnings and not (diff < eps):
        print(
            "Maximum number of iterations (",
            MaxNumberOfIterations,
            ") has been reached. diff = ",
            diff,
        )

    # If we
    # 1) Didn't use use energy during the embedment
    # 2) Didn't compute energy step by step due to verbose being false
    # or
    # 3) FinalEnergy != "Penalized"

    if FixNodesAtPoints != []:
        partition, dists = PartitionData_cp(
            Xcp,
            NewNodePositions,
            MaxBlockSize,
            SquaredXcp,
            TrimmingRadius,
        )
    else:
        partition, dists = move_partition, move_dists
    # if FixNodesAtPoints != []:
    #    move_partition[move_partition >= 0] += len(FixNodesAtPoints)
    #    partition = np.concatenate((FixedNodesPart, move_partition))
    #    dists = np.concatenate((FixedNodesDists, move_dists))

    if (FinalEnergy != "Base") or (not (verbose) and (Mode != 2)):
        if FinalEnergy == "Base":
            ElasticEnergy, MSE, EP, RP = ComputePrimitiveGraphElasticEnergy(
                NewNodePositions, ElasticMatrix, dists, PointWeights
            )

        elif FinalEnergy == "Penalized":
            ElasticEnergy, MSE, EP, RP = ComputePenalizedPrimitiveGraphElasticEnergy_v2(
                NewNodePositions,
                ElasticMatrix,
                dists,
                alpha,
                beta,
                PointWeights=PointWeights,
            )

    return NewNodePositions, ElasticEnergy, partition, dists, MSE, EP, RP


def PrimitiveElasticGraphEmbedment_v2(
    X,
    NodePositions,
    ElasticMatrix,
    MaxNumberOfIterations=10,
    eps=0.01,
    Mode=1,
    FinalEnergy="Base",
    alpha=0,
    beta=0,
    prob=1,
    DisplayWarnings=True,
    PointWeights=None,
    MaxBlockSize=100000000,
    verbose=False,
    TrimmingRadius=float("inf"),
    SquaredX=None,
    FixNodesAtPoints=[],
    PseudotimeNodePositions=None,
    PseudotimeLambda=None,
    partition=None,
    dists=None,
    precomp_d=None,
):

    """
    #' Function fitting a primitive elastic graph to the data
    #'
    #' @param X is n-by-m matrix containing the positions of the n points in the m-dimensional space
    #' @param NodePositions is k-by-m matrix of positions of the graph nodes in the same space as X
    #' @param ElasticMatrix is a k-by-k symmetric matrix describing the connectivity and the elastic
    #' properties of the graph. Star elasticities (mu coefficients) are along the main diagonal
    #' (non-zero entries only for star centers), and the edge elasticity moduli are at non-diagonal elements.
    #' @param MaxNumberOfIterations is an integer number indicating the maximum number of iterations for the EM algorithm
    #' @param TrimmingRadius is a real value indicating the trimming radius, a parameter required for robust principal graphs
    #' (see https://github.com/auranic/Elastic-principal-graphs/wiki/Robust-principal-graphs)
    #' @param eps a real number indicating the minimal relative change in the nodenpositions
    #' to be considered the graph embedded (convergence criteria)
    #' @param verbose is a boolean indicating whether diagnostig informations should be plotted
    #' @param Mode integer, the energy mode. It can be 1 (difference is computed using the position of the nodes) and
    #' 2 (difference is computed using the changes in elestic energy of the configuraztions)
    #' @param SquaredX the sum (by node) of X squared. It not specified, it will be calculated by the fucntion
    #' @param FastSolve boolean, shuold the Fastsolve of Armadillo by enabled?
    #' @param DisplayWarnings boolean, should warning about convergence be displayed?
    #' @param FinalEnergy string indicating the final elastic emergy associated with the configuration. Currently it can be "Base" or "Penalized"
    #' @param alpha positive numeric, the value of the alpha parameter of the penalized elastic energy
    #' @param beta positive numeric, the value of the beta parameter of the penalized elastic energy
    #' @param prob numeric between 0 and 1. If less than 1 point will be sampled at each iteration. Prob indicate the probability of
    #' using each points. This is an *experimental* feature, which may helps speeding up the computation if a large number of points is present.
    #'
    #' @return
    #' @export
    #'
    #' @examples
    """

    if prob < 1:
        raise ValueError("probPoint < 1 option not implemented yet")
    N = X.shape[0]

    # Auxiliary computations
    if PointWeights is None:
        PointWeights = np.ones((N, 1))
    if SquaredX is None:
        SquaredX = (X ** 2).sum(axis=1).reshape((N, 1))
    SpringLaplacianMatrix = ComputeSpringLaplacianMatrix(ElasticMatrix)

    if FixNodesAtPoints != []:
        # -----Index data and graph into moving nodes and fixed nodes------#
        # flat_FixNodesAtPoints = [
        #    item for sublist in FixNodesAtPoints for item in sublist
        # ]  # fixed datapoints

        # partition, dists = PartitionData(
        #    X,
        #    NodePositions,
        #    MaxBlockSize,
        #    SquaredX,
        #    TrimmingRadius,
        # )
        move_data_idx = np.where(~np.isin(partition, range(len(FixNodesAtPoints))))[0]

        # move_data_idx = [i for i in range(len(X)) if i not in flat_FixNodesAtPoints]
        move_nodes_idx = np.arange(len(FixNodesAtPoints), len(NodePositions))
        partition, dists, _ = RePartitionData(
            None,
            None,
            None,
            NodeIndicesArray=move_nodes_idx,
            opType="removenode",
            precomp_d=precomp_d,
            SquaredX=None,
        )
        move_partition, move_dists = partition[move_data_idx], dists[move_data_idx]
    else:
        move_data_idx = np.arange(len(X))
        move_nodes_idx = np.arange(len(NodePositions))
        move_partition = partition
        move_dists = dists

    # fitted subset
    move_X, move_PointWeights, move_SquaredX, move_NodePositions = (
        X[move_data_idx, :],
        PointWeights[move_data_idx],
        SquaredX[move_data_idx],
        NodePositions[move_nodes_idx],
    )

    if verbose or Mode == 2:
        # if FixNodesAtPoints != []:
        #    dists = np.concatenate((FixedNodesDists, move_dists))
        OldElasticEnergy, MSE, EP, RP = ComputePrimitiveGraphElasticEnergy(
            NodePositions, ElasticMatrix, move_dists
        )

    ElasticEnergy = 0
    for i in range(MaxNumberOfIterations):
        # Updated positions
        if FixNodesAtPoints != []:
            NewNodePositions = FitSubGraph2DataGivenPartition_v2(
                move_X,
                move_PointWeights,
                SpringLaplacianMatrix,
                NodePositions,
                move_partition,
                move_nodes_idx,
                PseudotimeNodePositions=PseudotimeNodePositions,
                PseudotimeLambda=PseudotimeLambda,
            )
        else:
            NewNodePositions = FitGraph2DataGivenPartition(
                move_X,
                move_PointWeights,
                SpringLaplacianMatrix,
                move_partition,
            )

        # Look at differences
        if verbose or Mode == 2:
            # if FixNodesAtPoints != []:
            #    dists = np.concatenate((FixedNodesDists, move_dists))
            ElasticEnergy, MSE, EP, RP = ComputePrimitiveGraphElasticEnergy(
                NewNodePositions, ElasticMatrix, move_dists
            )

        if Mode == 1:
            diff = ComputeRelativeChangeOfNodePositions(NodePositions, NewNodePositions)
        elif Mode == 2:
            diff = (OldElasticEnergy - ElasticEnergy) / ElasticEnergy

        # Print Info
        if verbose:
            print(
                "Iteration ",
                (i + 1),
                " difference of node position=",
                np.around(diff, 5),
                ", Energy=",
                np.around(ElasticEnergy, 5),
                ", MSE=",
                np.around(MSE, 5),
                ", EP=",
                np.around(EP, 5),
                ", RP=",
                np.around(RP, 5),
            )

        # Have we converged?
        if not np.isfinite(diff):
            diff = 0
        if diff < eps:
            break

        elif i < MaxNumberOfIterations - 1:
            move_NodePositions = NewNodePositions[move_nodes_idx]
            move_partition, move_dists = PartitionData(
                move_X,
                move_NodePositions,
                MaxBlockSize,
                move_SquaredX,
                TrimmingRadius,
            )

            NodePositions = NewNodePositions
            OldElasticEnergy = ElasticEnergy

    if DisplayWarnings and not (diff < eps):
        print(
            "Maximum number of iterations (",
            MaxNumberOfIterations,
            ") has been reached. diff = ",
            diff,
        )

    # If we
    # 1) Didn't use use energy during the embedment
    # 2) Didn't compute energy step by step due to verbose being false
    # or
    # 3) FinalEnergy != "Penalized"

    if FixNodesAtPoints != []:
        partition, dists = PartitionData(
            X,
            NewNodePositions,
            MaxBlockSize,
            SquaredX,
            TrimmingRadius,
        )
    else:
        partition, dists = move_partition, move_dists
    # if FixNodesAtPoints != []:
    #    move_partition[move_partition >= 0] += len(FixNodesAtPoints)
    #    partition = np.concatenate((FixedNodesPart, move_partition))
    #    dists = np.concatenate((FixedNodesDists, move_dists))

    if (FinalEnergy != "Base") or (not (verbose) and (Mode != 2)):
        if FinalEnergy == "Base":
            ElasticEnergy, MSE, EP, RP = ComputePrimitiveGraphElasticEnergy(
                NewNodePositions, ElasticMatrix, dists, PointWeights
            )

        elif FinalEnergy == "Penalized":
            ElasticEnergy, MSE, EP, RP = ComputePenalizedPrimitiveGraphElasticEnergy_v2(
                NewNodePositions,
                ElasticMatrix,
                dists,
                alpha,
                beta,
                PointWeights,
                PseudotimeNodePositions,
            )

    return NewNodePositions, ElasticEnergy, partition, dists, MSE, EP, RP


def PrimitiveElasticGraphEmbedment_v3(
    X,
    NodePositions,
    ElasticMatrix,
    MaxNumberOfIterations=10,
    eps=0.01,
    Mode=1,
    FinalEnergy="Base",
    alpha=0,
    beta=0,
    prob=1,
    DisplayWarnings=True,
    PointWeights=None,
    MaxBlockSize=100000000,
    verbose=False,
    TrimmingRadius=float("inf"),
    SquaredX=None,
    FixNodesAtPoints=[],
    PseudotimeNodePositions=None,
    PseudotimeLambda=None,
    Label=None,
    LabelLambda=None,
):

    """
    #' Function fitting a primitive elastic graph to the data
    #'
    #' @param X is n-by-m matrix containing the positions of the n points in the m-dimensional space
    #' @param NodePositions is k-by-m matrix of positions of the graph nodes in the same space as X
    #' @param ElasticMatrix is a k-by-k symmetric matrix describing the connectivity and the elastic
    #' properties of the graph. Star elasticities (mu coefficients) are along the main diagonal
    #' (non-zero entries only for star centers), and the edge elasticity moduli are at non-diagonal elements.
    #' @param MaxNumberOfIterations is an integer number indicating the maximum number of iterations for the EM algorithm
    #' @param TrimmingRadius is a real value indicating the trimming radius, a parameter required for robust principal graphs
    #' (see https://github.com/auranic/Elastic-principal-graphs/wiki/Robust-principal-graphs)
    #' @param eps a real number indicating the minimal relative change in the nodenpositions
    #' to be considered the graph embedded (convergence criteria)
    #' @param verbose is a boolean indicating whether diagnostig informations should be plotted
    #' @param Mode integer, the energy mode. It can be 1 (difference is computed using the position of the nodes) and
    #' 2 (difference is computed using the changes in elestic energy of the configuraztions)
    #' @param SquaredX the sum (by node) of X squared. It not specified, it will be calculated by the fucntion
    #' @param FastSolve boolean, shuold the Fastsolve of Armadillo by enabled?
    #' @param DisplayWarnings boolean, should warning about convergence be displayed?
    #' @param FinalEnergy string indicating the final elastic emergy associated with the configuration. Currently it can be "Base" or "Penalized"
    #' @param alpha positive numeric, the value of the alpha parameter of the penalized elastic energy
    #' @param beta positive numeric, the value of the beta parameter of the penalized elastic energy
    #' @param prob numeric between 0 and 1. If less than 1 point will be sampled at each iteration. Prob indicate the probability of
    #' using each points. This is an *experimental* feature, which may helps speeding up the computation if a large number of points is present.
    #'
    #' @return
    #' @export
    #'
    #' @examples
    """

    if prob < 1:
        raise ValueError("probPoint < 1 option not implemented yet")
    N = X.shape[0]

    # Auxiliary computations
    if PointWeights is None:
        PointWeights = np.ones((N, 1))
    if SquaredX is None:
        SquaredX = (X ** 2).sum(axis=1).reshape((N, 1))
    SpringLaplacianMatrix = ComputeSpringLaplacianMatrix(ElasticMatrix)

    if FixNodesAtPoints != []:
        # -----Index data and graph into moving nodes and fixed nodes------#
        # flat_FixNodesAtPoints = [
        #    item for sublist in FixNodesAtPoints for item in sublist
        # ]  # fixed datapoints

        partition, dists = PartitionData(
            X,
            NodePositions,
            MaxBlockSize,
            SquaredX,
            TrimmingRadius,
        )
        move_data_idx = np.where(~np.isin(partition, range(len(FixNodesAtPoints))))[0]

        # move_data_idx = [i for i in range(len(X)) if i not in flat_FixNodesAtPoints]
        move_nodes_idx = np.arange(len(FixNodesAtPoints), len(NodePositions))
    else:
        move_data_idx = np.arange(len(X))
        move_nodes_idx = np.arange(len(NodePositions))

    # fitted subset
    move_X, move_PointWeights, move_SquaredX, move_NodePositions = (
        X[move_data_idx, :],
        PointWeights[move_data_idx],
        SquaredX[move_data_idx],
        NodePositions[move_nodes_idx],
    )

    # Main iterative EM cycle: partition, fit given the partition, repeat
    move_partition, move_dists = PartitionData(
        move_X,
        move_NodePositions,
        MaxBlockSize,
        move_SquaredX,
        TrimmingRadius,
    )

    if verbose or Mode == 2:
        # if FixNodesAtPoints != []:
        #    dists = np.concatenate((FixedNodesDists, move_dists))
        OldElasticEnergy, MSE, EP, RP = ComputePrimitiveGraphElasticEnergy(
            NodePositions, ElasticMatrix, move_dists
        )

    ElasticEnergy = 0
    for i in range(MaxNumberOfIterations):
        # Updated positions
        if FixNodesAtPoints != []:
            NewNodePositions = FitSubGraph2DataGivenPartition_v2(
                move_X,
                move_PointWeights,
                SpringLaplacianMatrix,
                NodePositions,
                move_partition,
                move_nodes_idx,
                PseudotimeNodePositions=PseudotimeNodePositions,
                PseudotimeLambda=PseudotimeLambda,
            )
        else:
            NewNodePositions = FitGraph2DataGivenPartition(
                move_X,
                move_PointWeights,
                SpringLaplacianMatrix,
                move_partition,
            )

        # Look at differences
        if verbose or Mode == 2:
            # if FixNodesAtPoints != []:
            #    dists = np.concatenate((FixedNodesDists, move_dists))
            ElasticEnergy, MSE, EP, RP = ComputePrimitiveGraphElasticEnergy(
                NewNodePositions, ElasticMatrix, move_dists
            )

        if Mode == 1:
            diff = ComputeRelativeChangeOfNodePositions(NodePositions, NewNodePositions)
        elif Mode == 2:
            diff = (OldElasticEnergy - ElasticEnergy) / ElasticEnergy

        # Print Info
        if verbose:
            print(
                "Iteration ",
                (i + 1),
                " difference of node position=",
                np.around(diff, 5),
                ", Energy=",
                np.around(ElasticEnergy, 5),
                ", MSE=",
                np.around(MSE, 5),
                ", EP=",
                np.around(EP, 5),
                ", RP=",
                np.around(RP, 5),
            )

        # Have we converged?
        if not np.isfinite(diff):
            diff = 0
        if diff < eps:
            break

        elif i < MaxNumberOfIterations - 1:
            move_NodePositions = NewNodePositions[move_nodes_idx]
            move_partition, move_dists = PartitionData(
                move_X,
                move_NodePositions,
                MaxBlockSize,
                move_SquaredX,
                TrimmingRadius,
            )

            NodePositions = NewNodePositions
            OldElasticEnergy = ElasticEnergy

    if DisplayWarnings and not (diff < eps):
        print(
            "Maximum number of iterations (",
            MaxNumberOfIterations,
            ") has been reached. diff = ",
            diff,
        )

    # If we
    # 1) Didn't use use energy during the embedment
    # 2) Didn't compute energy step by step due to verbose being false
    # or
    # 3) FinalEnergy != "Penalized"

    if FixNodesAtPoints != []:
        partition, dists = PartitionData(
            X,
            NewNodePositions,
            MaxBlockSize,
            SquaredX,
            TrimmingRadius,
        )
    else:
        partition, dists = move_partition, move_dists
    # if FixNodesAtPoints != []:
    #    move_partition[move_partition >= 0] += len(FixNodesAtPoints)
    #    partition = np.concatenate((FixedNodesPart, move_partition))
    #    dists = np.concatenate((FixedNodesDists, move_dists))

    if (FinalEnergy != "Base") or (not (verbose) and (Mode != 2)):
        if FinalEnergy == "Base":
            ElasticEnergy, MSE, EP, RP = ComputePrimitiveGraphElasticEnergy(
                NewNodePositions, ElasticMatrix, dists
            )

        elif FinalEnergy == "Penalized":
            ElasticEnergy, MSE, EP, RP = ComputePenalizedPrimitiveGraphElasticEnergy_v3(
                NewNodePositions,
                ElasticMatrix,
                partition.flatten(),
                dists,
                alpha,
                beta,
                PseudotimeNodePositions,
                Label,
                LabelLambda,
            )

    return NewNodePositions, ElasticEnergy, partition, dists, MSE, EP, RP


def PrimitiveElasticGraphEmbedment_cp_v2(
    X,
    NodePositions,
    ElasticMatrix,
    MaxNumberOfIterations=10,
    eps=0.01,
    Mode=1,
    FinalEnergy="Base",
    alpha=0,
    beta=0,
    prob=1,
    DisplayWarnings=True,
    PointWeights=None,
    MaxBlockSize=100000000,
    verbose=False,
    TrimmingRadius=float("inf"),
    SquaredX=None,
    FixNodesAtPoints=[],
    PseudotimeNodePositions=None,
    PseudotimeLambda=None,
    partition=None,
    dists=None,
    precomp_d=None,
    Xcp=None,
    SquaredXcp=None,
):

    """
    #' Function fitting a primitive elastic graph to the data
    #'
    #' @param X is n-by-m matrix containing the positions of the n points in the m-dimensional space
    #' @param NodePositions is k-by-m matrix of positions of the graph nodes in the same space as X
    #' @param ElasticMatrix is a k-by-k symmetric matrix describing the connectivity and the elastic
    #' properties of the graph. Star elasticities (mu coefficients) are along the main diagonal
    #' (non-zero entries only for star centers), and the edge elasticity moduli are at non-diagonal elements.
    #' @param MaxNumberOfIterations is an integer number indicating the maximum number of iterations for the EM algorithm
    #' @param TrimmingRadius is a real value indicating the trimming radius, a parameter required for robust principal graphs
    #' (see https://github.com/auranic/Elastic-principal-graphs/wiki/Robust-principal-graphs)
    #' @param eps a real number indicating the minimal relative change in the nodenpositions
    #' to be considered the graph embedded (convergence criteria)
    #' @param verbose is a boolean indicating whether diagnostig informations should be plotted
    #' @param Mode integer, the energy mode. It can be 1 (difference is computed using the position of the nodes) and
    #' 2 (difference is computed using the changes in elestic energy of the configuraztions)
    #' @param SquaredX the sum (by node) of X squared. It not specified, it will be calculated by the fucntion
    #' @param FastSolve boolean, shuold the Fastsolve of Armadillo by enabled?
    #' @param DisplayWarnings boolean, should warning about convergence be displayed?
    #' @param FinalEnergy string indicating the final elastic emergy associated with the configuration. Currently it can be "Base" or "Penalized"
    #' @param alpha positive numeric, the value of the alpha parameter of the penalized elastic energy
    #' @param beta positive numeric, the value of the beta parameter of the penalized elastic energy
    #' @param prob numeric between 0 and 1. If less than 1 point will be sampled at each iteration. Prob indicate the probability of
    #' using each points. This is an *experimental* feature, which may helps speeding up the computation if a large number of points is present.
    #'
    #' @return
    #' @export
    #'
    #' @examples
    """

    if prob < 1:
        raise ValueError("probPoint < 1 option not implemented yet")
    N = X.shape[0]

    # Auxiliary computations
    if PointWeights is None:
        PointWeights = np.ones((N, 1))
    if SquaredX is None:
        SquaredX = (X ** 2).sum(axis=1).reshape((N, 1))
    SpringLaplacianMatrix = ComputeSpringLaplacianMatrix(ElasticMatrix)

    if FixNodesAtPoints != []:
        # -----Index data and graph into moving nodes and fixed nodes------#
        # flat_FixNodesAtPoints = [
        #    item for sublist in FixNodesAtPoints for item in sublist
        # ]  # fixed datapoints

        # partition, dists = PartitionData(
        #    X,
        #    NodePositions,
        #    MaxBlockSize,
        #    SquaredX,
        #    TrimmingRadius,
        # )
        move_data_idx = np.where(~np.isin(partition, range(len(FixNodesAtPoints))))[0]

        # move_data_idx = [i for i in range(len(X)) if i not in flat_FixNodesAtPoints]
        move_nodes_idx = np.arange(len(FixNodesAtPoints), len(NodePositions))
        partition, dists, _ = RePartitionData(
            None,
            None,
            None,
            NodeIndicesArray=move_nodes_idx,
            opType="removenode",
            precomp_d=precomp_d,
            SquaredX=None,
        )
        move_partition, move_dists = partition[move_data_idx], dists[move_data_idx]
    else:
        move_data_idx = np.arange(len(X))
        move_nodes_idx = np.arange(len(NodePositions))
        move_partition = partition
        move_dists = dists

    # fitted subset
    move_X, move_Xcp, move_PointWeights, move_SquaredXcp, move_NodePositions = (
        X[move_data_idx, :],
        Xcp[move_data_idx, :],
        PointWeights[move_data_idx],
        SquaredXcp[move_data_idx],
        NodePositions[move_nodes_idx],
    )

    if verbose or Mode == 2:
        # if FixNodesAtPoints != []:
        #    dists = np.concatenate((FixedNodesDists, move_dists))
        OldElasticEnergy, MSE, EP, RP = ComputePrimitiveGraphElasticEnergy(
            NodePositions,
            ElasticMatrix,
            move_dists,
            move_PointWeights,
        )

    ElasticEnergy = 0
    for i in range(MaxNumberOfIterations):
        # Updated positions
        if FixNodesAtPoints != []:
            NewNodePositions = FitSubGraph2DataGivenPartition_v2(
                move_X,
                move_PointWeights,
                SpringLaplacianMatrix,
                NodePositions,
                move_partition,
                move_nodes_idx,
                PseudotimeNodePositions=PseudotimeNodePositions,
                PseudotimeLambda=PseudotimeLambda,
            )
        else:
            # NewNodePositions = FitGraph2DataGivenPartition(
            #    move_X,
            #    move_PointWeights,
            #    SpringLaplacianMatrix,
            #    move_partition,
            # )
            NewNodePositions = FitSubGraph2DataGivenPartition_v2(
                move_X,
                move_PointWeights,
                SpringLaplacianMatrix,
                NodePositions,
                move_partition,
                move_nodes_idx,
                PseudotimeNodePositions=PseudotimeNodePositions,
                PseudotimeLambda=PseudotimeLambda,
            )

        # Look at differences
        if verbose or Mode == 2:
            # if FixNodesAtPoints != []:
            #    dists = np.concatenate((FixedNodesDists, move_dists))
            ElasticEnergy, MSE, EP, RP = ComputePrimitiveGraphElasticEnergy(
                NewNodePositions, ElasticMatrix, move_dists, move_PointWeights
            )

        if Mode == 1:
            diff = ComputeRelativeChangeOfNodePositions(NodePositions, NewNodePositions)
        elif Mode == 2:
            diff = (OldElasticEnergy - ElasticEnergy) / ElasticEnergy

        # Print Info
        if verbose:
            print(
                "Iteration ",
                (i + 1),
                " difference of node position=",
                np.around(diff, 5),
                ", Energy=",
                np.around(ElasticEnergy, 5),
                ", MSE=",
                np.around(MSE, 5),
                ", EP=",
                np.around(EP, 5),
                ", RP=",
                np.around(RP, 5),
            )

        # Have we converged?
        if not np.isfinite(diff):
            diff = 0
        if diff < eps:
            break

        elif i < MaxNumberOfIterations - 1:
            move_NodePositions = NewNodePositions[move_nodes_idx]
            move_partition, move_dists = PartitionData_cp(
                move_Xcp,
                move_NodePositions,
                MaxBlockSize,
                move_SquaredXcp,
                TrimmingRadius,
            )

            NodePositions = NewNodePositions
            OldElasticEnergy = ElasticEnergy

    if DisplayWarnings and not (diff < eps):
        print(
            "Maximum number of iterations (",
            MaxNumberOfIterations,
            ") has been reached. diff = ",
            diff,
        )

    # If we
    # 1) Didn't use use energy during the embedment
    # 2) Didn't compute energy step by step due to verbose being false
    # or
    # 3) FinalEnergy != "Penalized"

    if FixNodesAtPoints != []:
        partition, dists = PartitionData_cp(
            Xcp,
            NewNodePositions,
            MaxBlockSize,
            SquaredXcp,
            TrimmingRadius,
        )
    else:
        partition, dists = move_partition, move_dists
    # if FixNodesAtPoints != []:
    #    move_partition[move_partition >= 0] += len(FixNodesAtPoints)
    #    partition = np.concatenate((FixedNodesPart, move_partition))
    #    dists = np.concatenate((FixedNodesDists, move_dists))

    if (FinalEnergy != "Base") or (not (verbose) and (Mode != 2)):
        if FinalEnergy == "Base":
            ElasticEnergy, MSE, EP, RP = ComputePrimitiveGraphElasticEnergy(
                NewNodePositions,
                ElasticMatrix,
                dists,
                PointWeights,
            )

        elif FinalEnergy == "Penalized":
            ElasticEnergy, MSE, EP, RP = ComputePenalizedPrimitiveGraphElasticEnergy_v2(
                NewNodePositions,
                ElasticMatrix,
                dists,
                alpha,
                beta,
                PointWeights,
                PseudotimeNodePositions,
            )

    return NewNodePositions, ElasticEnergy, partition, dists, MSE, EP, RP
