# Importing the required libraries
from math import log

import numpy as np
from scipy import linalg


def make_sample_data(
    n_samples, n_features, cov_class_1, cov_class_2, shift=None, seed=0
):
    """Generates sample data for testing LDA and QDA models.

    Args:
        n_samples: (int) - Number of samples per class.
        n_features: (int) - Number of features.
        cov_class_1: (np.ndarray) - Covariance matrix for class 1.
        cov_class_2: (np.ndarray) - Covariance matrix for class 2.
        shift: (list), optional - Shift applied to class 2 data (default is [1, 1]).
        seed: (int), optional - Random seed for reproducibility (default is 0).

    Returns:
        X: (np.ndarray) - Generated feature data.
        y: (np.ndarray) - Generated target labels.
    """
    if shift is None:
        shift = [1, 1]
    rng = np.random.RandomState(seed)
    X = np.concatenate(
        [
            # Data points for class 1, generated by multiplying a random matrix with the covariance matrix of class 1
            rng.randn(n_samples, n_features) @ cov_class_1,
            # Data points for class 2, generated by multiplying a random matrix with the covariance matrix of class 2 and adding [1, 1]
            rng.randn(n_samples, n_features) @ cov_class_2 + np.array(shift),
        ]
    )
    y = np.concatenate([np.zeros(n_samples), np.ones(n_samples)])
    return X, y


def _validate_data(X, y):
    """Validates input data.

    Args:
        X : array-like of shape (n_samples, n_features): Training data.
        y : array-like of shape (n_samples,) or (n_samples, n_targets): Target values.

    Must be:
        - array-like
        - same number of samples
        - Not empty
    """
    if not isinstance(X, np.ndarray) or not isinstance(y, np.ndarray):
        raise ValueError("X and y must be numpy arrays")

    if X.shape[0] != y.shape[0]:
        raise ValueError("X and y must have the same number of samples")

    if X.shape[0] == 0 or y.shape[0] == 0:
        raise ValueError("X and y must not be empty")

    if np.any(np.isnan(X)) or np.any(np.isnan(y)):
        raise ValueError("X and y must not contain NaN values")

    if np.any(np.isinf(X)) or np.any(np.isinf(y)):
        raise ValueError("X and y must not contain infinite values")


class LinearDiscriminantAnalysis:
    """Implements Linear Discriminant Analysis.

    A classifier with a linear decision boundary, generated by fitting class conditional densities to the data and using Bayes' rule.

    Args:
        solver: (str) - {'svd', 'lsqr', 'eigen'}, default='svd'. Solver to use for the LDA.
        priors: (np.ndarray), optional - Prior probabilities of the classes (default is None).
    """

    def __init__(self, solver="svd", priors=None):
        """Initializes the Linear Discriminant Analysis model.

        Args:
            solver: (str) - {'svd', 'lsqr', 'eigen'}, default='svd'. Solver to use for the LDA.
            priors: (np.ndarray), optional - Prior probabilities of the classes (default is None).
        """
        self.solver = solver
        self.priors = priors

    def fit(self, X, y):
        """Fits the LDA model to the training data.

        Args:
            X: (np.ndarray) - Training feature data.
            y: (np.ndarray) - Training target data.
        """
        _validate_data(X, y)

        self.classes_ = np.unique(y)  # Unique classes, i.e., distinct class labels
        self.means_ = {}  # Mean of each feature per class
        self.covariance_ = np.cov(X, rowvar=False)  # Covariance matrix of all features
        if self.priors is None:  # Prior probabilities of each class
            self.priors_ = {}
        else:
            self.priors_ = self.priors

        # Compute mean and prior for each class, and covariance
        for cls in self.classes_:
            X_cls = X[y == cls]  # Data points corresponding to class cls
            self.means_[cls] = np.mean(X_cls, axis=0)  # Mean of each feature per class
            self.priors_[cls] = (
                X_cls.shape[0] / X.shape[0]
            )  # Prior probability of class cls

        if self.solver == "svd":
            self._fit_svd(X, y)
        elif self.solver == "lsqr":
            self._fit_lsqr(X, y)
        elif self.solver == "eigen":
            self._fit_eigen(X, y)
        else:
            raise ValueError(f"Solver '{self.solver}' is not yet supported.")

    def _fit_svd(self, X, y):
        """Fits the LDA model using Singular Value Decomposition (SVD).

        Args:
            X: (np.ndarray) - Training feature data.
            y: (np.ndarray) - Training target data.
        """
        X_centered = X - np.mean(X, axis=0)  # Center the data
        U, S, Vt = np.linalg.svd(
            X_centered, full_matrices=False
        )  # Perform SVD, U and Vt are the left and right singular vectors, S is the singular values
        rank = np.sum(
            S > 1e-10
        )  # Compute the rank of the matrix, i.e., the number of singular values greater than 1e-10

        # Select only the top components
        U = U[:, :rank]
        S = S[:rank]
        Vt = Vt[:rank, :]

        self.scalings_ = (
            Vt.T / S
        )  # Compute the transformation matrix, i.e., the scalings
        self.means_ = {
            cls: mean @ self.scalings_ for cls, mean in self.means_.items()
        }  # Transform the means, i.e., the mean of each feature per class
        self.covariance_ = np.diag(
            1 / S**2
        )  # Transform the covariance matrix, i.e., the inverse of the singular values squared

    def _fit_lsqr(self, X, y):
        """Fits the LDA model using Least Squares (LSQR).

        Args:
            X: (np.ndarray) - Training feature data.
            y: (np.ndarray) - Training target data.
        """
        # Create matrix for class-specific means
        _mean_matrix = np.vstack([self.means_[cls] for cls in self.classes_])

        # Solve least squares problem
        X_centered = X - np.mean(X, axis=0)
        Y = np.zeros((X.shape[0], len(self.classes_)))
        for i, cls in enumerate(self.classes_):
            Y[:, i] = (y == cls).astype(float)

        # Solve the normal equations using least squares
        coef, residuals, rank, singular_values = linalg.lstsq(X_centered, Y)

        # Compute scalings
        self.scalings_ = coef

        # Transform class means and store them back in the dictionary
        self.means_ = {
            cls: (self.means_[cls] @ self.scalings_) for cls in self.classes_
        }

        # Update the covariance matrix to match the transformed data
        self.covariance_ = np.cov(X_centered @ self.scalings_, rowvar=False)

    def _fit_eigen(self, X, y):
        """Fits the LDA model using eigenvalue decomposition.

        Args:
            X: (np.ndarray) - Training feature data.
            y: (np.ndarray) - Training target data.
        """
        # Center the data by subtracting the global mean
        _X_centered = X - np.mean(X, axis=0)

        # Compute the total scatter matrix (within-class + between-class)
        # We'll compute this by first computing the within-class scatter matrix
        Sw = np.zeros((X.shape[1], X.shape[1]))
        Sb = np.zeros((X.shape[1], X.shape[1]))

        for cls in self.classes_:
            X_cls = X[y == cls]
            X_cls_centered = X_cls - self.means_[cls]

            # Within-class scatter matrix
            Sw += X_cls_centered.T @ X_cls_centered

            # Between-class scatter matrix
            cls_mean_centered = self.means_[cls] - np.mean(X, axis=0)
            Sb += len(X_cls) * (
                cls_mean_centered[:, np.newaxis] @ cls_mean_centered[np.newaxis, :]
            )

        # Solve the generalized eigenvalue problem: Sb @ w = λ * Sw @ w
        # This requires inverting the within-class scatter matrix
        eigenvalues, eigenvectors = linalg.eigh(np.linalg.inv(Sw) @ Sb)

        # Sort eigenvalues and eigenvectors in descending order
        idx = np.argsort(eigenvalues)[::-1]
        eigenvalues = eigenvalues[idx]
        eigenvectors = eigenvectors[:, idx]

        # Select top eigenvectors
        n_components = len(self.classes_)
        self.scalings_ = eigenvectors[:, :n_components]

        # Transform class means
        self.means_ = {
            cls: (self.means_[cls] @ self.scalings_) for cls in self.classes_
        }

        # Transform the covariance matrix
        # Use the inverse of the eigenvalues as the diagonal of the transformed covariance
        self.covariance_ = np.diag(1.0 / (eigenvalues[:n_components] + 1e-11))

    def predict(self, X):
        """Predicts class labels for the input data.

        Args:
            X: (np.ndarray) - Test feature data.

        Returns:
            predictions: (np.ndarray) - Predicted class labels.
        """
        scores = self.decision_function(X)  # Compute the decision function
        return self.classes_[
            np.argmax(scores, axis=1)
        ]  # Return the class with the highest score

    def decision_function(self, X):
        """Computes the log-likelihood of each class for the input data. The decision function is the log-likelihood of each class.

        Args:
            X: (np.ndarray) - Test feature data.

        Returns:
            scores: (np.ndarray) - Log-likelihood of each class for the input samples.
        """
        # Cache the inverse of the covariance matrix
        inv_covariance = np.linalg.inv(self.covariance_)
        scores = []
        # Compute log-likelihood of each class
        for cls in self.classes_:
            mean = self.means_[cls]  # Mean of each feature per class
            prior = self.priors_[cls]  # Prior probability of class cls

            # Score is the log-likelihood of each class
            score = (
                X @ inv_covariance @ mean.T
                - 0.5 * mean @ inv_covariance @ mean.T
                + log(prior)
            )

            scores.append(score)  # Append the score to the list of scores
        return np.array(
            scores
        ).T  # Return the scores as a numpy array, with each row corresponding to a sample and each column corresponding to a class


class QuadraticDiscriminantAnalysis:
    """Implements Quadratic Discriminant Analysis.

    The quadratic term allows for more flexibility in modeling the class conditional
    A classifier with a quadratic decision boundary, generated by fitting class conditional densities to the data and using Bayes' rule.

    Args:
        priors: (np.ndarray), optional - Prior probabilities of the classes (default is None).
        reg_param: (float), optional - Regularization parameter (default is 0.0).
    """

    def __init__(self, priors=None, reg_param=0.0):
        """Initialize the Quadratic Discriminant Analysis model with the specified prior probabilities and regularization parameter.

        Args:
            priors: (np.ndarray), optional - Prior probabilities of the classes (default is None).
            reg_param: (float), optional - Regularization parameter (default is 0.0).
        """
        self.priors = priors

        assert reg_param >= 0.0, "Regularization parameter must be non-negative."
        self.reg_param = reg_param

    def fit(self, X, y):
        """Fit the model according to the given training data. Uses the means and covariance matrices of each class.

        Args:
            X: (np.ndarray) - Training feature data.
            y: (np.ndarray) - Training target data.
        """
        _validate_data(X, y)

        self.classes_ = np.unique(y)  # Unique classes, i.e., distinct class labels
        self.means_ = {}  # Mean of each feature per class
        self.covariances_ = {}  # Covariance matrix of each class
        if self.priors is None:  # Prior probabilities of each class
            self.priors_ = {}
        else:
            self.priors_ = self.priors

        # Compute mean and prior for each class, and covariance
        for cls in self.classes_:
            X_cls = X[y == cls]  # Data points corresponding to class cls
            self.means_[cls] = np.mean(X_cls, axis=0)  # Mean of each feature per class
            self.priors_[cls] = (
                X_cls.shape[0] / X.shape[0]
            )  # Prior probability of class cls

            cov = np.cov(X_cls, rowvar=False)  # Covariance matrix of all features

            if self.reg_param > 0.0:
                self.covariances_[cls] = (
                    cov + np.eye(cov.shape[0]) * self.reg_param
                )  # Add regularization term to diagonal
            else:
                self.covariances_[cls] = (
                    cov + np.eye(cov.shape[0]) * 1e-6
                )  # Add small value to diagonal

    def predict(self, X):
        """Perform classification on an array of test vectors X.

        Args:
            X: (np.ndarray) - Test feature data.

        Returns:
            predictions: (np.ndarray) - Predicted class labels.
        """
        scores = self.decision_function(X)
        return self.classes_[np.argmax(scores, axis=1)]

    def decision_function(self, X):
        """Apply decision function to an array of samples.

        The decision function is the log-likelihood of each class.

        Args:
            X: (np.ndarray) - Test feature data.

        Returns:
            scores: (np.ndarray) - Log-likelihood of each class for the input samples.
        """
        scores = []
        # Compute log-likelihood of each class
        for cls in self.classes_:
            mean = self.means_[cls]  # Mean of each feature per class
            cov = self.covariances_[cls]  # Covariance matrix of each class
            prior = self.priors_[cls]  # Prior probability of class cls

            # Compute the inverse of the covariance matrix and the log-determinant of the covariance matrix
            inv_cov = np.linalg.inv(cov)
            log_det_cov = np.log(np.linalg.det(cov))

            # Compute the log-likelihood
            score = -0.5 * np.sum(
                (X - mean) @ inv_cov * (X - mean), axis=1
            )  # Quadratic term
            score -= 0.5 * log_det_cov  # Log-determinant term
            score += log(prior)  # Prior term

            scores.append(score)  # Append the score to the list of scores
        return np.array(
            scores
        ).T  # Return the scores as a numpy array, with each row corresponding to a sample and each column corresponding to a class


class Perceptron:
    """Implements the Perceptron algorithm for binary and multiclass classification.

    Args:
        max_iter: (int) - Maximum number of iterations (default is 1000).
        learning_rate: (float) - Learning rate for weight updates (default is 0.01).
    """

    def __init__(self, max_iter=1000, learning_rate=0.01):
        """Initializes the classifier with the specified maximum number of iterations and learning rate.

        Args:
            max_iter (int, optional): The maximum number of iterations for the training process. Defaults to 1000.
            learning_rate (float, optional): The learning rate for the optimization algorithm. Defaults to 0.01.
        """
        self.max_iter = max_iter
        self.learning_rate = learning_rate
        self.weights = None
        self.bias = None
        self.classes_ = None

    def fit(self, X, y):
        """Fits the Perceptron model to the training data.

        Args:
            X: (np.ndarray) - Training feature data of shape (n_samples, n_features).
            y: (np.ndarray) - Training target data of shape (n_samples,).
        """
        _validate_data(X, y)

        self.classes_ = np.unique(y)
        n_samples, n_features = X.shape

        if len(self.classes_) == 2:
            # Binary classification
            self.weights = np.zeros(n_features)
            self.bias = 0

            # Convert labels to {1, -1}
            y_ = np.where(y == self.classes_[1], 1, -1)

            for _ in range(self.max_iter):
                for idx, x_i in enumerate(X):
                    if y_[idx] * (np.dot(x_i, self.weights) + self.bias) <= 0:
                        self.weights += self.learning_rate * y_[idx] * x_i
                        self.bias += self.learning_rate * y_[idx]
        else:
            # Multiclass classification (One-vs-Rest)
            self.weights = {}
            self.bias = {}

            for cls in self.classes_:
                # Create binary labels for the current class
                y_binary = np.where(y == cls, 1, -1)
                weights = np.zeros(n_features)
                bias = 0

                for _ in range(self.max_iter):
                    for idx, x_i in enumerate(X):
                        if y_binary[idx] * (np.dot(x_i, weights) + bias) <= 0:
                            weights += self.learning_rate * y_binary[idx] * x_i
                            bias += self.learning_rate * y_binary[idx]

                self.weights[cls] = weights
                self.bias[cls] = bias

    def predict(self, X):
        """Predicts class labels for the input data.

        Args:
            X: (np.ndarray) - Test feature data.

        Returns:
            predictions: (np.ndarray) - Predicted class labels.
        """
        if len(self.classes_) == 2:
            # Binary classification
            linear_output = np.dot(X, self.weights) + self.bias
            return np.where(linear_output > 0, self.classes_[1], self.classes_[0])
        else:
            # Multiclass classification
            scores = np.zeros((X.shape[0], len(self.classes_)))
            for idx, cls in enumerate(self.classes_):
                scores[:, idx] = np.dot(X, self.weights[cls]) + self.bias[cls]
            return self.classes_[np.argmax(scores, axis=1)]


class LogisticRegression:
    """Implements Logistic Regression using gradient descent. Supports binary and multiclass classification.

    Args:
        learning_rate: (float) - Learning rate for gradient updates (default is 0.01).
        max_iter: (int) - Maximum number of iterations (default is 1000).
    """

    valid_params = ["learning_rate", "max_iter"]

    def __init__(self, learning_rate=0.01, max_iter=1000):
        """Initializes the classifier with specified hyperparameters.

        Args:
            learning_rate (float, optional): The step size for updating weights during training. Defaults to 0.01.
            max_iter (int, optional): The maximum number of iterations for the training process. Defaults to 1000.
        """
        self.learning_rate = learning_rate
        self.max_iter = max_iter
        self.weights = None
        self.bias = None
        self.classes_ = None

    def fit(self, X, y):
        """Fits the Logistic Regression model to the training data.

        Args:
            X: (np.ndarray) - Training feature data of shape (n_samples, n_features).
            y: (np.ndarray) - Training target data of shape (n_samples,).
        """
        _validate_data(X, y)

        self.classes_ = np.unique(y)
        n_samples, n_features = X.shape

        if len(self.classes_) == 2:
            # Binary classification
            self.weights = np.zeros(n_features)
            self.bias = 0

            for _ in range(self.max_iter):
                linear_model = np.dot(X, self.weights) + self.bias
                y_predicted = self._sigmoid(linear_model)

                # Compute gradients
                dw = (1 / n_samples) * np.dot(X.T, (y_predicted - y))
                db = (1 / n_samples) * np.sum(y_predicted - y)

                # Update weights and bias
                self.weights -= self.learning_rate * dw
                self.bias -= self.learning_rate * db
        else:
            # Multiclass classification (One-vs-Rest)
            self.weights = {}
            self.bias = {}

            for cls in self.classes_:
                # Create binary labels for the current class
                y_binary = np.where(y == cls, 1, 0)
                weights = np.zeros(n_features)
                bias = 0

                for _ in range(self.max_iter):
                    linear_model = np.dot(X, weights) + bias
                    y_predicted = self._sigmoid(linear_model)

                    # Compute gradients
                    dw = (1 / n_samples) * np.dot(X.T, (y_predicted - y_binary))
                    db = (1 / n_samples) * np.sum(y_predicted - y_binary)

                    # Update weights and bias
                    weights -= self.learning_rate * dw
                    bias -= self.learning_rate * db

                self.weights[cls] = weights
                self.bias[cls] = bias

    def predict(self, X):
        """Predicts class labels for the input data.

        Args:
            X: (np.ndarray) - Test feature data.

        Returns:
            predictions: (np.ndarray) - Predicted class labels.
        """
        if len(self.classes_) == 2:
            # Binary classification
            linear_model = np.dot(X, self.weights) + self.bias
            y_predicted = self._sigmoid(linear_model)
            return np.where(y_predicted > 0.5, self.classes_[1], self.classes_[0])
        else:
            # Multiclass classification
            scores = np.zeros((X.shape[0], len(self.classes_)))
            for idx, cls in enumerate(self.classes_):
                linear_model = np.dot(X, self.weights[cls]) + self.bias[cls]
                scores[:, idx] = self._sigmoid(linear_model)
            return self.classes_[np.argmax(scores, axis=1)]

    def _sigmoid(self, z):
        """Applies the sigmoid function."""
        return 1 / (1 + np.exp(-z))

    # TODO: Add get_params and set_params to all estimators
    def get_params(self, **params):
        """Get the parameters of this logistic regression model.

        Args:
            **params: dict, Parameter names mapped to their values.

        Returns:
            params: dict, Parameter names mapped to their values.
        """
        params = {"learning_rate": self.learning_rate, "max_iter": self.max_iter}
        return params

    def set_params(self, **params):
        """Set the parameters of this logistic regression model.

        Args:
            **params: dict, Parameter names mapped to their values.

        Returns:
            self: LogisticRegression instance.
        """
        if not isinstance(params, dict):
            raise TypeError("params must be a dict")
        for key in params:
            if key not in self.valid_params:
                raise ValueError(f"Invalid parameter: {key}")

        if "learning_rate" in params:
            self.learning_rate = params["learning_rate"]

        if "max_iter" in params:
            self.max_iter = params["max_iter"]

        return self
