from typing import Optional
import torch

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 for feature dimension D={D}")

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

    C = torch.mm(X_centered.T, X_centered)

    eigenvalues, V = torch.linalg.eigh(C)

    idx = torch.argsort(eigenvalues, descending=True)
    eigenvalues = eigenvalues[idx]
    V = V[:, idx]

    if var is not None:
        cum_var_ratio = torch.cumsum(eigenvalues, dim=0) / trace_K
        mask = cum_var_ratio <= var
        n_to_keep = int(torch.sum(mask).item()) + 1
        debug(f"Variance threshold {var:.2f} requires {n_to_keep} components.")

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

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

    alpha = torch.mm(X_centered, V) / torch.clamp(eigenvalues, min=1e-12)

    return eigenvalues, alpha


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

        eigenvalues, eigenvectors = torch.linalg.eigh(K_dense)

        idx = torch.argsort(eigenvalues, descending=True)
        eigenvalues = eigenvalues[idx]
        eigenvectors = eigenvectors[:, idx]

        if var is not None:
            cum_var_ratio = torch.cumsum(eigenvalues, dim=0) / trace_K
            mask = cum_var_ratio <= var
            n_to_keep = int(torch.sum(mask).item()) + 1

            if n_components is not None:
                n_to_keep = min(n_to_keep, n_components)
        elif n_components is not None:
            n_to_keep = min(n_components, kernel.N)
        else:
            n_to_keep = max_k

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

        eigenvectors = eigenvectors / torch.sqrt(
            torch.clamp(eigenvalues, min=1e-12)
        )
        return eigenvalues, eigenvectors

    def matmat(v):
        return kernel @ v

    q = None
    current_k = 5 if n_components is None else n_components
    debug(
        f"Computing Kernel PCA using iterative LOBPCG (starting k={current_k})"
    )

    while True:
        current_k = min(current_k, max_k)

        if q is None or q.shape[1] != current_k:
            q = torch.randn(
                (kernel.N, current_k),
                device=kernel.X.device,
                dtype=kernel.X.dtype,
            )

        eigenvalues, eigenvectors = torch.lobpcg(A=matmat, k=current_k, X=q)

        idx = torch.argsort(eigenvalues, descending=True)
        eigenvalues = eigenvalues[idx]
        eigenvectors = eigenvectors[:, idx]

        if var is not None:
            cum_var_ratio = torch.cumsum(eigenvalues, dim=0) / trace_K
            debug(
                f"LOBPCG k={current_k} covers {cum_var_ratio[-1]:.4f} variance"
            )

            mask = cum_var_ratio <= var
            n_to_keep = int(torch.sum(mask).item()) + 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"Variance threshold not met. Increasing k to {current_k * 2}"
                )
                q = torch.randn(
                    (kernel.N, current_k * 2),
                    device=kernel.X.device,
                    dtype=kernel.X.dtype,
                )
                q[:, :current_k] = eigenvectors

                current_k = min(current_k * 2, max_k)
        else:
            break

    eigenvectors = eigenvectors / torch.sqrt(
        torch.clamp(eigenvalues, min=1e-12)
    )

    return eigenvalues, eigenvectors


def PCA(
    X: torch.Tensor | Kernel,
    n_components: Optional[int] = None,
    var: Optional[float] = None,
    is_standard: bool = False,
):
    if is_standard and isinstance(X, torch.Tensor):
        N, D = X.shape
        debug(f"Running Standard PCA decomposition for N={N}, D={D}")
        X_centered = X - X.mean(0)

        C = torch.mm(X_centered.T, X_centered) / (N - 1)

        eigenvalues, V = torch.linalg.eigh(C)

        eigenvalues = torch.flip(eigenvalues, dims=[0])
        V = torch.flip(V, dims=[1])

        calc_k = n_components if n_components else D
        eigenvalues, V = eigenvalues[:calc_k], V[:, :calc_k]

        if var is not None:
            cum_var_ratio = torch.cumsum(eigenvalues, 0) / torch.sum(
                eigenvalues
            )
            n_to_keep = torch.sum(cum_var_ratio <= var).item() + 1
            debug(f"Variance check: keeping top {n_to_keep} components")
            eigenvalues, V = eigenvalues[:n_to_keep], 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)
