"""
Base protocols for backend implementations.

Defines the interface that all backends must implement.
"""

from typing import Protocol, Optional, Dict, Tuple, Any
from numpy.typing import NDArray


class PCABackend(Protocol):
    """Protocol for PCA computation backends."""

    backend_name: str
    """Name of the backend (e.g., 'numpy', 'jax')"""

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

        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
        """
        ...

    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)
        """
        ...

    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
        """
        ...

    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)
        """
        ...


class GPBackend(Protocol):
    """Protocol for GP computation backends."""

    backend_name: str
    """Name of the backend (e.g., 'numpy', 'jax')"""

    def build_gp_covariance(
        self,
        log_k: NDArray,
        length_scale: float,
        sigma: float,
        noise_level: Optional[float] = None,
        noise_cov: Optional[NDArray] = None,
        kernel_config: Optional[Any] = None,
    ) -> NDArray:
        """
        Build GP covariance matrix.

        Args:
            log_k: $\\log(k)$ values, shape (n,)
            length_scale: RBF kernel length scale in $\\log(k)$ space (backward compat)
            sigma: Signal standard deviation (GP amplitude) (backward compat)
            noise_level: Diagonal noise standard deviation
            noise_cov: Full N×N posterior covariance matrix
            kernel_config: KernelConfig for advanced kernels (RQ, Periodic, etc.)

        Returns:
            K: Covariance matrix, shape (n, n)
        """
        ...

    def compute_log_marginal_likelihood(
        self,
        delta_values: NDArray,
        log_k: NDArray,
        length_scale: float,
        sigma: float,
        noise_level: Optional[float] = None,
        noise_cov: Optional[NDArray] = None,
    ) -> float:
        """
        Compute log-marginal likelihood for given GP hyperparameters.

        Args:
            delta_values: Observed $\\delta(k)$ values, shape (n,)
            log_k: $\\log(k)$ values, shape (n,)
            length_scale: RBF kernel length scale
            sigma: Signal standard deviation
            noise_level: Diagonal noise standard deviation
            noise_cov: Full N×N posterior covariance matrix

        Returns:
            lml: Log-marginal likelihood
        """
        ...

    def compute_lml_landscape(
        self,
        delta_values: NDArray,
        log_k: NDArray,
        sigma_range: Tuple[float, float],
        length_scale_range: Tuple[float, float],
        n_sigma: int = 50,
        n_length: int = 50,
        noise_level: Optional[float] = None,
        posterior_cov: Optional[NDArray] = None,
    ) -> Dict:
        """
        Compute LML landscape over hyperparameter grid.

        Args:
            delta_values: Observed $\\delta(k)$ values, shape (n,)
            log_k: $\\log(k)$ values, shape (n,)
            sigma_range: (min, max) for signal amplitude grid
            length_scale_range: (min, max) for length scale grid
            n_sigma: Number of sigma grid points
            n_length: Number of length scale grid points
            noise_level: Fixed noise standard deviation
            posterior_cov: Full posterior covariance

        Returns:
            Dictionary with:
                - 'lml_grid': Grid of LML values, shape (n_sigma, n_length)
                - 'sigma_grid': Sigma values, shape (n_sigma,)
                - 'length_scale_grid': Length scale values, shape (n_length,)
                - 'optimal_sigma': Optimal sigma value
                - 'optimal_length_scale': Optimal length scale value
                - 'max_lml': Maximum LML value
        """
        ...
