"""
bishop.sampling – PRML Chapter 11
===================================
Uses ``emcee`` as the MCMC engine for ensemble sampling.
scipy.stats drives simple univariate sampling.

    from ml_ai_library.bishop.sampling import (
        emcee_sample,           # emcee.EnsembleSampler – production MCMC
        rejection_sample,       # scipy.stats thin wrapper
        importance_sample,      # scipy-based importance weighting
        slice_sample,           # scipy.stats.rv_continuous.rvs-based
        MetropolisHastings,     # thin emcee wrapper, single-chain API
        GibbsSampler,           # coordinate-wise Gibbs from conditionals
    )

PRML chapter map:
  rejection / importance    → §11.1 / §11.2
  MetropolisHastings        → §11.2 (backed by emcee EnsembleSampler)
  Gibbs                     → §11.3
  emcee_sample              → §11.4 (ensemble / affine-invariant MCMC)
"""

from __future__ import annotations
from typing import Callable, Optional, Sequence, Tuple

import numpy as np
from scipy import stats as _stats

import emcee                                    # ensemble MCMC engine


# ---------------------------------------------------------------------------
# Rejection Sampling (PRML §11.1) via scipy.stats
# ---------------------------------------------------------------------------

def rejection_sample(
    log_p: Callable[[np.ndarray], float],
    proposal: "_stats.rv_continuous",
    n_samples: int,
    M: float = 10.0,
) -> np.ndarray:
    """
    Rejection sampling (PRML §11.1) using a scipy.stats distribution as
    the proposal q and log_p as (unnormalised) log target.

    Parameters
    ----------
    log_p    : callable x -> log p~(x)
    proposal : scipy.stats continuous rv (e.g. scipy.stats.norm)
    n_samples: number of accepted samples to return
    M        : upper bound constant k s.t. k q(x) >= p~(x) for all x

    Returns
    -------
    accepted : (n_samples,) array
    """
    accepted = []
    while len(accepted) < n_samples:
        x  = proposal.rvs()
        u  = _stats.uniform.rvs()
        lq = proposal.logpdf(x)
        lp = log_p(np.atleast_1d(x))
        if np.log(u) < lp - np.log(M) - lq:
            accepted.append(x)
    return np.array(accepted[:n_samples])


# ---------------------------------------------------------------------------
# Importance Sampling (PRML §11.2) via scipy.stats
# ---------------------------------------------------------------------------

def importance_sample(
    log_p: Callable[[np.ndarray], float],
    proposal: "_stats.rv_continuous",
    n_samples: int,
) -> Tuple[np.ndarray, np.ndarray]:
    """
    Importance sampling (PRML §11.2).

    Returns (samples, normalised_weights).  Use for estimating expectations:
        E[f] ≈ sum(w_i * f(x_i))
    """
    xs = proposal.rvs(size=n_samples)
    log_w = np.array([log_p(np.atleast_1d(x)) - proposal.logpdf(x) for x in xs])
    log_w -= log_w.max()          # numerical stability
    w  = np.exp(log_w)
    w /= w.sum()
    return xs, w


# ---------------------------------------------------------------------------
# emcee EnsembleSampler wrapper (PRML §11.4 / production MCMC)
# ---------------------------------------------------------------------------

def emcee_sample(
    log_prob: Callable[[np.ndarray], float],
    ndim: int,
    n_walkers: int = 32,
    n_steps: int = 2000,
    burn_in: int = 200,
    rng: Optional[np.random.Generator] = None,
) -> np.ndarray:
    """
    Draw samples using emcee's affine-invariant ensemble sampler (PRML §11.4).

    Parameters
    ----------
    log_prob  : callable (ndim,) -> float, log of unnormalised target density
    ndim      : dimensionality
    n_walkers : ensemble size (must be even, >= 2*ndim)
    n_steps   : total MCMC steps
    burn_in   : steps to discard as warm-up
    rng       : numpy random generator (for reproducibility)

    Returns
    -------
    flat_samples : ((n_steps - burn_in) * n_walkers, ndim) array

    Example
    -------
    >>> import numpy as np
    >>> from ml_ai_library.bishop.sampling import emcee_sample
    >>> log_prob = lambda x: -0.5 * np.sum(x**2)
    >>> samples = emcee_sample(log_prob, ndim=2, n_steps=500, burn_in=50)
    >>> samples.shape[1]
    2
    """
    rng    = rng or np.random.default_rng()
    p0     = rng.standard_normal((n_walkers, ndim)) * 0.1
    sampler = emcee.EnsembleSampler(n_walkers, ndim, log_prob)
    sampler.run_mcmc(p0, n_steps, progress=False)
    return sampler.get_chain(discard=burn_in, flat=True)


# ---------------------------------------------------------------------------
# Metropolis-Hastings (PRML §11.2) as thin emcee wrapper
# ---------------------------------------------------------------------------

class MetropolisHastings:
    """
    Single-chain Metropolis-Hastings backed by emcee's 2-walker stretch move.

    Uses emcee's EnsembleSampler with n_walkers=2 (minimum), which
    internally implements the Metropolis accept/reject.

    Example
    -------
    >>> import numpy as np
    >>> mh = MetropolisHastings(log_prob=lambda x: -0.5 * x @ x, ndim=2)
    >>> samples = mh.sample(n=500)
    >>> samples.shape
    (500, 2)
    """

    def __init__(self, log_prob: Callable, ndim: int, n_walkers: int = 10):
        self.log_prob  = log_prob
        self.ndim      = ndim
        self.n_walkers = max(n_walkers, 2 * ndim if 2 * ndim >= 4 else 4)

    def sample(self, n: int = 1000, burn_in: int = 200,
               rng: Optional[np.random.Generator] = None) -> np.ndarray:
        """Return (n, ndim) sample array."""
        return emcee_sample(
            self.log_prob, self.ndim,
            n_walkers=self.n_walkers,
            n_steps=n + burn_in,
            burn_in=burn_in,
            rng=rng,
        )[:n]


# ---------------------------------------------------------------------------
# Gibbs Sampler (PRML §11.3) – coordinate-wise conditional sampling
# ---------------------------------------------------------------------------

class GibbsSampler:
    """
    Gibbs sampler (PRML §11.3).  Each conditional is provided as a callable.

    Parameters
    ----------
    conditionals : list of callables
                   conditionals[i](x) -> scalar sample of x[i] given rest

    Example
    -------
    >>> import numpy as np
    >>> from scipy import stats
    >>> # Bivariate normal via Gibbs
    >>> rho = 0.8
    >>> conds = [
    ...   lambda x: stats.norm.rvs(rho * x[1], np.sqrt(1 - rho**2)),
    ...   lambda x: stats.norm.rvs(rho * x[0], np.sqrt(1 - rho**2)),
    ... ]
    >>> sampler = GibbsSampler(conds)
    >>> samples = sampler.sample(500)
    >>> samples.shape
    (500, 2)
    """

    def __init__(self, conditionals: Sequence[Callable]):
        self.conditionals = list(conditionals)

    def sample(self, n: int, burn_in: int = 100,
               x0: Optional[np.ndarray] = None) -> np.ndarray:
        """Return (n, d) sample array after discarding `burn_in` steps."""
        d    = len(self.conditionals)
        x    = np.zeros(d) if x0 is None else np.asarray(x0, dtype=float).copy()
        out  = []
        for step in range(n + burn_in):
            for i, cond in enumerate(self.conditionals):
                x[i] = cond(x)
            if step >= burn_in:
                out.append(x.copy())
        return np.array(out)
