"""Chain inspection utilities for MCMC parameter extraction."""

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

from .config import PowerlawConfig
from .exceptions import ChainParameterError


class ChainInspector:
    """
    Encapsulates logic for extracting parameters from MCMC chains.

    Handles both getdist MCSamples objects and dict-like structures.
    Provides consistent interface for parameter access with error handling.

    Example:
        >>> from primefeat import ChainInspector
        >>> inspector = ChainInspector(chain)
        >>> if inspector.has_parameter('delta_1'):
        ...     deltas = inspector.get_parameter('delta_1')
        >>> A_s, n_s = inspector.extract_powerlaw_params()
    """

    def __init__(self, chain: Union[Dict, Any]):
        """
        Initialize chain inspector.

        Args:
            chain: MCMC chain (getdist MCSamples or dict-like)
        """
        self.chain = chain
        self._param_names: Optional[List[str]] = None
        self._n_samples: Optional[int] = None

    @property
    def param_names(self) -> List[str]:
        """Get available parameter names (cached)."""
        if self._param_names is None:
            self._param_names = self._extract_param_names()
        return self._param_names

    @property
    def n_samples(self) -> int:
        """Get number of samples in chain (cached)."""
        if self._n_samples is None:
            self._n_samples = self._extract_n_samples()
        return self._n_samples

    def has_parameter(self, param_name: str) -> bool:
        """
        Check if parameter exists in chain.

        Args:
            param_name: Name of parameter to check

        Returns:
            True if parameter exists in chain
        """
        return param_name in self.param_names

    def get_parameter(
        self,
        param_name: str,
        fallback: Optional[float] = None,
        warn_on_fallback: bool = True,
    ) -> NDArray[np.floating]:
        """
        Extract parameter values from chain.

        Args:
            param_name: Name of parameter to extract
            fallback: Value to use if parameter not found
            warn_on_fallback: Whether to print warning when using fallback

        Returns:
            Array of parameter values (length n_samples)

        Raises:
            ChainParameterError: If parameter not found and no fallback provided
        """
        if param_name not in self.param_names:
            if fallback is not None:
                if warn_on_fallback:
                    print(
                        f"Parameter '{param_name}' not found in chain. "
                        f"Using fallback value: {fallback}"
                    )
                return np.full(self.n_samples, fallback)

            # Show first 10 available parameters
            available = self.param_names[:10]
            available_str = ", ".join(available)
            if len(self.param_names) > 10:
                available_str += f", ... ({len(self.param_names) - 10} more)"

            raise ChainParameterError(
                f"Parameter '{param_name}' not found in chain.\n"
                f"Available parameters: {available_str}\n"
                f"Provide fallback value or check parameter name."
            )

        try:
            return np.array(self.chain[param_name])
        except Exception as e:
            raise ChainParameterError(
                f"Failed to extract parameter '{param_name}': {e}"
            ) from e

    def extract_powerlaw_params(
        self,
        As_param_name: Optional[str] = None,
        ns_param_name: str = "n_s",
        fallback_As: Optional[float] = None,
        fallback_ns: Optional[float] = None,
    ) -> Tuple[NDArray[np.floating], NDArray[np.floating]]:
        """
        Extract $A_s$ and $n_s$ with automatic name detection and conversion.

        Detects common naming conventions:
            - `ln10^{10}A_s`, `ln_A_s_1e10`, `logA` (converts from log)
            - `A_s`, `As` (direct values)

        Args:
            As_param_name: Explicit $A_s$ parameter name (auto-detects if None)
            ns_param_name: Name of $n_s$ parameter
            fallback_As: Fallback $A_s$ value if not found
            fallback_ns: Fallback $n_s$ value if not found

        Returns:
            Tuple of ($A_s$ array, $n_s$ array)

        Raises:
            ChainParameterError: If parameters not found and no fallback provided
        """
        # Extract n_s
        n_s_array = self.get_parameter(
            ns_param_name, fallback=fallback_ns, warn_on_fallback=True
        )

        # Auto-detect A_s parameter name (only if not using fallback)
        if As_param_name is None:
            # Try to auto-detect, but use fallback if detection fails
            try:
                As_param_name = self._detect_As_param_name()
            except ChainParameterError:
                if fallback_As is not None:
                    # Use fallback instead of raising
                    print(
                        f"Could not auto-detect A_s parameter. "
                        f"Using fallback value: {fallback_As}"
                    )
                    return np.full(self.n_samples, fallback_As), n_s_array
                else:
                    raise

        # Extract A_s with conversion
        As_raw = self.get_parameter(
            As_param_name, fallback=fallback_As, warn_on_fallback=True
        )

        # Convert if logarithmic
        As_array = self._convert_As_values(As_raw, As_param_name)

        return As_array, n_s_array

    def _extract_param_names(self) -> List[str]:
        """Extract parameter names from chain."""
        try:
            # getdist MCSamples object
            return self.chain.getParamNames().list()
        except AttributeError:
            # dict-like object
            if hasattr(self.chain, "keys"):
                return list(self.chain.keys())
            else:
                raise ChainParameterError(
                    "Cannot extract parameter names from chain. "
                    "Chain must be getdist MCSamples or dict-like object."
                )

    def _extract_n_samples(self) -> int:
        """Extract number of samples from chain."""
        try:
            # getdist MCSamples object
            return self.chain.numrows
        except AttributeError:
            # dict-like object
            try:
                first_param = next(iter(self.param_names))
                return len(self.chain[first_param])
            except (StopIteration, KeyError, TypeError):
                raise ChainParameterError("Cannot determine chain size")

    def _detect_As_param_name(self) -> str:
        """
        Auto-detect $A_s$ parameter name from common variants.

        Returns:
            Detected parameter name

        Raises:
            ChainParameterError: If no A_s variant found
        """
        # Use canonical variant list from PowerlawConfig
        config = PowerlawConfig()
        As_variants = config.As_name_variants

        for variant in As_variants:
            if variant in self.param_names:
                print(f"Auto-detected A_s parameter: '{variant}'")
                return variant

        raise ChainParameterError(
            f"Could not auto-detect A_s parameter. Tried: {As_variants}\n"
            f"Available parameters: {self.param_names[:10]}...\n"
            f"Specify As_param_name explicitly."
        )

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

        Delegates to PowerlawConfig.convert_As() for consistent conversion logic.

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

        Returns:
            Converted $A_s$ values in standard units
        """
        config = PowerlawConfig()
        return config.convert_As(values, param_name)
