Bayesian MCP

"""Distribution definitions and utilities for the Bayesian engine."""

from typing import Dict, Any, Optional, Union
import numpy as np
import pymc as pm
import arviz as az


def create_distribution(
    distribution_type: str,
    name: str,
    **params
) -> Any:
    """Create a PyMC distribution based on type and parameters.
    
    Args:
        distribution_type: Type of distribution ('normal', 'beta', etc.)
        name: Name of the random variable
        **params: Distribution parameters
    
    Returns:
        PyMC distribution object
    
    Raises:
        ValueError: If distribution type is not supported
    """
    dist_map = {
        'normal': pm.Normal,
        'beta': pm.Beta,
        'gamma': pm.Gamma,
        'uniform': pm.Uniform,
        'exponential': pm.Exponential,
        'poisson': pm.Poisson,
        'bernoulli': pm.Bernoulli,
        'binomial': pm.Binomial,
        'categorical': pm.Categorical,
        'dirichlet': pm.Dirichlet,
        'multinomial': pm.Multinomial,
        'lognormal': pm.Lognormal,
        'studentt': pm.StudentT,
        'halfnormal': pm.HalfNormal,
        'weibull': pm.Weibull,
    }
    
    if distribution_type not in dist_map:
        raise ValueError(f"Unsupported distribution: {distribution_type}")
        
    return dist_map[distribution_type](name, **params)


def extract_posterior_stats(trace: az.InferenceData, var_name: str) -> Dict[str, Any]:
    """Extract summary statistics for a posterior variable.
    
    Args:
        trace: ArviZ InferenceData object containing posterior samples
        var_name: Name of the variable to extract
    
    Returns:
        Dictionary with posterior statistics
    """
    posterior = trace.posterior[var_name]
    
    # Get HDI (highest density interval)
    hdi = az.hdi(trace, var_names=[var_name])
    
    # Get ESS (effective sample size)
    ess = az.ess(trace, var_names=[var_name])
    
    # Get r_hat (convergence diagnostic)
    rhat = az.rhat(trace, var_names=[var_name])
    
    # Create result dictionary
    result = {
        "mean": float(posterior.mean().values),
        "std": float(posterior.std().values),
        "median": float(np.median(posterior.values)),
        "hdi_low": float(hdi[var_name].values[0]),
        "hdi_high": float(hdi[var_name].values[1]),
        "ess": float(ess[var_name].mean().values),
        "r_hat": float(rhat[var_name].mean().values)
    }
    
    return result