"""Parallel tempering (replica-exchange) MCMC for beam-parameter inference.

Fully vectorised: one JIT-compiled Metropolis kernel (vmap) and swap step
between adjacent chains. No dynamic Python logic inside jitted code.
"""

from typing import Any, Callable

import jax
import jax.numpy as jnp
from jax.scipy.stats import norm


def build_beta_ladder(
    n_chains: int,
    beta_max: float = 1.0,
    beta_cut1: float = 0.97,
    beta_cut2: float = 0.0026,
    beta_min: float = 0.0025,
    frac_top: float = 0.2,
    frac_mid: float = 0.7,
    frac_bottom: float = 0.1,
) -> jnp.ndarray:
    """Build a temperature ladder (inverse temperatures beta) for parallel tempering.

    Parameters
    ----------
    n_chains : int
        Number of chains (replicas).
    beta_max : float
        Coldest chain (beta = 1 / T, so beta_max = 1 is the target posterior).
    beta_cut1, beta_cut2, beta_min : float
        Breakpoints; must satisfy beta_max > beta_cut1 > beta_cut2 > beta_min.
    frac_top, frac_mid, frac_bottom : float
        Fraction of chains in top/mid/bottom segments (must sum to 1).

    Returns
    -------
    jnp.ndarray
        Shape (n_chains,), strictly decreasing from cold to hot.
    """
    if not (beta_max > beta_cut1 > beta_cut2 > beta_min):
        raise ValueError("Require beta_max > beta_cut1 > beta_cut2 > beta_min")

    n_top = int(n_chains * frac_top)
    n_mid = int(n_chains * frac_mid)
    n_bottom = n_chains - n_top - n_mid

    beta_top = jnp.linspace(beta_max, beta_cut1, n_top, endpoint=False)
    beta_mid = jnp.geomspace(beta_cut1, beta_cut2, n_mid, endpoint=False)
    beta_bottom = jnp.linspace(beta_cut2, beta_min, n_bottom, endpoint=True)

    betas = jnp.concatenate([beta_top, beta_mid, beta_bottom])
    betas = jnp.sort(betas)[::-1]
    return betas


def _log_gaussian_prior(theta: jnp.ndarray, mean: jnp.ndarray, std: jnp.ndarray) -> jnp.ndarray:
    """Log density of independent Gaussian prior (untruncated)."""
    z = (theta - mean) / std
    return -0.5 * jnp.sum(z * z)


def _log_uniform_prior(theta: jnp.ndarray, lower: jnp.ndarray, upper: jnp.ndarray) -> jnp.ndarray:
    """Log density of uniform prior: 0 inside [lower, upper], -inf outside."""
    inside = jnp.all((theta >= lower) & (theta <= upper))
    return jnp.where(inside, 0.0, -jnp.inf)


def _reflect(x: jnp.ndarray, lower: jnp.ndarray, upper: jnp.ndarray) -> jnp.ndarray:
    """Reflect proposal back into [lower, upper]."""
    return jnp.where(
        x < lower,
        2.0 * lower - x,
        jnp.where(x > upper, 2.0 * upper - x, x),
    )


def _propose(
    key: jax.Array,
    theta: jnp.ndarray,
    proposal_std: jnp.ndarray,
    lower: jnp.ndarray,
    upper: jnp.ndarray,
) -> tuple[jax.Array, jnp.ndarray]:
    """Random-walk proposal with reflection at bounds."""
    key, sub = jax.random.split(key)
    eps = jax.random.normal(sub, shape=theta.shape) * proposal_std
    theta_prop = _reflect(theta + eps, lower, upper)
    return key, theta_prop


def sample_from_prior(prior: dict[str, Any], key: jax.Array) -> tuple[jax.Array, jnp.ndarray]:
    """Sample from a prior (uniform or truncated Gaussian). JAX-friendly.

    Parameters
    ----------
    prior : dict
        Must contain "type" ("uniform" or "gaussian") and:
        - uniform: "lower", "upper" (arrays)
        - gaussian: "mean", "std", "lower", "upper" (arrays)
    key : jax.Array
        PRNG key.

    Returns
    -------
    key_out : jax.Array
        Updated PRNG key.
    sample : jnp.ndarray
        One sample, shape matching lower/upper.
    """
    typ = prior["type"]

    if typ == "uniform":
        lower = jnp.asarray(prior["lower"], dtype=jnp.float32)
        upper = jnp.asarray(prior["upper"], dtype=jnp.float32)
        key, sub = jax.random.split(key)
        sample = jax.random.uniform(sub, shape=lower.shape, minval=lower, maxval=upper)
        return key, sample

    if typ == "gaussian":
        mean = jnp.asarray(prior["mean"], dtype=jnp.float32)
        std = jnp.asarray(prior["std"], dtype=jnp.float32)
        lower = jnp.asarray(prior["lower"], dtype=jnp.float32)
        upper = jnp.asarray(prior["upper"], dtype=jnp.float32)
        key, sub = jax.random.split(key)
        a = (lower - mean) / std
        b = (upper - mean) / std
        u = jax.random.uniform(sub, shape=mean.shape)
        cdf_a = norm.cdf(a)
        cdf_b = norm.cdf(b)
        u_scaled = cdf_a + u * (cdf_b - cdf_a)
        z = norm.ppf(jnp.clip(u_scaled, 1e-10, 1.0 - 1e-10))
        sample = mean + std * z
        return key, sample

    raise ValueError(f"Unsupported prior type: {typ!r}")


def _build_proposal_std_chain(
    proposal_std: jnp.ndarray,
    betas: jnp.ndarray,
    n_chains: int,
    dim: int,
) -> jnp.ndarray:
    """Construct per-chain proposal std, validating shape and scaling by temperature.

    Accepts either a base (D,) vector that is scaled per chain, or a full
    (n_chains, D) array with one vector per chain.
    """
    if proposal_std.ndim == 1:
        # Scale base proposal by temperature (colder chains move more cautiously).
        alpha = 0.5
        return proposal_std * (betas[:, None] ** (-alpha))

    if proposal_std.ndim == 2:
        if proposal_std.shape != (n_chains, dim):
            raise ValueError(
                f"proposal_std shape {proposal_std.shape} expected ({n_chains}, {dim})"
            )
        return proposal_std

    raise ValueError(
        f"proposal_std must be 1D (D,) or 2D (n_chains, D), got {proposal_std.shape}"
    )


def _build_log_prior_fn(prior: dict[str, Any], lower: jnp.ndarray, upper: jnp.ndarray) -> Callable[[jnp.ndarray], jnp.ndarray]:
    """Create a log-prior function from a prior dictionary."""
    prior_type = prior["type"]

    if prior_type == "gaussian":
        mean = jnp.asarray(prior["mean"], dtype=jnp.float32)
        std = jnp.asarray(prior["std"], dtype=jnp.float32)
        return lambda th: _log_gaussian_prior(th, mean, std)

    if prior_type == "uniform":
        return lambda th: _log_uniform_prior(th, lower, upper)

    raise ValueError(f"Unsupported prior type: {prior_type!r}")


def _make_mh_step_all_chains(
    log_prior_fn: Callable[[jnp.ndarray], jnp.ndarray],
    base_loglike_fn: Callable[[jnp.ndarray], jnp.ndarray],
    betas: jnp.ndarray,
    proposal_std_chain: jnp.ndarray,
    lower: jnp.ndarray,
    upper: jnp.ndarray,
) -> Callable[[jax.Array, jnp.ndarray], tuple[jax.Array, jnp.ndarray, jnp.ndarray]]:
    """Factory for the JIT-compiled Metropolis–Hastings step across all chains."""

    @jax.jit
    def mh_step_all_chains(keys: jax.Array, thetas: jnp.ndarray) -> tuple[jax.Array, jnp.ndarray, jnp.ndarray]:
        def mh_step_one_chain(
            k: jax.Array,
            th: jnp.ndarray,
            beta: jnp.ndarray,
            std_vec: jnp.ndarray,
        ):
            k, th_prop = _propose(k, th, std_vec, lower, upper)
            log_post = log_prior_fn(th) + beta * base_loglike_fn(th)
            log_post_prop = log_prior_fn(th_prop) + beta * base_loglike_fn(th_prop)
            log_alpha = log_post_prop - log_post
            k, sub = jax.random.split(k)
            accept = jnp.log(jax.random.uniform(sub)) < log_alpha
            th_new = jnp.where(accept, th_prop, th)
            return k, th_new, accept

        keys_out, thetas_out, accepts = jax.vmap(mh_step_one_chain)(
            keys,
            thetas,
            betas,
            proposal_std_chain,
        )
        return keys_out, thetas_out, accepts

    return mh_step_all_chains


def _swap_samples_one_round(
    thetas: jnp.ndarray,
    key: jax.Array,
    round_idx: int,
    betas: jnp.ndarray,
    untempered_log_post: Callable[[jnp.ndarray], jnp.ndarray],
) -> tuple[jnp.ndarray, jax.Array, jnp.ndarray, jnp.ndarray]:
    """Perform one replica-exchange round (even or odd adjacent pairs)."""
    n_chains = thetas.shape[0]
    even_round = (round_idx % 2) == 0

    if even_round:
        pairs_i = jnp.arange(0, n_chains - 1, 2)
    else:
        pairs_i = jnp.arange(1, n_chains - 1, 2)
    pairs_j = pairs_i + 1

    thetas_i = thetas[pairs_i]
    thetas_j = thetas[pairs_j]
    betas_i = betas[pairs_i]
    betas_j = betas[pairs_j]

    logp_i = jax.vmap(untempered_log_post)(thetas_i)
    logp_j = jax.vmap(untempered_log_post)(thetas_j)
    deltas = (logp_j - logp_i) * (betas_i - betas_j)

    key, sub = jax.random.split(key)
    uniforms = jax.random.uniform(sub, deltas.shape)
    accept = jnp.log(uniforms) < deltas
    mask = accept[:, None]

    new_i = jnp.where(mask, thetas_j, thetas_i)
    new_j = jnp.where(mask, thetas_i, thetas_j)
    thetas = thetas.at[pairs_i].set(new_i)
    thetas = thetas.at[pairs_j].set(new_j)

    return thetas, key, accept, pairs_i


def _run_replica_exchange_loop(
    key: jax.Array,
    theta0_all: jnp.ndarray,
    betas: jnp.ndarray,
    mh_step_all_chains: Callable[[jax.Array, jnp.ndarray], tuple[jax.Array, jnp.ndarray, jnp.ndarray]],
    untempered_log_post: Callable[[jnp.ndarray], jnp.ndarray],
    steps_per_swap: int,
    n_swap_intervals: int,
    burn_in_steps: int,
) -> tuple[jnp.ndarray, float, jnp.ndarray, jnp.ndarray]:
    """Main PT loop: burn-in, MH updates, swaps, and diagnostics."""
    n_chains = theta0_all.shape[0]

    key, sub = jax.random.split(key)
    keys = jax.random.split(sub, n_chains)
    thetas = theta0_all

    swap_rates: list[jnp.ndarray] = []
    cold_trace: list[jnp.ndarray] = []
    acceptance_history: list[jnp.ndarray] = []

    n_pairs_total = n_chains - 1
    swap_counts = jnp.zeros(n_pairs_total)
    swap_totals = jnp.zeros(n_pairs_total)
    round_idx = 0

    # Burn-in with independent MH updates only
    if burn_in_steps > 0:
        for _ in range(burn_in_steps):
            keys, thetas, _ = mh_step_all_chains(keys, thetas)

    # Replica-exchange phase
    for _ in range(n_swap_intervals):
        accepts_per_swap = []
        for _ in range(steps_per_swap):
            keys, thetas, accepts = mh_step_all_chains(keys, thetas)
            accepts_per_swap.append(accepts)
            cold_trace.append(thetas[0])

        accepts_per_swap = jnp.stack(accepts_per_swap)
        mean_accepts_per_chain = jnp.mean(accepts_per_swap, axis=0)
        acceptance_history.append(mean_accepts_per_chain)

        key, subkey = jax.random.split(key)
        thetas, key, swap_accepts, pairs_i = _swap_samples_one_round(
            thetas=thetas,
            key=subkey,
            round_idx=round_idx,
            betas=betas,
            untempered_log_post=untempered_log_post,
        )
        swap_rates.append(jnp.mean(swap_accepts))

        swap_counts = swap_counts.at[pairs_i].add(1)
        swap_totals = swap_totals.at[pairs_i].add(swap_accepts.astype(jnp.float32))
        round_idx += 1

    acceptance_history = jnp.stack(acceptance_history)
    mean_acceptance_per_chain = jnp.mean(acceptance_history, axis=0)
    swap_accept_means = swap_totals / (swap_counts + 1e-8)

    samples_cold = jnp.stack(cold_trace)
    mean_swap_rate = float(jnp.mean(jnp.array(swap_rates)))

    return samples_cold, mean_swap_rate, mean_acceptance_per_chain, swap_accept_means


def run_parallel_tempering(
    *,
    key: jax.Array,
    theta0_all: jnp.ndarray,
    prior: dict[str, Any],
    base_loglike_fn: Callable[[jnp.ndarray], jnp.ndarray],
    betas: jnp.ndarray,
    steps_per_swap: int = 10,
    n_swap_intervals: int = 2000,
    proposal_std: jnp.ndarray | None = None,
    burn_in_steps: int = 1000,
) -> tuple[jnp.ndarray, float, jnp.ndarray, jnp.ndarray]:
    """Parallel tempering MCMC with random-walk Metropolis and adjacent-pair swaps.

    Parameters
    ----------
    key : jax.Array
        PRNG key.
    theta0_all : jnp.ndarray
        Shape (n_chains, D), initial states per chain.
    prior : dict
        Prior dict for sample_from_prior; must have "type", "lower", "upper"
        and for gaussian "mean", "std".
    base_loglike_fn : callable
        Log-likelihood for untempered posterior (theta -> scalar).
    betas : jnp.ndarray
        Shape (n_chains,), inverse temperatures (cold to hot).
    steps_per_swap : int
        MH steps between swap attempts.
    n_swap_intervals : int
        Number of swap intervals (total cold samples = n_swap_intervals * steps_per_swap).
    proposal_std : jnp.ndarray or None
        Proposal standard deviation, shape (D,) or (n_chains, D). Required.
    burn_in_steps : int
        MH steps before first swap (no swaps during burn-in).

    Returns
    -------
    samples_cold : jnp.ndarray
        Shape (n_swap_intervals * steps_per_swap, D), cold chain trace.
    mean_swap_rate : float
        Overall swap acceptance rate.
    mean_acceptance_per_chain : jnp.ndarray
        Shape (n_chains,), per-chain MH acceptance rate.
    swap_accept_means : jnp.ndarray
        Shape (n_chains - 1,), per adjacent-pair swap acceptance rate.
    """
    n_chains, dim = theta0_all.shape
    if betas.shape[0] != n_chains:
        raise ValueError(
            f"betas length must equal number of chains: got {betas.shape[0]}, expected {n_chains}"
        )

    # Validate and prepare proposal_std
    if proposal_std is None:
        raise ValueError("proposal_std must be provided")

    proposal_std = jnp.asarray(proposal_std, dtype=jnp.float32)
    lower = jnp.asarray(prior["lower"], dtype=jnp.float32)
    upper = jnp.asarray(prior["upper"], dtype=jnp.float32)

    proposal_std_chain = _build_proposal_std_chain(
        proposal_std=proposal_std,
        betas=betas,
        n_chains=n_chains,
        dim=dim,
    )

    log_prior_fn = _build_log_prior_fn(prior=prior, lower=lower, upper=upper)

    def untempered_log_post(theta: jnp.ndarray) -> jnp.ndarray:
        return log_prior_fn(theta) + base_loglike_fn(theta)

    mh_step_all_chains = _make_mh_step_all_chains(
        log_prior_fn=log_prior_fn,
        base_loglike_fn=base_loglike_fn,
        betas=betas,
        proposal_std_chain=proposal_std_chain,
        lower=lower,
        upper=upper,
    )

    return _run_replica_exchange_loop(
        key=key,
        theta0_all=theta0_all,
        betas=betas,
        mh_step_all_chains=mh_step_all_chains,
        untempered_log_post=untempered_log_post,
        steps_per_swap=steps_per_swap,
        n_swap_intervals=n_swap_intervals,
        burn_in_steps=burn_in_steps,
    )
