from builtins import super
import numpy as np
from scipy.linalg import cho_factor, cho_solve
from scipy.spatial import cKDTree
from scipy.sparse import csr_matrix
from .emregistration import EMRegistration
from .utility import is_positive_semi_definite


class AffineRegistration(EMRegistration):
    """
    Affine registration.

    Attributes
    ----------
    B: numpy array (semi-positive definite)
        DxD affine transformation matrix.

    t: numpy array
        1xD initial translation vector.
    """
    # Additional parameters used in this class, but not inputs.
    # YPY: float
    #     Denominator value used to update the scale factor.
    #     Defined in Fig. 2 and Eq. 8 of https://arxiv.org/pdf/0905.2635.pdf.

    # X_hat: numpy array
    #     Centered target point cloud.
    #     Defined in Fig. 2 of https://arxiv.org/pdf/0905.2635.pdf

    def __init__(self, B=None, t=None, use_kdtree=True, k=10, *args, **kwargs):
        super().__init__(*args, **kwargs)
        
        if B is not None and ((B.ndim != 2) or (B.shape[0] != self.D) or (B.shape[1] != self.D) or not is_positive_semi_definite(B)):
            raise ValueError(
                'The rotation matrix can only be initialized to {}x{} positive semi definite matrices. Instead got: {}.'.format(self.D, self.D, B))

        if t is not None and ((t.ndim != 2) or (t.shape[0] != 1) or (t.shape[1] != self.D)):
            raise ValueError(
                'The translation vector can only be initialized to 1x{} positive semi definite matrices. Instead got: {}.'.format(self.D, t))
        
        self.B = np.eye(self.D) if B is None else B
        self.t = np.atleast_2d(np.zeros((1, self.D))) if t is None else t

        self.YPY = None
        self.X_hat = None
        self.A = None

        # sparse E-step acceleration
        self.use_kdtree = bool(use_kdtree)
        self.k = int(k)
        if self.use_kdtree:
            self.kdtree = cKDTree(self.X)

    def expectation(self):
        if self.use_kdtree:
            distances, indices = self.kdtree.query(self.TY, k=self.k)
            distances = np.clip(distances, np.finfo(self.X.dtype).eps, None)
            P_vals = np.exp(-distances**2 / (2 * self.sigma2))
            rows = np.arange(self.M).repeat(self.k)
            P_sparse = csr_matrix((P_vals.ravel(), (rows, indices.ravel())), shape=(self.M, self.N))
            c_term = (2 * np.pi * self.sigma2)**(self.D / 2) * self.w / (1 - self.w) * self.M / self.N
            den_col = np.array(P_sparse.sum(axis=0)).ravel() + c_term
            den_col = np.clip(den_col, np.finfo(self.X.dtype).eps, None)
            inv_den_col = 1.0 / den_col
            self.P = P_sparse.multiply(inv_den_col[np.newaxis, :])
            self.Pt1 = np.array(self.P.sum(axis=0)).ravel()
            self.P1 = np.array(self.P.sum(axis=1)).ravel()
            self.Np = float(self.P1.sum())
            self.PX = self.P @ self.X
        else:
            super().expectation()

    def update_transform(self):
        """
        Calculate a new estimate of the rigid transformation.

        """

        # source and target point cloud means
        muX = np.divide(np.sum(self.PX, axis=0), self.Np)
        muY = np.divide(np.sum((self.P.T @ self.Y), axis=0), self.Np)

        self.X_hat = self.X - np.tile(muX, (self.N, 1))
        Y_hat = self.Y - np.tile(muY, (self.M, 1))

        self.A = (self.X_hat.T @ (self.P.T @ Y_hat))

        # YPY = Y_hat^T diag(P1) Y_hat  -> use broadcasting
        YPY = (self.P1[:, None] * Y_hat).T @ Y_hat
        self.YPY = YPY

        # Solve for B with Cholesky (add small jitter for stability)
        jitter = 1e-8
        c, low = cho_factor(self.YPY + jitter * np.eye(self.D))
        self.B = cho_solve((c, low), self.A.T).T
        self.t = (muX - (self.B.T @ muY.T).T)

    def transform_point_cloud(self, Y=None):
        """
        Update a point cloud using the new estimate of the affine transformation.
        
        Attributes
        ----------
        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 None:
            self.TY = np.dot(self.Y, self.B) + np.tile(self.t, (self.M, 1))
            return
        else:
            return np.dot(Y, self.B) + np.tile(self.t, (Y.shape[0], 1))

    def update_variance(self):
        """
        Update the variance of the mixture model using the new estimate of the affine transformation.
        See the update rule for sigma2 in Fig. 3 of of https://arxiv.org/pdf/0905.2635.pdf.

        """
        qprev = self.q

        trAB = np.trace(np.dot(self.A, self.B))
        xPx = np.dot(np.transpose(self.Pt1), np.sum(
            np.multiply(self.X_hat, self.X_hat), axis=1))
        trBYPYP = np.trace(np.dot(np.dot(self.B, self.YPY), self.B))
        self.q = (xPx - 2 * trAB + trBYPYP) / (2 * self.sigma2) + \
            self.D * self.Np/2 * np.log(self.sigma2)
        self.diff = np.abs(self.q - qprev)

        self.sigma2 = (xPx - trAB) / (self.Np * self.D)

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

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

        Returns
        -------
        B: numpy array
            DxD affine transformation matrix.

        t: numpy array
            1xD translation vector.

        """
        return self.B, self.t
