from typing import Optional

import numpy as np
from scipy.sparse.linalg import LinearOperator, eigsh
from scipy.linalg import eigh

from vicentin.kernels import Kernel, LinearKernel
from vicentin.kernels.kernel import CenteredKernel
from vicentin.utils.logging import info, debug


def _pca_primal(kernel, n_components, max_k, trace_K, var):
    X = kernel.X
    D = X.shape[1]
    debug(f"Computing Primal PCA on feature space D={D}")

    X_mean = X.mean(axis=0)
    X_centered = X - X_mean

    C = X_centered.T @ X_centered

    if var is not None:
        calc_k = D
    elif n_components is not None:
        calc_k = min(n_components, D)
    else:
        calc_k = min(max_k, D)

    eigenvalues, V = eigh(C, subset_by_index=[D - calc_k, D - 1])

    idx = np.argsort(eigenvalues)[::-1]
    eigenvalues = eigenvalues[idx]
    V = V[:, idx]

    if var is not None:
        cum_var_ratio = np.cumsum(eigenvalues) / trace_K
        mask = cum_var_ratio <= var
        n_to_keep = np.sum(mask) + 1
        debug(f"Retaining {n_to_keep} components for {var:.2f} variance.")

        if n_components is not None:
            n_to_keep = min(n_to_keep, n_components)

        eigenvalues = eigenvalues[:n_to_keep]
        V = V[:, :n_to_keep]

    alpha = X_centered @ V / np.maximum(eigenvalues, 1e-12)

    return eigenvalues, alpha


def _pca_kernel(kernel, n_components, max_k, trace_K, var):
    if kernel.is_dense:
        debug("Computing Kernel PCA with dense eigendecomposition (eigh)")
        K_dense = kernel(kernel.X)
        N = kernel.N

        if var is not None:
            calc_k = N
        elif n_components is not None:
            calc_k = min(n_components, N)
        else:
            calc_k = max_k

        eigenvalues, eigenvectors = eigh(
            K_dense, subset_by_index=[N - calc_k, N - 1]
        )

        idx = np.argsort(eigenvalues)[::-1]
        eigenvalues = eigenvalues[idx]
        eigenvectors = eigenvectors[:, idx]

        if var is not None:
            cum_var_ratio = np.cumsum(eigenvalues) / trace_K
            mask = cum_var_ratio <= var
            n_to_keep = np.sum(mask) + 1

            if n_components is not None:
                n_to_keep = min(n_to_keep, n_components)

            eigenvalues = eigenvalues[:n_to_keep]
            eigenvectors = eigenvectors[:, :n_to_keep]

        eigenvectors = eigenvectors / np.sqrt(np.maximum(eigenvalues, 1e-12))
        return eigenvalues, eigenvectors

    v0 = None

    def matmat(V):
        return (kernel @ V).reshape(kernel.N, -1)

    matvec = lambda v: (kernel @ v).reshape(-1)

    K_op = LinearOperator(
        shape=kernel.shape, matvec=matvec, matmat=matmat, dtype=kernel.X.dtype
    )

    current_k = 5 if n_components is None else n_components
    debug(
        f"Computing Kernel PCA with sparse iterative solver (starting k={current_k})"
    )

    while True:
        current_k = min(current_k, max_k)
        eigenvalues, eigenvectors = eigsh(K_op, k=current_k, which="LM", v0=v0)

        idx = np.argsort(eigenvalues)[::-1]
        eigenvalues = eigenvalues[idx]
        eigenvectors = eigenvectors[:, idx]

        if var is not None:
            cum_var_ratio = np.cumsum(eigenvalues) / trace_K
            debug(
                f"Sparse k={current_k} explains {cum_var_ratio[-1]:.4f} variance"
            )

            mask = cum_var_ratio <= var
            n_to_keep = np.sum(mask) + 1

            if n_components is not None:
                n_to_keep = min(n_to_keep, n_components)

            if cum_var_ratio[n_to_keep - 1] >= var or current_k >= max_k:
                eigenvalues = eigenvalues[:n_to_keep]
                eigenvectors = eigenvectors[:, :n_to_keep]
                break
            else:
                debug(
                    f"Threshold not met. Doubling components to k={current_k * 2}"
                )
                v0 = eigenvectors[:, 0]
                current_k = min(current_k * 2, max_k)
        else:
            break

    eigenvectors = eigenvectors / np.sqrt(np.maximum(eigenvalues, 1e-12))
    return eigenvalues, eigenvectors


def PCA(
    X: np.ndarray | Kernel,
    n_components: Optional[int] = None,
    var: Optional[float] = None,
    is_standard: bool = False,
):
    if is_standard and isinstance(X, np.ndarray):
        N, D = X.shape
        debug(f"Running Standard PCA (NumPy) on {N}x{D} data")
        X_centered = X - X.mean(axis=0)
        C = (X_centered.T @ X_centered) / (N - 1)

        calc_k = n_components if n_components else D
        eigenvalues, V = eigh(C, subset_by_index=[D - calc_k, D - 1])

        idx = np.argsort(eigenvalues)[::-1]
        eigenvalues = eigenvalues[idx]
        V = V[:, idx]

        if var is not None:
            total_var = np.sum(eigenvalues)
            cum_var_ratio = np.cumsum(eigenvalues) / total_var
            n_to_keep = np.sum(cum_var_ratio <= var) + 1
            debug(f"Keeping {n_to_keep} components for variance threshold")
            eigenvalues = eigenvalues[:n_to_keep]
            V = V[:, :n_to_keep]

        return eigenvalues, V

    if isinstance(X, Kernel):
        kernel = X
    else:
        kernel = LinearKernel(X).center()

    if not isinstance(kernel, CenteredKernel):
        kernel = kernel.center()

    trace_K = kernel.trace()
    max_k = kernel.N - 1

    if (
        isinstance(kernel.K, LinearKernel)
        and kernel.X.shape[0] > kernel.X.shape[1]
    ):
        return _pca_primal(kernel, n_components, max_k, trace_K, var)

    return _pca_kernel(kernel, n_components, max_k, trace_K, var)
