"""
NumPy/Numba backend for GP computations.

This is a streamlined version of the original GP code in gp.py,
implementing only the core functions needed for the backend interface.

Supported kernels:
- RBF (ExpSquared): Infinitely smooth, single length scale
- Rational Quadratic: Multi-scale structure
- Periodic (ExpSineSquared): Exactly repeating patterns
- Locally Periodic: RBF × Periodic product kernel
"""

import numpy as np
from scipy.linalg import cho_factor, cho_solve
from sklearn.gaussian_process.kernels import (
    RBF,
    ConstantKernel,
    RationalQuadratic,
    ExpSineSquared,
)
from typing import Optional, Dict, Tuple, Any
from numpy.typing import NDArray

# Import KernelType and KernelConfig from main gp module
try:
    from ...gp import KernelType, KernelConfig, build_kernel

    HAS_KERNEL_CONFIG = True
except ImportError:
    HAS_KERNEL_CONFIG = False
    KernelType = None
    KernelConfig = None
    build_kernel = None


class NumPyGPBackend:
    """NumPy/Numba backend for GP computations."""

    def __init__(self):
        self.backend_name = "numpy"

    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["KernelConfig"] = None,
    ) -> NDArray:
        """
        Build GP covariance matrix using sklearn kernels.

        Supports multiple kernel types through KernelConfig:
        - RBF (default): Infinitely smooth
        - Rational Quadratic: Multi-scale structure
        - Periodic: Exactly repeating patterns
        - Locally Periodic: RBF × Periodic product

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

        Returns:
            K: Covariance matrix, shape (n, n)
        """
        log_k = np.asarray(log_k).reshape(-1, 1)
        n = len(log_k)

        # Determine which kernel to use
        if kernel_config is not None:
            # Advanced kernel via KernelConfig
            if not HAS_KERNEL_CONFIG:
                raise ImportError("KernelConfig not available")
            signal_kernel = build_kernel(kernel_config)
            K_signal = signal_kernel(log_k)
        else:
            # Backward compatibility: simple RBF
            # Signal kernel: σ² * RBF(ℓ)
            signal_kernel = ConstantKernel(sigma**2) * RBF(length_scale)
            K_signal = signal_kernel(log_k)

        # Noise covariance
        if noise_cov is not None:
            K_noise = noise_cov
        else:
            K_noise = (noise_level**2) * np.eye(n) if noise_level else np.zeros((n, n))

        K = K_signal + K_noise
        return K

    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.

        LML = -0.5 * delta^T K^{-1} delta - 0.5 * log|K| - 0.5 * n * log(2π)

        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
        """
        delta_values = np.asarray(delta_values).ravel()
        n = len(delta_values)

        # Build covariance matrix
        K = self.build_gp_covariance(log_k, length_scale, sigma, noise_level, noise_cov)

        # Cholesky factorization for numerical stability
        try:
            L, lower = cho_factor(K, lower=True)
        except np.linalg.LinAlgError:
            # Singular matrix - return very low likelihood
            return -np.inf

        # log|K| = 2 * sum(log(diag(L)))
        log_det_K = 2 * np.sum(np.log(np.diag(L)))

        # Compute delta^T K^{-1} delta using Cholesky solve
        K_inv_delta = cho_solve((L, lower), delta_values)
        quad_form = delta_values @ K_inv_delta

        # Log-marginal likelihood
        lml = -0.5 * quad_form - 0.5 * log_det_K - 0.5 * n * np.log(2 * np.pi)

        return float(lml)

    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 (sequential evaluation).

        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', 'optimal_sigma', 'optimal_length_scale', etc.
        """
        sigma_grid = np.linspace(sigma_range[0], sigma_range[1], n_sigma)
        length_scale_grid = np.linspace(
            length_scale_range[0], length_scale_range[1], n_length
        )

        lml_grid = np.zeros((n_sigma, n_length))

        # Sequential grid search (can be slow for large grids)
        for i, sigma in enumerate(sigma_grid):
            for j, length_scale in enumerate(length_scale_grid):
                lml = self.compute_log_marginal_likelihood(
                    delta_values, log_k, length_scale, sigma, noise_level, posterior_cov
                )
                lml_grid[i, j] = lml

        # Find optimal hyperparameters
        i_opt, j_opt = np.unravel_index(np.argmax(lml_grid), lml_grid.shape)

        return {
            "lml_grid": lml_grid,
            "sigma_grid": sigma_grid,
            "length_scale_grid": length_scale_grid,
            "optimal_sigma": float(sigma_grid[i_opt]),
            "optimal_length_scale": float(length_scale_grid[j_opt]),
            "max_lml": float(lml_grid[i_opt, j_opt]),
        }


# Singleton instance
_numpy_backend = None


def get_numpy_backend() -> NumPyGPBackend:
    """Get NumPy backend instance (lazy initialization)."""
    global _numpy_backend
    if _numpy_backend is None:
        _numpy_backend = NumPyGPBackend()
    return _numpy_backend
