import numpy as np

class LDA_fs:
    """
    Performs a Linear Discriminant Analysis (LDA)

    Methods
    =======
    fit_transform():
        Fits the model to the data X and Y, derives the transformation matrix W
    and projects the feature matrix X onto the m LDA axes
    """

    def __init__(self, m):
        """
        Parameters
        ==========
        m : int
            Number of LDA axes onto which the data will be projected

        Returns
        =======
        None
        """
        self.m = m

    def fit_transform(self, X, Y):
        """
        Parameters
        ==========
        X : array(n_samples, n_features)
            Feature matrix of the dataset
        Y = array(n_samples)
            Label vector of the dataset

        Returns
        =======
        X_transform : New feature matrix projected onto the m LDA axes

        """

        # Get number of features (columns)
        self.n_features = X.shape[1]
        # Get unique class labels
        class_labels = np.unique(Y)
        # Get the overall mean vector (independent of the class labels)
        mean_overall = np.mean(X, axis=0)  # Mean of each feature
        # Initialize both scatter matrices with zeros
        SW = np.zeros((self.n_features, self.n_features))  # Within scatter matrix
        SB = np.zeros((self.n_features, self.n_features))  # Between scatter matrix

        # Iterate over all classes and select the corresponding data
        for c in class_labels:
            # Filter X for class c
            X_c = X[Y == c]
            # Calculate the mean vector for class c
            mean_c = np.mean(X_c, axis=0)
            # Calculate within-class scatter for class c
            SW += (X_c - mean_c).T.dot((X_c - mean_c))
            # Number of samples in class c
            n_c = X_c.shape[0]
            # Difference between the overall mean and the mean of class c --> between-class scatter
            mean_diff = (mean_c - mean_overall).reshape(self.n_features, 1)
            SB += n_c * (mean_diff).dot(mean_diff.T)

        # Determine SW^-1 * SB
        A = np.linalg.inv(SW).dot(SB)
        # Get the eigenvalues and eigenvectors of (SW^-1 * SB)
        eigenvalues, eigenvectors = np.linalg.eig(A)
        # Keep only the real parts of eigenvalues and eigenvectors
        eigenvalues = np.real(eigenvalues)
        eigenvectors = np.real(eigenvectors.T)

        # Sort the eigenvalues descending (high to low)
        idxs = np.argsort(np.abs(eigenvalues))[::-1]
        self.eigenvalues = np.abs(eigenvalues[idxs])
        self.eigenvectors = eigenvectors[idxs]
        # Store the first m eigenvectors as transformation matrix W
        self.W = self.eigenvectors[0:self.m]

        # Transform the feature matrix X onto LD axes
        return np.dot(X, self.W.T)