"""
Mathematical spaces for Boolean function analysis.

Provides conversions between different representations:
- BOOLEAN_CUBE: {0, 1}^n
- PLUS_MINUS_CUBE: {-1, +1}^n
- REAL: ℝ^n
- LOG: Log-probability space
- GAUSSIAN: Standard normal space

WARNING: Some conversions are lossy (continuous → discrete).
These will emit warnings unless explicitly silenced.
"""

import warnings
from enum import Enum, auto
from typing import Union

import numpy as np


class Space(Enum):
    BOOLEAN_CUBE = auto()  # {0,1}^n
    PLUS_MINUS_CUBE = auto()  # {-1,1}^n
    REAL = auto()  # ℝ^n
    LOG = auto()  # Log space
    GAUSSIAN = auto()  # Gaussian space

    @staticmethod
    def translate(
        input: Union[int, float, np.ndarray],
        source_space: "Space",
        target_space: "Space",
    ) -> Union[int, float, np.ndarray]:
        """
        Translate a scalar or array from one space to another.

        Args:
            input: Input value(s) to translate
            source_space: Source mathematical space
            target_space: Target mathematical space

        Returns:
            Translated value(s) in target space

        Examples:
            >>> Space.translate([0, 1], Space.BOOLEAN_CUBE, Space.PLUS_MINUS_CUBE)
            array([-1,  1])
            >>> Space.translate([-1, 1], Space.PLUS_MINUS_CUBE, Space.BOOLEAN_CUBE)
            array([0, 1])
        """
        if source_space == target_space:
            return input

        arr = np.asarray(input)

        # Boolean (0/1) → ±1
        if source_space == Space.BOOLEAN_CUBE and target_space == Space.PLUS_MINUS_CUBE:
            return 2 * arr - 1

        # ±1 → Boolean (0/1)
        if source_space == Space.PLUS_MINUS_CUBE and target_space == Space.BOOLEAN_CUBE:
            return ((arr + 1) // 2).astype(int)

        # Real-valued input → Boolean (mod 2) - LOSSY!
        if source_space == Space.REAL and target_space == Space.BOOLEAN_CUBE:
            warnings.warn(
                "Lossy conversion: REAL → BOOLEAN_CUBE uses round() % 2. "
                "Magnitude information is lost.",
                UserWarning,
                stacklevel=2,
            )
            return (np.round(arr) % 2).astype(int)

        # Boolean → Real (just cast)
        if source_space == Space.BOOLEAN_CUBE and target_space == Space.REAL:
            return arr.astype(float)

        # ±1 → Real
        if source_space == Space.PLUS_MINUS_CUBE and target_space == Space.REAL:
            return arr.astype(float)

        # Real → ±1 (sign function) - LOSSY!
        if source_space == Space.REAL and target_space == Space.PLUS_MINUS_CUBE:
            warnings.warn(
                "Lossy conversion: REAL → PLUS_MINUS_CUBE uses sign(). "
                "Magnitude information is lost.",
                UserWarning,
                stacklevel=2,
            )
            return np.where(arr >= 0, 1, -1)

        # LOG space conversions - LOSSY!
        if source_space == Space.LOG and target_space == Space.BOOLEAN_CUBE:
            warnings.warn(
                "Lossy conversion: LOG → BOOLEAN_CUBE uses threshold at 0.5. "
                "Probability magnitude is lost.",
                UserWarning,
                stacklevel=2,
            )
            return (np.exp(arr) > 0.5).astype(int)

        if source_space == Space.BOOLEAN_CUBE and target_space == Space.LOG:
            # Convert to log probabilities (avoid log(0))
            prob = np.clip(arr.astype(float), 1e-10, 1 - 1e-10)
            return np.log(prob)

        # GAUSSIAN space conversions - LOSSY!
        if source_space == Space.GAUSSIAN and target_space == Space.BOOLEAN_CUBE:
            warnings.warn(
                "Lossy conversion: GAUSSIAN → BOOLEAN_CUBE uses CDF threshold. "
                "Distribution information is lost.",
                UserWarning,
                stacklevel=2,
            )
            from scipy.stats import norm

            return (norm.cdf(arr) > 0.5).astype(int)

        if source_space == Space.BOOLEAN_CUBE and target_space == Space.GAUSSIAN:
            # Inverse normal CDF
            from scipy.stats import norm

            prob = np.clip(arr.astype(float), 1e-10, 1 - 1e-10)
            return norm.ppf(prob)

        # Cross-conversions through intermediate spaces
        if source_space == Space.LOG and target_space == Space.PLUS_MINUS_CUBE:
            # LOG -> BOOLEAN -> ±1
            boolean = Space.translate(arr, Space.LOG, Space.BOOLEAN_CUBE)
            return Space.translate(boolean, Space.BOOLEAN_CUBE, Space.PLUS_MINUS_CUBE)

        if source_space == Space.PLUS_MINUS_CUBE and target_space == Space.LOG:
            # ±1 -> BOOLEAN -> LOG
            boolean = Space.translate(arr, Space.PLUS_MINUS_CUBE, Space.BOOLEAN_CUBE)
            return Space.translate(boolean, Space.BOOLEAN_CUBE, Space.LOG)

        raise NotImplementedError(
            f"Translation from {source_space.name} to {target_space.name} not implemented."
        )

    @staticmethod
    def get_canonical_space() -> "Space":
        """Get the canonical space for internal computations."""
        return Space.BOOLEAN_CUBE

    @staticmethod
    def is_discrete(space: "Space") -> bool:
        """Check if space uses discrete values."""
        return space in {Space.BOOLEAN_CUBE, Space.PLUS_MINUS_CUBE}

    @staticmethod
    def is_continuous(space: "Space") -> bool:
        """Check if space uses continuous values."""
        return space in {Space.REAL, Space.LOG, Space.GAUSSIAN}

    @staticmethod
    def get_default_threshold(space: "Space") -> float:
        """Get default threshold for converting continuous to discrete."""
        if space == Space.REAL:
            return 0.5
        elif space == Space.LOG:
            return 0.0  # log(0.5) ≈ -0.693, but 0.0 is simpler
        elif space == Space.GAUSSIAN:
            return 0.0  # Standard normal median
        else:
            return 0.5


class Measure:
    """
    A probability measure on the Boolean hypercube.

    Supports uniform (p=0.5) and p-biased measures where each bit is 1
    independently with probability p.

    This unifies p-biased analysis: instead of calling separate functions
    from analysis/p_biased.py, pass a Measure to standard methods.

    Examples:
        >>> uniform = Measure.uniform()
        >>> biased = Measure.p_biased(0.3)
        >>> biased.p       # 0.3
        >>> biased.sigma    # sqrt(0.3 * 0.7)
    """

    def __init__(self, p: float = 0.5):
        if not 0 < p < 1:
            raise ValueError(f"Bias p must be in (0, 1), got {p}")
        self._p = p

    @classmethod
    def uniform(cls) -> "Measure":
        """The uniform measure (p = 1/2)."""
        return cls(0.5)

    @classmethod
    def p_biased(cls, p: float) -> "Measure":
        """P-biased measure: each bit is 1 with probability p."""
        return cls(p)

    @property
    def p(self) -> float:
        """Bias parameter."""
        return self._p

    @property
    def sigma(self) -> float:
        """Standard deviation of a single bit: sqrt(p(1-p))."""
        return np.sqrt(self._p * (1 - self._p))

    @property
    def is_uniform(self) -> bool:
        """True if this is the uniform measure (p = 0.5)."""
        return self._p == 0.5

    def sample(self, n: int, rng=None) -> np.ndarray:
        """Sample a random input from this measure.

        Args:
            n: Number of variables
            rng: NumPy random generator (optional)

        Returns:
            Binary array of length n
        """
        if rng is None:
            rng = np.random.default_rng()
        return (rng.random(n) < self._p).astype(int)

    def sample_batch(self, n: int, n_samples: int, rng=None) -> np.ndarray:
        """Sample multiple inputs from this measure.

        Args:
            n: Number of variables
            n_samples: Number of samples
            rng: NumPy random generator (optional)

        Returns:
            Array of shape (n_samples, n) with binary entries
        """
        if rng is None:
            rng = np.random.default_rng()
        return (rng.random((n_samples, n)) < self._p).astype(int)

    def weight(self, x) -> float:
        """Probability of a specific input under this measure.

        Args:
            x: Input as integer index or binary array

        Returns:
            mu_p(x) = p^|x| * (1-p)^(n-|x|)
        """
        if isinstance(x, (int, np.integer)):
            k = bin(int(x)).count("1")
            # Need n to compute, but can infer from context
            return self._p**k * (1 - self._p) ** (0)  # Partial, needs n
        x = np.asarray(x)
        k = int(np.sum(x))
        n = len(x)
        return self._p**k * (1 - self._p) ** (n - k)

    def __repr__(self) -> str:
        if self.is_uniform:
            return "Measure.uniform()"
        return f"Measure.p_biased({self._p})"

    def __eq__(self, other) -> bool:
        if not isinstance(other, Measure):
            return NotImplemented
        return self._p == other._p

    def __hash__(self) -> int:
        return hash(self._p)
