"""
Configuration module for backend selection and binning parameters.

Controls which backend (NumPy or JAX) is used for computations, and provides
configuration dataclasses for binning and power-law parameters.
"""

from dataclasses import dataclass
from functools import cached_property
from typing import List, Literal, Optional
import os

import numpy as np
from numpy.typing import NDArray

from .exceptions import BinningError

BackendType = Literal["numpy", "jax", "auto"]


@dataclass
class PrimefeatConfig:
    """Global configuration for PrimeFeat."""

    backend: BackendType = "auto"
    """
    Computation backend:
    - "auto": JAX if available and GPU detected, else NumPy
    - "jax": Force JAX (raises error if unavailable)
    - "numpy": Force NumPy/Numba
    """

    jax_device: Optional[str] = None
    """JAX device: "gpu", "cpu", or None for auto-detection"""

    jax_precision: Literal["float32", "float64"] = "float64"
    """JAX numerical precision (float32 is faster but less accurate)"""

    cache_compilations: bool = True
    """Cache JAX compilations across sessions"""

    verbose: bool = False
    """Print backend selection and compilation info"""


# Global config instance
_config = PrimefeatConfig()


def get_config() -> PrimefeatConfig:
    """Get current configuration."""
    return _config


def set_backend(backend: BackendType) -> None:
    """
    Set computation backend globally.

    Args:
        backend: "numpy", "jax", or "auto"

    Examples:
        >>> import primefeat as pf
        >>> pf.set_backend("jax")  # Use JAX backend
        >>> pf.set_backend("numpy")  # Use NumPy backend
        >>> pf.set_backend("auto")  # Auto-detect (default)
    """
    if backend not in ["numpy", "jax", "auto"]:
        raise ValueError(
            f"Invalid backend: {backend}. Must be 'numpy', 'jax', or 'auto'"
        )

    _config.backend = backend
    if _config.verbose:
        print(f"PrimeFeat backend set to: {backend}")


def get_backend() -> BackendType:
    """
    Get active backend (resolves 'auto').

    Returns:
        "numpy" or "jax" (never "auto")
    """
    if _config.backend == "auto":
        # Auto-detect: prefer JAX if available + GPU
        try:
            import jax

            devices = jax.devices()
            if any(d.platform == "gpu" for d in devices):
                if _config.verbose:
                    print("Auto-detected JAX with GPU - using JAX backend")
                return "jax"
            else:
                if _config.verbose:
                    print("JAX available but no GPU - using NumPy backend")
                return "numpy"
        except ImportError:
            if _config.verbose:
                print("JAX not available - using NumPy backend")
            return "numpy"
    return _config.backend


# Environment variable override
_env_backend = os.getenv("PRIMEFEAT_BACKEND")
if _env_backend in ["numpy", "jax"]:
    _config.backend = _env_backend


"""Configuration dataclasses for binning and power-law parameters."""


@dataclass(frozen=True)
class BinningConfig:
    """
    Configuration for binned primordial power spectrum.

    Immutable configuration object that encapsulates all binning parameters
    and provides derived quantities (bin centers, log bin centers, etc.).

    Matches CLASS `binned_Pk` implementation with logarithmic spacing.

    Args:
        k_start: Minimum $k$ for binning range (Mpc$^{-1}$)
        k_end: Maximum $k$ for binning range (Mpc$^{-1}$)
        nbins: Number of bins
        param_pattern: Parameter name pattern (use {i} for bin index)

    Example:
        >>> from primefeat import BinningConfig
        >>> config = BinningConfig(nbins=20, k_start=1e-3, k_end=0.23)
        >>> print(config.bin_centers)
        >>> print(config.bin_param_names)
    """

    k_start: float = 1e-2
    k_end: float = 1e-1
    nbins: int = 10
    param_pattern: str = "delta_{i}"

    @cached_property
    def bin_centers(self) -> NDArray[np.floating]:
        """
        Compute bin centers as CLASS does.

        Uses logarithmic spacing: bin centers are at geometric mean
        of bin edges in log space.

        Returns:
            Array of bin center $k$ values (length nbins)
        """
        log10_k_min = np.log10(self.k_start)
        log10_k_max = np.log10(self.k_end)
        delta_log10_k = (log10_k_max - log10_k_min) / self.nbins

        return np.array(
            [10 ** (log10_k_min + (i + 0.5) * delta_log10_k) for i in range(self.nbins)]
        )

    @cached_property
    def log_bin_centers(self) -> NDArray[np.floating]:
        """
        Natural log of bin centers (cached).

        Returns:
            $\\ln(k)$ for each bin center
        """
        return np.log(self.bin_centers)

    @cached_property
    def bin_param_names(self) -> List[str]:
        """
        Generate parameter names for all bins.

        Returns:
            Parameter names like ['delta_1', 'delta_2', ...]
        """
        return [self.param_pattern.format(i=i) for i in range(1, self.nbins + 1)]

    @cached_property
    def delta_log10_k(self) -> float:
        """
        Bin spacing in $\\log_{10}(k)$ space.

        Returns:
            $\\Delta \\log_{10}(k)$ per bin
        """
        return (np.log10(self.k_end) - np.log10(self.k_start)) / self.nbins

    def validate(self) -> None:
        """
        Validate configuration parameters.

        Raises:
            BinningError: If configuration is invalid
        """
        if self.k_start <= 0:
            raise BinningError(f"k_start must be positive, got {self.k_start}")

        if self.k_end <= self.k_start:
            raise BinningError(
                f"k_end ({self.k_end}) must be > k_start ({self.k_start})"
            )

        if self.nbins < 1:
            raise BinningError(f"nbins must be >= 1, got {self.nbins}")

        if "{i}" not in self.param_pattern:
            raise BinningError(
                f"param_pattern must contain '{{i}}', got '{self.param_pattern}'"
            )

    def __repr__(self) -> str:
        """String representation."""
        return (
            f"BinningConfig(k_start={self.k_start}, k_end={self.k_end}, "
            f"nbins={self.nbins}, param_pattern='{self.param_pattern}')"
        )


@dataclass(frozen=True)
class PowerlawConfig:
    """
    Configuration for power-law primordial spectrum parameters.

    Encapsulates $A_s$ and $n_s$ extraction settings including auto-detection,
    naming conventions, and fallback values.

    Args:
        As_param_name: Name of $A_s$ parameter (None = auto-detect)
        ns_param_name: Name of $n_s$ parameter
        k_pivot: Pivot scale in Mpc$^{-1}$ (Planck standard: 0.05)
        fallback_As: Fallback $A_s$ value if not found in chain
        fallback_ns: Fallback $n_s$ value if not found in chain

    Example:
        >>> from primefeat import PowerlawConfig
        >>> config = PowerlawConfig(k_pivot=0.05)
        >>> # Auto-detects A_s parameter name from common variants
        >>> config = PowerlawConfig(As_param_name='ln10^{10}A_s', k_pivot=0.002)
    """

    As_param_name: Optional[str] = None
    ns_param_name: str = "n_s"
    k_pivot: float = 0.05
    fallback_As: Optional[float] = None
    fallback_ns: Optional[float] = None

    @property
    def As_name_variants(self) -> List[str]:
        """
        Common $A_s$ parameter naming conventions.

        Returns:
            List of parameter name variants to try
        """
        return [
            "ln10^{10}A_s",
            "ln_A_s_1e10",
            "logA",
            "log(10^10 A_s)",
            "A_s",
            "As",
        ]

    def is_logarithmic(self, param_name: str) -> bool:
        """
        Check if parameter name indicates logarithmic encoding.

        Args:
            param_name: Parameter name to check

        Returns:
            True if name suggests logarithmic encoding
        """
        return any(x in param_name.lower() for x in ["ln", "log"])

    def convert_As(
        self, values: NDArray[np.floating], param_name: str
    ) -> NDArray[np.floating]:
        """
        Convert $A_s$ values based on naming convention.

        Args:
            values: Raw parameter values from chain
            param_name: Parameter name (determines conversion)

        Returns:
            Converted $A_s$ values in standard units

        Example:
            >>> config = PowerlawConfig()
            >>> # ln(10^10 A_s) = 3.0 -> A_s = exp(3)/1e10
            >>> As = config.convert_As(np.array([3.0]), 'ln10^{10}A_s')
        """
        if not self.is_logarithmic(param_name):
            return values

        if any(x in param_name for x in ["1e10", "10^10", "10^{10}"]):
            # ln(10^10 A_s) -> A_s
            return np.exp(values) / 1e10
        else:
            # log10(A_s) -> A_s
            return 10**values

    def __repr__(self) -> str:
        """String representation."""
        return (
            f"PowerlawConfig(As_param_name={self.As_param_name!r}, "
            f"ns_param_name={self.ns_param_name!r}, k_pivot={self.k_pivot}, "
            f"fallback_As={self.fallback_As}, fallback_ns={self.fallback_ns})"
        )
