"""

Note
----

This file is copied (possibly with major modifications) from the
sources of the pycpd project - https://github.com/siavashk/pycpd.
It remains licensed as the rest of PyCPD (MIT license as of October 2010).

# ## ### ### ### ### ### ### ### ### ### ### ### ### ### ### ### ### ### ### ##
#
#   See COPYING file distributed along with the PyCPD package for the
#   copyright and license terms.
#
# ## ### ### ### ### ### ### ### ### ### ### ### ### ### ### ### ### ### ### ##
"""

import numbers
from warnings import warn

import numpy as np

from dipy.testing.decorators import warning_for_keywords


def gaussian_kernel(X, beta, Y=None):
    if Y is None:
        Y = X
    diff = X[:, None, :] - Y[None, :, :]
    diff = np.square(diff)
    diff = np.sum(diff, 2)
    return np.exp(-diff / (2 * beta**2))


def low_rank_eigen(G, num_eig):
    """Calculate num_eig eigenvectors and eigenvalues of gaussian matrix G.

    Enables lower dimensional solving.

    """
    S, Q = np.linalg.eigh(G)
    eig_indices = list(np.argsort(np.abs(S))[::-1][:num_eig])
    Q = Q[:, eig_indices]  # eigenvectors
    S = S[eig_indices]  # eigenvalues.
    return Q, S


def initialize_sigma2(X, Y):
    """Initialize the variance (sigma2).

    Parameters
    ----------
    X: numpy array
        NxD array of points for target.

    Y: numpy array
        MxD array of points for source.

    Returns
    -------
    sigma2: float
        Initial variance.
    """
    (N, D) = X.shape
    (M, _) = Y.shape
    diff = X[None, :, :] - Y[:, None, :]
    err = diff**2
    return np.sum(err) / (D * M * N)


@warning_for_keywords()
def lowrankQS(G, beta, num_eig, *, eig_fgt=False):
    """Calculate eigenvectors and eigenvalues of gaussian matrix G.

    !!!
    This function is a placeholder for implementing the fast
    gauss transform. It is not yet implemented.
    !!!

    Parameters
    ----------
    G: numpy array
        Gaussian kernel matrix.

    beta: float
        Width of the Gaussian kernel.

    num_eig: int
        Number of eigenvectors to use in lowrank calculation of G

    eig_fgt: bool
        If True, use fast gauss transform method to speed up.

    """
    # if we do not use FGT we construct affinity matrix G and find the
    # first eigenvectors/values directly

    if eig_fgt is False:
        S, Q = np.linalg.eigh(G)
        eig_indices = list(np.argsort(np.abs(S))[::-1][:num_eig])
        Q = Q[:, eig_indices]  # eigenvectors
        S = S[eig_indices]  # eigenvalues.

        return Q, S

    elif eig_fgt is True:
        raise Exception("Fast Gauss Transform Not Implemented!")


class DeformableRegistration:
    """
    Deformable point cloud registration.

    Attributes
    ----------
    X: numpy array
        NxD array of target points.

    Y: numpy array
        MxD array of source points.

    TY: numpy array
        MxD array of transformed source points.

    sigma2: float (positive)
        Initial variance of the Gaussian mixture model.

    N: int
        Number of target points.

    M: int
        Number of source points.

    D: int
        Dimensionality of source and target points

    iteration: int
        The current iteration throughout registration.

    max_iterations: int
        Registration will terminate once the algorithm has taken this
        many iterations.

    tolerance: float (positive)
        Registration will terminate once the difference between
        consecutive objective function values falls within this tolerance.

    w: float (between 0 and 1)
        Contribution of the uniform distribution to account for outliers.
        Valid values span 0 (inclusive) and 1 (exclusive).

    q: float
        The objective function value that represents the misalignment between
        source and target point clouds.

    diff: float (positive)
        The absolute difference between the current and previous objective
        function values.

    P: numpy array
        MxN array of probabilities.
        P[m, n] represents the probability that the m-th source point
        corresponds to the n-th target point.

    Pt1: numpy array
        Nx1 column array. Multiplication result between the transpose of P
        and a column vector of all 1s.

    P1: numpy array
        Mx1 column array.
        Multiplication result between P and a column vector of all 1s.

    Np: float (positive)
        The sum of all elements in P.

    alpha: float (positive)
        Represents the trade-off between the goodness of maximum likelihoo
        fit and regularization.

    beta: float(positive)
        Width of the Gaussian kernel.

    low_rank: bool
        Whether to use low rank approximation.

    num_eig: int
        Number of eigenvectors to use in lowrank calculation.
    """

    def __init__(
        self,
        X,
        Y,
        *args,
        sigma2=None,
        alpha=None,
        beta=None,
        low_rank=False,
        num_eig=100,
        max_iterations=None,
        tolerance=None,
        w=None,
        **kwargs,
    ):
        if not isinstance(X, np.ndarray) or X.ndim != 2:
            raise ValueError("The target point cloud (X) must be at a 2D numpy array.")

        if not isinstance(Y, np.ndarray) or Y.ndim != 2:
            raise ValueError("The source point cloud (Y) must be a 2D numpy array.")

        if X.shape[1] != Y.shape[1]:
            msg = "Both point clouds need to have the same number "
            msg += "of dimensions."
            raise ValueError(msg)

        if sigma2 is not None and (
            not isinstance(sigma2, numbers.Number) or sigma2 <= 0
        ):
            msg = f"Expected a positive value for sigma2 instead got: {sigma2}"
            raise ValueError(msg)

        if max_iterations is not None and (
            not isinstance(max_iterations, numbers.Number) or max_iterations < 0
        ):
            msg = "Expected a positive integer for max_iterations "
            msg += f"instead got: {max_iterations}"
            raise ValueError(msg)
        elif isinstance(max_iterations, numbers.Number) and not isinstance(
            max_iterations, int
        ):
            msg = "Received a non-integer value for max_iterations: "
            msg += f"{max_iterations}. Casting to integer."
            warn(msg, stacklevel=2)
            max_iterations = int(max_iterations)

        if tolerance is not None and (
            not isinstance(tolerance, numbers.Number) or tolerance < 0
        ):
            msg = "Expected a positive float for tolerance "
            msg += f"instead got: {tolerance}"
            raise ValueError(msg)

        if w is not None and (not isinstance(w, numbers.Number) or w < 0 or w >= 1):
            msg = "Expected a value between 0 (inclusive) and 1 (exclusive) "
            msg += f"for w instead got: {w}"
            raise ValueError(msg)

        self.X = X
        self.Y = Y
        self.TY = Y
        self.sigma2 = initialize_sigma2(X, Y) if sigma2 is None else sigma2
        (self.N, self.D) = self.X.shape
        (self.M, _) = self.Y.shape
        self.tolerance = 0.001 if tolerance is None else tolerance
        self.w = 0.0 if w is None else w
        self.max_iterations = 100 if max_iterations is None else max_iterations
        self.iteration = 0
        self.diff = np.inf
        self.q = np.inf
        self.P = np.zeros((self.M, self.N))
        self.Pt1 = np.zeros((self.N,))
        self.P1 = np.zeros((self.M,))
        self.PX = np.zeros((self.M, self.D))
        self.Np = 0

        if alpha is not None and (not isinstance(alpha, numbers.Number) or alpha <= 0):
            msg = "Expected a positive value for regularization parameter "
            msg += f"alpha. Instead got: {alpha}"
            raise ValueError(msg)

        if beta is not None and (not isinstance(beta, numbers.Number) or beta <= 0):
            msg = "Expected a positive value for the width of the coherent "
            msg += f"Gaussian kernel. Instead got: {beta}"

        self.alpha = 2 if alpha is None else alpha
        self.beta = 2 if beta is None else beta
        self.W = np.zeros((self.M, self.D))
        self.G = gaussian_kernel(self.Y, self.beta)
        self.low_rank = low_rank
        self.num_eig = num_eig
        if self.low_rank is True:
            self.Q, self.S = low_rank_eigen(self.G, self.num_eig)
            self.inv_S = np.diag(1.0 / self.S)
            self.S = np.diag(self.S)
            self.E = 0.0

    @warning_for_keywords()
    def register(self, *, callback=lambda **kwargs: None):
        """
        Perform the EM registration.

        Parameters
        ----------
        callback: function
            A function that will be called after each iteration.
            Can be used to visualize the registration process.

        Returns
        -------
        self.TY: numpy array
            MxD array of transformed source points.

        registration_parameters:
            Returned params dependent on registration method used.
        """
        self.transform_point_cloud()
        while self.iteration < self.max_iterations and self.diff > self.tolerance:
            self.iterate()
            if callable(callback):
                kwargs = {
                    "iteration": self.iteration,
                    "error": self.q,
                    "X": self.X,
                    "Y": self.TY,
                }
                callback(**kwargs)

        return self.TY, self.get_registration_parameters()

    def update_transform(self):
        """
        Calculate a new estimate of the deformable transformation.
        See Eq. 22 of https://arxiv.org/pdf/0905.2635.pdf.

        """
        if self.low_rank is False:
            A = np.dot(np.diag(self.P1), self.G) + self.alpha * self.sigma2 * np.eye(
                self.M
            )
            B = self.PX - np.dot(np.diag(self.P1), self.Y)
            self.W = np.linalg.solve(A, B)

        elif self.low_rank is True:
            # Matlab code equivalent can be found here:
            # https://github.com/markeroon/matlab-computer-vision-routines/tree/master/third_party/CoherentPointDrift
            dP = np.diag(self.P1)
            dPQ = np.matmul(dP, self.Q)
            F = self.PX - np.matmul(dP, self.Y)

            self.W = (
                1
                / (self.alpha * self.sigma2)
                * (
                    F
                    - np.matmul(
                        dPQ,
                        (
                            np.linalg.solve(
                                (
                                    self.alpha * self.sigma2 * self.inv_S
                                    + np.matmul(self.Q.T, dPQ)
                                ),
                                (np.matmul(self.Q.T, F)),
                            )
                        ),
                    )
                )
            )
            QtW = np.matmul(self.Q.T, self.W)
            self.E = self.E + self.alpha / 2 * np.trace(
                np.matmul(QtW.T, np.matmul(self.S, QtW))
            )

    @warning_for_keywords()
    def transform_point_cloud(self, *, Y=None):
        """Update a point cloud using the new estimate of the deformable
        transformation.

        Parameters
        ----------
        Y: numpy array, optional
            Array of points to transform - use to predict on new set of points.
            Best for predicting on new points not used to run initial
            registration. If None, self.Y used.

        Returns
        -------
        If Y is None, returns None.
        Otherwise, returns the transformed Y.


        """
        if Y is not None:
            G = gaussian_kernel(X=Y, beta=self.beta, Y=self.Y)
            return Y + np.dot(G, self.W)
        else:
            if self.low_rank is False:
                self.TY = self.Y + np.dot(self.G, self.W)

            elif self.low_rank is True:
                self.TY = self.Y + np.matmul(
                    self.Q, np.matmul(self.S, np.matmul(self.Q.T, self.W))
                )
                return

    def update_variance(self):
        """Update the variance of the mixture model.

        This is using the new estimate of the deformable transformation.
        See the update rule for sigma2 in
        Eq. 23 of of https://arxiv.org/pdf/0905.2635.pdf.

        """
        qprev = self.sigma2

        # The original CPD paper does not explicitly calculate the objective
        # functional. This functional will include terms from both the negative
        # log-likelihood and the Gaussian kernel used for regularization.
        self.q = np.inf

        xPx = np.dot(
            np.transpose(self.Pt1), np.sum(np.multiply(self.X, self.X), axis=1)
        )
        yPy = np.dot(
            np.transpose(self.P1), np.sum(np.multiply(self.TY, self.TY), axis=1)
        )
        trPXY = np.sum(np.multiply(self.TY, self.PX))

        self.sigma2 = (xPx - 2 * trPXY + yPy) / (self.Np * self.D)

        if self.sigma2 <= 0:
            self.sigma2 = self.tolerance / 10

        # Here we use the difference between the current and previous
        # estimate of the variance as a proxy to test for convergence.
        self.diff = np.abs(self.sigma2 - qprev)

    def get_registration_parameters(self):
        """Return the current estimate of the deformable transformation
        parameters.

        Returns
        -------
        self.G: numpy array
            Gaussian kernel matrix.

        self.W: numpy array
            Deformable transformation matrix.
        """
        return self.G, self.W

    def iterate(self):
        """Perform one iteration of the EM algorithm."""
        self.expectation()
        self.maximization()
        self.iteration += 1

    def expectation(self):
        """Compute the expectation step of the EM algorithm."""
        # (M, N)
        P = np.sum((self.X[None, :, :] - self.TY[:, None, :]) ** 2, axis=2)
        P = np.exp(-P / (2 * self.sigma2))
        c = (
            (2 * np.pi * self.sigma2) ** (self.D / 2)
            * self.w
            / (1.0 - self.w)
            * self.M
            / self.N
        )

        den = np.sum(P, axis=0, keepdims=True)  # (1, N)
        den = np.clip(den, np.finfo(self.X.dtype).eps, None) + c

        self.P = np.divide(P, den)
        self.Pt1 = np.sum(self.P, axis=0)
        self.P1 = np.sum(self.P, axis=1)
        self.Np = np.sum(self.P1)
        self.PX = np.matmul(self.P, self.X)

    def maximization(self):
        """Compute the maximization step of the EM algorithm."""
        self.update_transform()
        self.transform_point_cloud()
        self.update_variance()
