"""Well-defined MCMC entrypoint for beam-parameter inference."""

from dataclasses import dataclass
from typing import Literal, Optional

import jax
import jax.numpy as jnp
from jax import lax
import numpy as np
from jax.scipy.special import gammaln

from beam_infer.base_model import BeamForwardModel

ArrayLike = np.ndarray | jnp.ndarray
PriorType = Literal["uniform", "gaussian"]


@dataclass(frozen=True)
class MCMCConfig:
    """Configuration for Metropolis-Hastings parameter inference."""

    lower_bounds: ArrayLike
    upper_bounds: ArrayLike
    total_counts_obs: Optional[float] = None
    initial_theta: Optional[ArrayLike] = None
    prior_type: PriorType = "uniform"
    prior_mean: Optional[ArrayLike] = None
    prior_std: Optional[ArrayLike] = None
    num_samples: int = 4000
    burn_in: int = 1000
    thin: int = 2
    proposal_std: Optional[ArrayLike] = None
    seed: int = 0
    use_scan: bool = True  # Use lax.scan instead of Python loop for better performance (default: True)


@dataclass(frozen=True)
class MCMCSamples:
    """Raw MCMC samples before statistics computation."""

    samples: jnp.ndarray  # Shape: (num_samples, dim)
    acceptance_flags: jnp.ndarray  # Shape: (total_steps,)
    final_state: jnp.ndarray  # Shape: (dim,)


@dataclass(frozen=True)
class MCMCResult:
    """Normalized result object returned by MCMC entrypoint."""

    samples: np.ndarray
    posterior_mean: np.ndarray
    posterior_std: np.ndarray
    acceptance_rate: float
    final_theta: np.ndarray


def _as_jax_array(value: ArrayLike) -> jnp.ndarray:
    """Convert array-like value to JAX array."""
    return jnp.asarray(value, dtype=jnp.float32)


def _to_numpy(value: jnp.ndarray, dtype=np.float64) -> np.ndarray:
    """Convert JAX array to numpy at API boundary.
    
    Parameters
    ----------
    value : jnp.ndarray
        JAX array to convert
    dtype : numpy dtype, optional
        Target numpy dtype (default: np.float64)
        
    Returns
    -------
    np.ndarray
        Numpy array
    """
    return np.asarray(value, dtype=dtype)


def compute_mcmc_statistics(samples: MCMCSamples, total_steps: int) -> MCMCResult:
    """Compute posterior statistics from MCMC samples.
    
    Parameters
    ----------
    samples : MCMCSamples
        Raw MCMC samples
    total_steps : int
        Total number of MCMC steps (for acceptance rate calculation)
        
    Returns
    -------
    MCMCResult
        Computed statistics
    """
    return MCMCResult(
        samples=_to_numpy(samples.samples),
        posterior_mean=_to_numpy(jnp.mean(samples.samples, axis=0)),
        posterior_std=_to_numpy(jnp.std(samples.samples, axis=0)),
        acceptance_rate=float(jnp.mean(samples.acceptance_flags)),
        final_theta=_to_numpy(samples.final_state),
    )


def validate_mcmc_config(config: MCMCConfig, dim: int) -> None:
    """Validate MCMC configuration.
    
    Parameters
    ----------
    config : MCMCConfig
        MCMC configuration to validate
    dim : int
        Parameter dimension (number of parameters)
        
    Raises
    ------
    ValueError
        If configuration is invalid
    """
    if config.num_samples <= 0:
        raise ValueError(f"num_samples must be positive, got {config.num_samples}")
    if config.burn_in < 0:
        raise ValueError(f"burn_in must be non-negative, got {config.burn_in}")
    if config.thin <= 0:
        raise ValueError(f"thin must be positive, got {config.thin}")
    
    total_steps = config.burn_in + config.num_samples * config.thin
    if total_steps > 1_000_000:
        raise ValueError(
            f"Total steps ({total_steps}) exceeds reasonable limit. "
            f"Consider reducing num_samples or thin."
        )
    
    if config.proposal_std is not None:
        proposal_std_array = _as_jax_array(config.proposal_std)
        if proposal_std_array.shape != (dim,):
            raise ValueError(
                f"proposal_std shape {proposal_std_array.shape} "
                f"does not match parameter dimension {dim}"
            )


class BeamMCMCModel(BeamForwardModel):
    """Self-contained MCMC model (forward + priors + Poisson likelihood + MH sampling)."""

    @staticmethod
    def reflect(theta: jnp.ndarray, lower: jnp.ndarray, upper: jnp.ndarray):
        return jnp.where(theta < lower, 2.0 * lower - theta, jnp.where(theta > upper, 2.0 * upper - theta, theta))

    @staticmethod
    def log_uniform_prior(theta: jnp.ndarray, lower: jnp.ndarray, upper: jnp.ndarray):
        inside = jnp.all((theta >= lower) & (theta <= upper))
        return jnp.where(inside, 0.0, -jnp.inf)

    @staticmethod
    def log_gaussian_prior(theta: jnp.ndarray, mean: jnp.ndarray, std: jnp.ndarray):
        z = (theta - mean) / std
        return -0.5 * jnp.sum(z * z)

    def make_log_prior(self, config: MCMCConfig, lower: jnp.ndarray, upper: jnp.ndarray):
        if config.prior_type == "uniform":
            return lambda th: self.log_uniform_prior(th, lower, upper)

        if config.prior_mean is None or config.prior_std is None:
            raise ValueError("Gaussian prior requires prior_mean and prior_std.")

        mean = _as_jax_array(config.prior_mean)
        std = _as_jax_array(config.prior_std)
        if mean.shape != lower.shape or std.shape != lower.shape:
            raise ValueError("prior_mean/prior_std must match bounds shape.")

        return lambda th: self.log_gaussian_prior(th, mean, std) + self.log_uniform_prior(th, lower, upper)

    @staticmethod
    def build_initial_theta(config: MCMCConfig, lower: jnp.ndarray, upper: jnp.ndarray) -> jnp.ndarray:
        if config.initial_theta is not None:
            theta0 = _as_jax_array(config.initial_theta)
            if theta0.shape != lower.shape:
                raise ValueError("initial_theta must match bounds shape.")
            return theta0

        if config.prior_type == "gaussian" and config.prior_mean is not None:
            theta0 = _as_jax_array(config.prior_mean)
        else:
            theta0 = 0.5 * (lower + upper)

        return jnp.clip(theta0, lower, upper)

    def log_poisson_likelihood(
        self,
        theta: jnp.ndarray,
        i_obs: jnp.ndarray,
        x: jnp.ndarray,
        y: jnp.ndarray,
        t: jnp.ndarray,
        total_counts_obs: float,
    ) -> jnp.ndarray:
        i_obs_safe = jnp.clip(i_obs, min=1e-8)
        i_sim = self.simulate_image(x, y, t, theta)
        i_sim = i_sim / jnp.sum(i_sim)
        lam = jnp.clip(total_counts_obs * i_sim, min=1e-8)
        return jnp.sum(i_obs_safe * jnp.log(lam) - lam - gammaln(i_obs_safe + 1.0))

    def _run_metropolis_hastings_sampling(
        self,
        i_obs: jnp.ndarray,
        x: jnp.ndarray,
        y: jnp.ndarray,
        t: jnp.ndarray,
        lower: jnp.ndarray,
        upper: jnp.ndarray,
        config: MCMCConfig,
    ) -> MCMCSamples:
        """Run MCMC sampling and return raw samples.
        
        This method performs the actual MCMC sampling and returns raw samples
        without computing statistics. Use compute_mcmc_statistics() to get
        final results.
        """
        dim = int(lower.shape[0])
        
        # Validate configuration early
        validate_mcmc_config(config, dim)
        
        theta = self.build_initial_theta(config, lower, upper)

        if config.proposal_std is None:
            proposal_std = 0.05 * (upper - lower)
        else:
            proposal_std = _as_jax_array(config.proposal_std)
            # Shape validation already done in validate_mcmc_config

        total_counts_obs = (
            float(config.total_counts_obs)
            if config.total_counts_obs is not None
            else float(jnp.sum(i_obs))
        )
        if total_counts_obs <= 0:
            raise ValueError("total_counts_obs must be positive.")

        log_prior = self.make_log_prior(config, lower, upper)

        def log_posterior(th: jnp.ndarray) -> jnp.ndarray:
            return log_prior(th) + self.log_poisson_likelihood(th, i_obs, x, y, t, total_counts_obs)

        logp = log_posterior(theta)
        key = jax.random.PRNGKey(config.seed)
        total_steps = config.burn_in + config.num_samples * config.thin
        
        if config.use_scan:
            # Use lax.scan for efficient JIT compilation
            return self._run_metropolis_hastings_with_scan(
                theta, logp, key, proposal_std, lower, upper, log_posterior,
                config, total_steps, dim
            )
        else:
            # Use Python loop (original implementation)
            return self._run_metropolis_hastings_with_loop(
                theta, logp, key, proposal_std, lower, upper, log_posterior,
                config, total_steps, dim
            )

    def _run_metropolis_hastings_with_loop(
        self,
        theta: jnp.ndarray,
        logp: jnp.ndarray,
        key: jax.random.PRNGKey,
        proposal_std: jnp.ndarray,
        lower: jnp.ndarray,
        upper: jnp.ndarray,
        log_posterior: callable,
        config: MCMCConfig,
        total_steps: int,
        dim: int,
    ) -> MCMCSamples:
        """Run MCMC using Python loop (original implementation).
        
        Note: This approach is slower than lax.scan because:
        1. Python loop cannot be fully JIT compiled
        2. Each iteration has Python interpreter overhead
        3. Multiple GPU kernel launches instead of fused kernels
        4. Less efficient memory management
        
        Use this only for debugging or when exact reproducibility with old code is needed.
        """
        samples: list[jnp.ndarray] = []  # Keep as JAX arrays internally
        acceptance_flags: list[bool] = []  # Track acceptance for all steps

        for step in range(total_steps):
            key, noise_key, accept_key = jax.random.split(key, 3)
            noise = jax.random.normal(noise_key, shape=(dim,)) * proposal_std
            theta_prop = self.reflect(theta + noise, lower, upper)
            logp_prop = log_posterior(theta_prop)

            log_alpha = logp_prop - logp
            u = jnp.log(jax.random.uniform(accept_key))
            is_accept = u < log_alpha

            theta = jnp.where(is_accept, theta_prop, theta)
            logp = jnp.where(is_accept, logp_prop, logp)
            acceptance_flags.append(bool(is_accept))

            if step >= config.burn_in and (step - config.burn_in) % config.thin == 0:
                samples.append(theta)  # Keep as JAX array

        if not samples:
            raise RuntimeError("No MCMC samples collected. Check num_samples, burn_in, and thin.")

        # Stack as JAX arrays
        samples_jax = jnp.stack(samples, axis=0)
        acceptance_flags_jax = jnp.array(acceptance_flags, dtype=jnp.bool_)
        
        return MCMCSamples(
            samples=samples_jax,
            acceptance_flags=acceptance_flags_jax,
            final_state=theta,
        )

    def _run_metropolis_hastings_with_scan(
        self,
        theta: jnp.ndarray,
        logp: jnp.ndarray,
        key: jax.random.PRNGKey,
        proposal_std: jnp.ndarray,
        lower: jnp.ndarray,
        upper: jnp.ndarray,
        log_posterior: callable,
        config: MCMCConfig,
        total_steps: int,
        dim: int,
    ) -> MCMCSamples:
        """Run MCMC using lax.scan for efficient JIT compilation.
        
        This implementation is significantly faster than the Python loop version
        because:
        1. The entire loop is compiled into a single XLA computation graph
        2. Eliminates Python interpreter overhead during execution
        3. Enables XLA compiler optimizations (fusion, parallelization)
        4. Better GPU memory management and kernel fusion
        
        Typical speedups: 2-5x on CPU, 5-10x on GPU for large sample sizes.
        """
        
        def mcmc_step(carry: tuple, step_idx: jnp.ndarray) -> tuple[tuple, tuple]:
            """Single MCMC step for lax.scan."""
            theta_state, logp_state, key_state = carry
            
            # Split keys for this step
            key_state, noise_key, accept_key = jax.random.split(key_state, 3)
            
            # Propose new state
            noise = jax.random.normal(noise_key, shape=(dim,)) * proposal_std
            theta_prop = self.reflect(theta_state + noise, lower, upper)
            logp_prop = log_posterior(theta_prop)
            
            # Metropolis-Hastings acceptance
            log_alpha = logp_prop - logp_state
            u = jnp.log(jax.random.uniform(accept_key))
            is_accept = u < log_alpha
            
            # Update state
            theta_new = jnp.where(is_accept, theta_prop, theta_state)
            logp_new = jnp.where(is_accept, logp_prop, logp_state)
            
            # Determine if we should collect this sample (using JAX operations)
            # step_idx is already a JAX array, so use JAX comparisons
            after_burn_in = step_idx >= config.burn_in
            step_from_burn_in = step_idx - config.burn_in
            is_thin_step = (step_from_burn_in % config.thin) == 0
            should_collect = after_burn_in & is_thin_step
            
            # Return carry (state) and output (sample + acceptance flag + collect flag)
            new_carry = (theta_new, logp_new, key_state)
            output = (theta_new, is_accept, should_collect)
            
            return new_carry, output

        # Initialize carry: (theta, logp, key)
        initial_carry = (theta, logp, key)
        
        # Run scan
        final_carry, outputs = lax.scan(
            mcmc_step,
            initial_carry,
            jnp.arange(total_steps, dtype=jnp.int32)
        )
        
        theta_final, logp_final, key_final = final_carry
        samples_all, acceptance_flags_all, should_collect_flags = outputs
        
        # Filter samples: only keep those where should_collect is True
        collected_mask = should_collect_flags.astype(jnp.bool_)
        num_collected = int(jnp.sum(collected_mask))
        
        if num_collected == 0:
            raise RuntimeError("No MCMC samples collected. Check num_samples, burn_in, and thin.")
        
        # Collect samples using the mask
        samples_collected = samples_all[collected_mask]
        
        return MCMCSamples(
            samples=samples_collected,
            acceptance_flags=acceptance_flags_all.astype(jnp.bool_),
            final_state=theta_final,
        )

    def run_metropolis_hastings(
        self,
        i_obs: jnp.ndarray,
        x: jnp.ndarray,
        y: jnp.ndarray,
        t: jnp.ndarray,
        lower: jnp.ndarray,
        upper: jnp.ndarray,
        config: MCMCConfig,
    ) -> MCMCResult:
        """Run MCMC and return computed statistics.
        
        This is a convenience method that combines sampling and statistics computation.
        For more control, use _run_metropolis_hastings_sampling() and compute_mcmc_statistics() separately.
        """
        samples = self._run_metropolis_hastings_sampling(
            i_obs, x, y, t, lower, upper, config
        )
        total_steps = config.burn_in + config.num_samples * config.thin
        return compute_mcmc_statistics(samples, total_steps)


def run_mcmc_inference(
    image: ArrayLike,
    x_grid: ArrayLike,
    y_grid: ArrayLike,
    t_vals: ArrayLike,
    config: MCMCConfig,
) -> MCMCResult:
    """
    Entrypoint for stochastic inference with Metropolis-Hastings MCMC.

    Parameters
    ----------
    image:
        Observed 2D beam image (counts or normalized intensity).
    x_grid, y_grid:
        2D coordinate grids matching image shape.
    t_vals:
        1D time discretization used by the forward model.
    config:
        MCMC settings and prior definition.
    """
    i_obs = _as_jax_array(image)
    x = _as_jax_array(x_grid)
    y = _as_jax_array(y_grid)
    t = _as_jax_array(t_vals)
    lower = _as_jax_array(config.lower_bounds)
    upper = _as_jax_array(config.upper_bounds)

    if i_obs.shape != x.shape or i_obs.shape != y.shape:
        raise ValueError("image, x_grid, and y_grid must have the same 2D shape.")
    if t.ndim != 1:
        raise ValueError("t_vals must be a 1D array.")
    if lower.shape != upper.shape:
        raise ValueError("lower_bounds and upper_bounds must have matching shapes.")
    if jnp.any(lower > upper):
        raise ValueError("Bounds are invalid: lower bound exceeds upper bound.")

    model = BeamMCMCModel()
    return model.run_metropolis_hastings(
        i_obs=i_obs,
        x=x,
        y=y,
        t=t,
        lower=lower,
        upper=upper,
        config=config,
    )
