import numpy as np


def Ridge_parameters(X: np.ndarray, y: np.ndarray, lam: float) -> np.ndarray:
    """Computes the Ridge regression parameters.
    Args:
        X: The input data matrix of shape (n_samples, n_features).
        y: The target vector of shape (n_samples,).
        lam: The regularization parameter (lambda).
    Returns:
        The Ridge regression parameters of shape (n_features,).
    """
    # Assumes X is scaled and has no intercept column
    return np.linalg.inv(X.T @ X + lam * np.eye(X.shape[1])) @ X.T @ y


def OLS_parameters(X: np.ndarray, y: np.ndarray) -> np.ndarray:
    """Computes the Ordinary Least Squares (OLS) regression parameters.
    Args:
        X: The input data matrix of shape (n_samples, n_features).
        y: The target vector of shape (n_samples,).
    Returns:
        The OLS regression parameters of shape (n_features,).
    """
    return Ridge_parameters(X, y, lam=0.0)  # OLS is Ridge with lambda=0


class GradientDescent:
    def __init__(
        self, *args, learning_rate: float = 0.1, num_iterations: int = 1000, **kwargs
    ):
        """Gradient Descent Class

        Args:
            learning_rate (float, optional): Learning rate used for step updates. Defaults to 0.1.
            num_iterations (int, optional): Number of iterations for gradient descent. Defaults to 1000.
        """
        self._cost_history = np.zeros(num_iterations)
        self.learning_rate = learning_rate
        self.num_iterations = num_iterations

    @property
    def cost_history(self) -> np.ndarray:
        """Returns the cost history of the optimization process.

        Returns:
            np.ndarray: Array of cost values for each iteration.
        """
        return self._cost_history

    def get_epochs(self) -> np.ndarray:
        """Returns an array of epoch numbers from 0 to num_iterations - 1.

        Returns:
            np.ndarray: Array of epoch numbers.
        """
        return np.arange(self.num_iterations)

    def fit(self, X: np.ndarray, y: np.ndarray) -> np.ndarray:
        """Returns the optimal solution for an optimization problem using the gradient descent

        Args:
            X (np.ndarray): X values
            y (np.ndarray): y values

        Returns:
            np.ndarray: Optimal parameters
        """
        self.X = X
        self.y = y
        self.theta = np.zeros(X.shape[1])
        self._precomp()
        for t in range(self.num_iterations):
            self._comp_step()
            self._cost_history[t] = self._compute_cost()
            self._update_theta()
        return self.theta

    def _precomp(self):
        pass

    def _comp_step(self):
        pass

    def _compute_cost(self):
        pass

    def _update_theta(self):
        pass


class OLSGradientDescent(GradientDescent):
    def _precomp(self):
        self.XTX = self.X.T @ self.X
        self.XTy = self.X.T @ self.y
        self.n = len(self.y)

    def _comp_step(self):
        self.err = self.X @ self.theta - self.y

    def _compute_cost(self) -> float:
        return 0.5 * np.mean(np.square(self.err))

    def _compute_grad(self) -> np.ndarray:
        return (self.XTX @ self.theta - self.XTy) / self.n

    def _update_theta(self):
        self.theta -= self.learning_rate * self._compute_grad()


class RidgeGradientDescent(OLSGradientDescent):
    def __init__(self, *args, lam: float = 0.1, **kwargs):
        super().__init__(*args, **kwargs)
        self.lam = lam

    def _compute_cost(self) -> float:
        return super()._compute_cost() + 0.5 * self.lam * np.sum(np.square(self.theta))

    def _compute_grad(self) -> np.ndarray:
        return super()._compute_grad() + self.lam * self.theta


class OLSMomentum(OLSGradientDescent):
    def __init__(self, *args, delta: float = 1.0, **kwargs):
        super().__init__(*args, **kwargs)
        self.delta = delta

    def _precomp(self):
        self.last_theta = np.zeros_like(self.theta)
        return super()._precomp()

    def _update_theta(self):
        v = self.delta * (
            self.theta - self.last_theta - self.learning_rate * self._compute_grad()
        )
        self.last_theta = self.theta.copy()
        self.theta += v


class RidgeMomentum(OLSMomentum, RidgeGradientDescent):
    def __init__(self, *args, **kwargs):
        super().__init__(*args, **kwargs)


class OLSAdaGrad(OLSGradientDescent):
    def __init__(self, *args, learning_rate=0.1, num_iterations=1000, **kwargs):
        super().__init__(
            *args, learning_rate=learning_rate, num_iterations=num_iterations, **kwargs
        )

    def _precomp(self):
        self.quad_sum = np.zeros_like(self.theta)
        return super()._precomp()

    def _update_theta(self):
        grad = self._compute_grad()
        self.quad_sum += np.square(grad)
        self.theta -= (
            self.learning_rate * grad / (np.sqrt(self.quad_sum) + 1e-10)
        )  # to avoid division by zero


class RidgeAdaGrad(OLSAdaGrad, RidgeGradientDescent):
    def __init__(self, *args, **kwargs):
        super().__init__(*args, **kwargs)


class OLSRMSProp(OLSGradientDescent):
    def __init__(
        self, *args, learning_rate=0.1, num_iterations=1000, gamma=0.9, **kwargs
    ):
        super().__init__(
            *args, learning_rate=learning_rate, num_iterations=num_iterations, **kwargs
        )
        self.gamma = gamma

    def _precomp(self):
        self.prev_v = np.zeros_like(self.theta)
        return super()._precomp()

    def _update_theta(self):
        grad = self._compute_grad()
        v = self.gamma * self.prev_v + (1 - self.gamma) * np.square(grad)
        self.theta -= (
            self.learning_rate * grad / (np.sqrt(v) + 1e-10)
        )  # to avoid division by zero
        self.prev_v = v


class RidgeRMSProp(OLSRMSProp, RidgeGradientDescent):
    def __init__(self, *args, **kwargs):
        super().__init__(*args, **kwargs)


class OLSAdam(OLSGradientDescent):
    def __init__(
        self,
        *args,
        learning_rate=0.1,
        num_iterations=1000,
        beta1=0.9,
        beta2=0.999,
        **kwargs,
    ):
        super().__init__(
            *args, learning_rate=learning_rate, num_iterations=num_iterations, **kwargs
        )

        self.beta1 = beta1
        self.beta2 = beta2

    def _precomp(self):
        self.m = np.zeros_like(self.theta)
        self.v = np.zeros_like(self.theta)
        self.current_iteration = 0
        return super()._precomp()

    def _update_theta(self):
        self.current_iteration += 1
        grad = self._compute_grad()
        self.m = self.beta1 * self.m + (1 - self.beta1) * grad
        self.v = self.beta2 * self.v + (1 - self.beta2) * np.square(grad)

        m_hat = self.m / (1 - self.beta1 ** (self.current_iteration))
        v_hat = self.v / (1 - self.beta2 ** (self.current_iteration))

        self.theta -= (
            self.learning_rate * m_hat / (np.sqrt(v_hat) + 1e-10)
        )  # to avoid division by zero


class RidgeAdam(OLSAdam, RidgeGradientDescent):
    def __init__(self, *args, **kwargs):
        super().__init__(*args, **kwargs)


class LASSOGradientDescent(OLSGradientDescent):
    def __init__(self, *args, lam: float = 0.1, **kwargs):
        super().__init__(*args, **kwargs)
        self.lam = lam

    def _compute_cost(self) -> float:
        return super()._compute_cost() + self.lam * np.sum(np.abs(self.theta))

    def _compute_grad(self) -> np.ndarray:
        return super()._compute_grad() + self.lam * np.sign(self.theta)


class LASSOMomentum(OLSMomentum, LASSOGradientDescent):
    def __init__(self, *args, **kwargs):
        super().__init__(*args, **kwargs)


class LASSOAdaGrad(OLSAdaGrad, LASSOGradientDescent):
    def __init__(self, *args, **kwargs):
        super().__init__(*args, **kwargs)


class LASSORMSProp(OLSRMSProp, LASSOGradientDescent):
    def __init__(self, *args, **kwargs):
        super().__init__(*args, **kwargs)


class LASSOAdam(OLSAdam, LASSOGradientDescent):
    def __init__(self, *args, **kwargs):
        super().__init__(*args, **kwargs)


class OLSStochasticGradientDescent(OLSGradientDescent):
    def __init__(
        self, *args, batch_size: int = 100, batches_per_epoch: int = 1, **kwargs
    ):
        # print(self.__class__.__name__) # If you see this: debugging yaaaay, the programmer that wrote this line is stupid...
        super().__init__(*args, **kwargs)
        self.batch_size = batch_size
        self.batches_per_epoch = batches_per_epoch

    def _precomp(self):
        self.N = len(self.y)
        self.indices = np.arange(self.N)
        self.n = self.batch_size
        np.random.shuffle(self.indices)
        self.X = self.X[self.indices]
        self.y = self.y[self.indices]

    def _comp_step(self):
        index = np.random.randint(0, self.N)
        batch_indices = slice(index, index + self.batch_size)
        if index + self.batch_size > self.N:
            batch_indices = slice(index, self.N)

        X_batch = self.X[batch_indices]
        y_batch = self.y[batch_indices]
        self.err = X_batch @ self.theta - y_batch
        self.XTX = X_batch.T @ X_batch
        self.XTy = X_batch.T @ y_batch

    def get_epochs(self):
        return np.arange(self.num_iterations // self.batches_per_epoch)

    @property
    def cost_history(self) -> np.ndarray:
        """Returns the cost history of the optimization process.

        Returns:
            np.ndarray: Array of cost values for each epoch.
        """
        return np.mean(self._cost_history.reshape(-1, self.batches_per_epoch), axis=1)


class RidgeStochasticGradientDescent(
    OLSStochasticGradientDescent, RidgeGradientDescent
):
    def __init__(self, *args, **kwargs):
        super().__init__(*args, **kwargs)


class LASSOStochasticGradientDescent(
    OLSStochasticGradientDescent, LASSOGradientDescent
):
    def __init__(self, *args, **kwargs):
        super().__init__(*args, **kwargs)


class OLSMomentumSGD(OLSMomentum, OLSStochasticGradientDescent):
    def __init__(self, *args, **kwargs):
        super().__init__(*args, **kwargs)


class RidgeMomentumSGD(RidgeMomentum, RidgeStochasticGradientDescent):
    def __init__(self, *args, **kwargs):
        super().__init__(*args, **kwargs)


class LASSOMomentumSGD(LASSOMomentum, LASSOStochasticGradientDescent):
    def __init__(self, *args, **kwargs):
        super().__init__(*args, **kwargs)


class OLSAdaGradSGD(OLSAdaGrad, OLSStochasticGradientDescent):
    def __init__(self, *args, **kwargs):
        super().__init__(*args, **kwargs)


class RidgeAdaGradSGD(RidgeAdaGrad, RidgeStochasticGradientDescent):
    def __init__(self, *args, **kwargs):
        super().__init__(*args, **kwargs)


class LASSOAdaGradSGD(LASSOAdaGrad, LASSOStochasticGradientDescent):
    def __init__(self, *args, **kwargs):
        super().__init__(*args, **kwargs)


class OLSRMSPropSGD(OLSRMSProp, OLSStochasticGradientDescent):
    def __init__(self, *args, **kwargs):
        super().__init__(*args, **kwargs)


class RidgeRMSPropSGD(RidgeRMSProp, RidgeStochasticGradientDescent):
    def __init__(self, *args, **kwargs):
        super().__init__(*args, **kwargs)


class LASSORMSPropSGD(LASSORMSProp, LASSOStochasticGradientDescent):
    def __init__(self, *args, **kwargs):
        super().__init__(*args, **kwargs)


class OLSAdamSGD(OLSAdam, OLSStochasticGradientDescent):
    def __init__(self, *args, **kwargs):
        super().__init__(*args, **kwargs)


class RidgeAdamSGD(RidgeAdam, RidgeStochasticGradientDescent):
    def __init__(self, *args, **kwargs):
        super().__init__(*args, **kwargs)


class LASSOAdamSGD(LASSOAdam, LASSOStochasticGradientDescent):
    def __init__(self, *args, **kwargs):
        super().__init__(*args, **kwargs)


OLS_GD_OPTIMIZERS = [OLSGradientDescent, OLSMomentum, OLSAdaGrad, OLSRMSProp, OLSAdam]
RIDGE_GD_OPTIMIZERS = [
    RidgeGradientDescent,
    RidgeMomentum,
    RidgeAdaGrad,
    RidgeRMSProp,
    RidgeAdam,
]
LASSO_GD_OPTIMIZERS = [
    LASSOGradientDescent,
    LASSOMomentum,
    LASSOAdaGrad,
    LASSORMSProp,
    LASSOAdam,
]

OLS_SGD_OPTIMIZERS = [
    OLSStochasticGradientDescent,
    OLSMomentumSGD,
    OLSAdaGradSGD,
    OLSRMSPropSGD,
    OLSAdamSGD,
]
RIDGE_SGD_OPTIMIZERS = [
    RidgeStochasticGradientDescent,
    RidgeMomentumSGD,
    RidgeAdaGradSGD,
    RidgeRMSPropSGD,
    RidgeAdamSGD,
]
LASSO_SGD_OPTIMIZERS = [
    LASSOStochasticGradientDescent,
    LASSOMomentumSGD,
    LASSOAdaGradSGD,
    LASSORMSPropSGD,
    LASSOAdamSGD,
]