"""
JAX backend for GPU-accelerated PCA computations.

This implementation provides:
1. GPU acceleration for all operations
2. JIT compilation for speed
3. Automatic differentiation capabilities
4. Vectorized batch operations via vmap

The PCA is computed using eigendecomposition of the covariance matrix,
which is mathematically equivalent to SVD-based PCA but more suitable
for autodiff applications.
"""

import numpy as np
from typing import Dict, Tuple, Optional, Union
from numpy.typing import NDArray

# Check for JAX availability
try:
    import jax
    import jax.numpy as jnp
    from jax import jit, vmap

    HAS_JAX = True
except ImportError:
    HAS_JAX = False
    jax = None
    jnp = None
    jit = None
    vmap = None


def _check_jax_available():
    """Check if JAX is available and raise informative error if not."""
    if not HAS_JAX:
        raise ImportError(
            "JAX backend requires jax. Install with: pip install jax jaxlib"
        )


# JIT-compiled core functions (defined at module level for efficiency)
if HAS_JAX:
    # Configure JAX for numerical precision
    jax.config.update("jax_enable_x64", True)

    @jit
    def _center_data_jax(X: jnp.ndarray) -> Tuple[jnp.ndarray, jnp.ndarray]:
        """
        Center data by subtracting the mean.

        Args:
            X: Data matrix, shape (n_samples, n_features)

        Returns:
            X_centered: Centered data
            mean: Mean vector
        """
        mean = jnp.mean(X, axis=0)
        X_centered = X - mean
        return X_centered, mean

    @jit
    def _compute_covariance_jax(X_centered: jnp.ndarray) -> jnp.ndarray:
        """
        Compute covariance matrix from centered data.

        Uses unbiased estimator (n-1 normalization).

        Args:
            X_centered: Centered data, shape (n_samples, n_features)

        Returns:
            Covariance matrix, shape (n_features, n_features)
        """
        n_samples = X_centered.shape[0]
        # Covariance: (X^T @ X) / (n - 1)
        return jnp.dot(X_centered.T, X_centered) / (n_samples - 1)

    @jit
    def _compute_eigendecomp_jax(
        cov: jnp.ndarray,
    ) -> Tuple[jnp.ndarray, jnp.ndarray]:
        """
        Compute eigendecomposition of symmetric covariance matrix.

        Uses jnp.linalg.eigh which is optimized for symmetric matrices.
        Returns eigenvalues and eigenvectors sorted in descending order.

        Args:
            cov: Covariance matrix, shape (n_features, n_features)

        Returns:
            eigenvalues: Sorted eigenvalues (descending), shape (n_features,)
            eigenvectors: Corresponding eigenvectors (columns), shape (n_features, n_features)
        """
        # eigh returns eigenvalues in ascending order
        eigenvalues, eigenvectors = jnp.linalg.eigh(cov)

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

        return eigenvalues, eigenvectors

    @jit
    def _fit_pca_full_jax(X: jnp.ndarray) -> Dict[str, jnp.ndarray]:
        """
        Fit full PCA model using eigendecomposition (JIT-compiled).

        This computes all principal components. Use _select_components()
        to select top n_components after this function.

        Args:
            X: Data matrix, shape (n_samples, n_features)

        Returns:
            Dictionary with all components, mean, explained_variance, etc.
        """
        n_samples, n_features = X.shape

        # Center the data
        X_centered, mean = _center_data_jax(X)

        # Compute covariance matrix
        cov = _compute_covariance_jax(X_centered)

        # Eigendecomposition
        eigenvalues, eigenvectors = _compute_eigendecomp_jax(cov)

        # All components stored as rows (n_features, n_features)
        components = eigenvectors.T

        # Total variance for ratio computation
        total_variance = jnp.sum(eigenvalues)
        explained_variance_ratio = eigenvalues / total_variance

        # Singular values: sqrt(explained_variance * (n_samples - 1))
        singular_values = jnp.sqrt(eigenvalues * (n_samples - 1))

        return {
            "components": components,
            "mean": mean,
            "explained_variance": eigenvalues,
            "explained_variance_ratio": explained_variance_ratio,
            "singular_values": singular_values,
            "total_variance": total_variance,
        }

    def _fit_pca_jax(X: jnp.ndarray, n_components: int) -> Dict[str, jnp.ndarray]:
        """
        Fit PCA model using eigendecomposition.

        This is a wrapper that calls the JIT-compiled full PCA and then
        selects the top n_components.

        Args:
            X: Data matrix, shape (n_samples, n_features)
            n_components: Number of components to extract

        Returns:
            Dictionary with components, mean, explained_variance, etc.
        """
        # Fit full PCA (JIT-compiled)
        result = _fit_pca_full_jax(X)

        # Select top n_components (done outside JIT to handle dynamic slicing)
        return {
            "components": result["components"][:n_components],
            "mean": result["mean"],
            "explained_variance": result["explained_variance"][:n_components],
            "explained_variance_ratio": result["explained_variance_ratio"][
                :n_components
            ],
            "singular_values": result["singular_values"][:n_components],
            "total_variance": result["total_variance"],
        }

    @jit
    def _transform_jax(
        X: jnp.ndarray, components: jnp.ndarray, mean: jnp.ndarray
    ) -> jnp.ndarray:
        """
        Project data onto principal components (JIT-compiled).

        Args:
            X: Data matrix, shape (n_samples, n_features)
            components: Principal components, shape (n_components, n_features)
            mean: Data mean, shape (n_features,)

        Returns:
            Projected data, shape (n_samples, n_components)
        """
        X_centered = X - mean
        return jnp.dot(X_centered, components.T)

    @jit
    def _inverse_transform_jax(
        X_transformed: jnp.ndarray, components: jnp.ndarray, mean: jnp.ndarray
    ) -> jnp.ndarray:
        """
        Reconstruct data from principal components (JIT-compiled).

        Args:
            X_transformed: Projected data, shape (n_samples, n_components)
            components: Principal components, shape (n_components, n_features)
            mean: Data mean, shape (n_features,)

        Returns:
            Reconstructed data, shape (n_samples, n_features)
        """
        return jnp.dot(X_transformed, components) + mean


class JAXPCABackend:
    """JAX backend for GPU-accelerated PCA computations."""

    def __init__(self):
        _check_jax_available()
        self.backend_name = "jax"

        # Configure JAX for numerical precision
        jax.config.update("jax_enable_x64", True)

    def fit(
        self,
        X: NDArray,
        n_components: int,
    ) -> Dict[str, NDArray]:
        """
        Fit PCA model to data using JAX eigendecomposition.

        This method is autodiff-compatible: the returned arrays are JAX arrays
        that can be used with jax.grad() for gradient computation.

        Args:
            X: Data matrix, shape (n_samples, n_features)
            n_components: Number of principal components to compute

        Returns:
            Dictionary containing:
                - 'components': Principal component vectors, shape (n_components, n_features)
                - 'mean': Mean of the data, shape (n_features,)
                - 'explained_variance': Variance explained by each component
                - 'explained_variance_ratio': Fraction of variance explained
                - 'singular_values': Singular values corresponding to each component
        """
        # Convert to JAX array
        X_jax = jnp.asarray(X)
        n_samples, n_features = X_jax.shape

        # Limit n_components to valid range
        n_components = min(n_components, n_samples, n_features)

        # Fit using JIT-compiled function
        result = _fit_pca_jax(X_jax, n_components)

        # Convert back to NumPy for compatibility
        return {
            "components": np.asarray(result["components"]),
            "mean": np.asarray(result["mean"]),
            "explained_variance": np.asarray(result["explained_variance"]),
            "explained_variance_ratio": np.asarray(result["explained_variance_ratio"]),
            "singular_values": np.asarray(result["singular_values"]),
        }

    def fit_jax(
        self,
        X: jnp.ndarray,
        n_components: int,
    ) -> Dict[str, jnp.ndarray]:
        """
        Fit PCA model returning JAX arrays (for autodiff).

        This is the pure JAX version that returns JAX arrays directly,
        enabling use with jax.grad(), jax.jacobian(), etc.

        Args:
            X: Data matrix as JAX array, shape (n_samples, n_features)
            n_components: Number of principal components to compute

        Returns:
            Dictionary containing JAX arrays (differentiable)
        """
        n_samples, n_features = X.shape
        n_components = min(n_components, n_samples, n_features)
        return _fit_pca_jax(X, n_components)

    def transform(
        self,
        X: NDArray,
        components: NDArray,
        mean: NDArray,
    ) -> NDArray:
        """
        Project data onto principal components.

        Args:
            X: Data matrix, shape (n_samples, n_features)
            components: Principal components, shape (n_components, n_features)
            mean: Data mean for centering, shape (n_features,)

        Returns:
            Projected data, shape (n_samples, n_components)
        """
        X_jax = jnp.asarray(X)
        components_jax = jnp.asarray(components)
        mean_jax = jnp.asarray(mean)

        result = _transform_jax(X_jax, components_jax, mean_jax)
        return np.asarray(result)

    def transform_jax(
        self,
        X: jnp.ndarray,
        components: jnp.ndarray,
        mean: jnp.ndarray,
    ) -> jnp.ndarray:
        """
        Project data onto principal components (pure JAX for autodiff).

        Args:
            X: Data matrix as JAX array
            components: Principal components as JAX array
            mean: Data mean as JAX array

        Returns:
            Projected data as JAX array (differentiable)
        """
        return _transform_jax(X, components, mean)

    def fit_transform(
        self,
        X: NDArray,
        n_components: int,
    ) -> Tuple[NDArray, Dict[str, NDArray]]:
        """
        Fit PCA and transform data in one step.

        Args:
            X: Data matrix, shape (n_samples, n_features)
            n_components: Number of principal components to compute

        Returns:
            Tuple of:
                - transformed: Projected data, shape (n_samples, n_components)
                - fit_result: Dictionary from fit() method
        """
        # Convert to JAX
        X_jax = jnp.asarray(X)
        n_samples, n_features = X_jax.shape
        n_components = min(n_components, n_samples, n_features)

        # Fit
        result = _fit_pca_jax(X_jax, n_components)

        # Transform
        X_transformed = _transform_jax(X_jax, result["components"], result["mean"])

        # Convert back to NumPy
        fit_result = {
            "components": np.asarray(result["components"]),
            "mean": np.asarray(result["mean"]),
            "explained_variance": np.asarray(result["explained_variance"]),
            "explained_variance_ratio": np.asarray(result["explained_variance_ratio"]),
            "singular_values": np.asarray(result["singular_values"]),
        }

        return np.asarray(X_transformed), fit_result

    def fit_transform_jax(
        self,
        X: jnp.ndarray,
        n_components: int,
    ) -> Tuple[jnp.ndarray, Dict[str, jnp.ndarray]]:
        """
        Fit PCA and transform data (pure JAX for autodiff).

        Args:
            X: Data matrix as JAX array
            n_components: Number of components

        Returns:
            Tuple of transformed data and fit results (all JAX arrays)
        """
        n_samples, n_features = X.shape
        n_components = min(n_components, n_samples, n_features)

        result = _fit_pca_jax(X, n_components)
        X_transformed = _transform_jax(X, result["components"], result["mean"])

        return X_transformed, result

    def inverse_transform(
        self,
        X_transformed: NDArray,
        components: NDArray,
        mean: NDArray,
    ) -> NDArray:
        """
        Reconstruct data from principal components.

        Args:
            X_transformed: Projected data, shape (n_samples, n_components)
            components: Principal components, shape (n_components, n_features)
            mean: Data mean for centering, shape (n_features,)

        Returns:
            Reconstructed data, shape (n_samples, n_features)
        """
        X_jax = jnp.asarray(X_transformed)
        components_jax = jnp.asarray(components)
        mean_jax = jnp.asarray(mean)

        result = _inverse_transform_jax(X_jax, components_jax, mean_jax)
        return np.asarray(result)

    def inverse_transform_jax(
        self,
        X_transformed: jnp.ndarray,
        components: jnp.ndarray,
        mean: jnp.ndarray,
    ) -> jnp.ndarray:
        """
        Reconstruct data from principal components (pure JAX for autodiff).

        Args:
            X_transformed: Projected data as JAX array
            components: Principal components as JAX array
            mean: Data mean as JAX array

        Returns:
            Reconstructed data as JAX array (differentiable)
        """
        return _inverse_transform_jax(X_transformed, components, mean)

    def fit_transform_batch(
        self,
        X_batch: NDArray,
        n_components: int,
    ) -> Tuple[NDArray, Dict[str, NDArray]]:
        """
        Fit PCA to multiple datasets in parallel using vmap.

        This is useful for processing multiple chains independently.

        Args:
            X_batch: Batch of data matrices, shape (n_datasets, n_samples, n_features)
            n_components: Number of principal components to compute

        Returns:
            Tuple of:
                - transformed_batch: Projected data, shape (n_datasets, n_samples, n_components)
                - fit_results: Dictionary with batched results
        """
        X_batch_jax = jnp.asarray(X_batch)
        n_datasets, n_samples, n_features = X_batch_jax.shape
        n_components = min(n_components, n_samples, n_features)

        # Vectorize fit_transform over the batch dimension
        def fit_single(X):
            result = _fit_pca_jax(X, n_components)
            X_transformed = _transform_jax(X, result["components"], result["mean"])
            return X_transformed, result

        # vmap over first axis (datasets)
        batched_fit = vmap(fit_single)
        X_transformed_batch, results_batch = batched_fit(X_batch_jax)

        # Convert to NumPy
        fit_results = {
            "components": np.asarray(results_batch["components"]),
            "mean": np.asarray(results_batch["mean"]),
            "explained_variance": np.asarray(results_batch["explained_variance"]),
            "explained_variance_ratio": np.asarray(
                results_batch["explained_variance_ratio"]
            ),
            "singular_values": np.asarray(results_batch["singular_values"]),
        }

        return np.asarray(X_transformed_batch), fit_results


# Singleton instance
_jax_pca_backend = None


def get_jax_pca_backend() -> JAXPCABackend:
    """Get JAX PCA backend instance (lazy initialization)."""
    global _jax_pca_backend
    if _jax_pca_backend is None:
        _jax_pca_backend = JAXPCABackend()
    return _jax_pca_backend


def is_jax_available() -> bool:
    """Check if JAX is available."""
    return HAS_JAX
