import numpy as np
from stats import calculate_mean, calculate_cov


class PCA:
    def __init__(self, mean=None, cov=None):
        """
        Initializes q PCA object. Declares the mean, covariance, transformation matrix and eigen values attributes for
        learning during the training.
        """
        self.mean = None
        self.cov = None
        self.w = None
        self.eigen_values = None

    def calculate_pca(self, x_data):
        """
        Given x_data as the input data calculate a matrix to calculate a PCA matrix.
        This method returns a transformation matrix W.
        Notice that the input dataset x_data is assumed to be having examples along the rows.
        It will be of shape n * d where n is number of examples and d is the dimensionality of the dataset.
        :param x_data: Data using which the transformation matrix is calculated.
        :return: Transformation Matrix W. Each row in W corresponds to an eigen vector.
        """
        # Get the mean of the data
        self.mean = calculate_mean(x_data)
        # Get the covariance of the data
        self.cov = calculate_cov(x_data, self.mean)

        # Find the eigen vectors and values from the covariance matrix
        eigen_values, eigen_vectors = np.linalg.eig(self.cov)
        # Zip together the eigen values and the eigen vectors
        eigen = list(zip(eigen_values, eigen_vectors))
        # Sorting the zipped eigen values and vectors in descending order of eigen values.
        eigen.sort(key=lambda x: x[0], reverse=True)
        # Extract out the sorted eigen vectors in numpy array
        w = np.array(list(map(lambda x: x[1], eigen)))
        # Extract out the sorted eigen values in numpy array
        eigen_values = np.array(list(map(lambda x: x[0], eigen)))
        # saving the results.
        self.w = w
        self.eigen_values = eigen_values
