# Enhanced Multi-Modal Deep Latent Variable Models for Social Science Research
# ===============================================================================
#
# This module provides a comprehensive framework for multi-modal variational autoencoders (VAEs)
# designed for social science applications, particularly for analyzing political texts, speeches,
# voting patterns, and surveys. The framework supports various data types, fusion strategies,
# priors, and variational families.
#
# Key Features:
# ============
#
# 1. MULTIPLE DATA MODALITIES & LIKELIHOODS:
#    - Text embeddings: Gaussian likelihood for continuous embeddings (BERT, Doc2Vec, etc.)
#    - Discrete choices: Categorical likelihood for survey responses, vote choices
#    - Ordinal responses: Ordinal regression for Likert scales, ideology ratings
#    - Count data: Poisson, Negative Binomial for word counts, speech frequencies
#    - Binary data: Bernoulli for yes/no votes, presence/absence indicators
#    - Cross-entropy: For treating counts as pseudo-probabilities
#    - Multi-question support: Handle multiple categorical/ordinal questions simultaneously
#
# 2. FLEXIBLE FUSION STRATEGIES:
#    - Concatenate: Simple concatenation of all modalities (default, computationally efficient)
#    - Product of Experts (PoE): Multiply Gaussian posteriors for principled uncertainty combination
#    - Mixture of Experts (MoE): Three variants for handling heterogeneous modalities:
#      * moe_average: Simple averaging of expert posteriors
#      * moe_gating: Learned gating network for dynamic weighting and individual differences
#      * moe_learned: Fixed learned weights for expert combination
#    - MoE Weight Tracking: When using MoE strategies, weights are automatically tracked as
#      deterministic variables for posterior analysis and interpretability
#
# 3. CUSTOM AMORTIZED VARIATIONAL INFERENCE:
#    - Custom MultiModalGuide with modality-specific encoders for amortized inference
#    - Diagonal Gaussian: Independent normal distributions (computationally efficient)
#    - Full-rank Gaussian: Full covariance matrix for capturing correlations
#    - Inverse Autoregressive Flow (IAF): Flexible normalizing flows for complex posteriors
#    - Supports all fusion strategies (concatenation, PoE, MoE) with neural network encoders
#    - Efficient batch processing and memory management
#    - Automatic handling of modality-specific encoding and fusion
#
# 4. PRIOR DISTRIBUTIONS:
#    - Standard Gaussian: z ~ N(0, I) for basic latent structure
#    - Structured Gaussian: z ~ N(X @ gamma, Sigma) incorporating covariates/metadata
#    - Standard Logistic-Normal: z = LN(0, I) for simplex-valued latents
#    - Structured Logistic-Normal: z = ~ LN(X @ gamma, Sigma) for constrained latents
#
# 5. JOINT OUTCOME PREDICTION:
#    - Simultaneous training of VAE and outcome predictors
#    - Ideal point estimation with downstream prediction
#    - Support for multiple loss functions: MSE, cross-entropy, binary cross-entropy
#    - Configurable predictor architectures per outcome
#
# 6. COMPUTATIONAL EFFICIENCY:
#    - Mini-batch training for large datasets
#    - Memory-efficient batched inference
#    - Automatic batch size recommendations
#    - JAX-based implementation for GPU acceleration
#
# 7. POLITICAL SCIENCE APPLICATIONS:
#    - Ideal point estimation from speeches, voting records
#    - Multi-modal analysis of text + metadata + voting patterns
#    - Structured priors incorporating party affiliation, demographics
#    - Uncertainty quantification in political positioning
#    - Prediction of voting outcomes, policy preferences
#
# Usage Examples:
# ==============
#
# Basic Multi-Modal VAE:
# >>> modality_specs = {
# ...     "speech_embeddings": {"type": "gaussian", "dim": 768},
# ...     "voting_record": {"type": "bernoulli", "dim": 100},
# ...     "survey_responses": {"type": "ordinal", "dim": 5}
# ... }
# >>> model = DeepLatent(modality_specs, latent_dim=10)
# >>> model.fit(data)
#
# Structured Prior with Covariates:
# >>> model = create_structured_prior_vae(
# ...     modality_specs, covariate_dim=3, latent_dim=10
# ... )
# >>> model.fit(data, covariates=party_and_region_dummies)
#
# Product of Experts with Custom Guide:
# >>> model = create_multimodal_vae(
# ...     modality_specs, fusion_strategy="poe"
# ... )
#
# Joint Training with Outcome Prediction:
# >>> predictor_specs = {
# ...     "vote_choice": {"output_dim": 3, "loss_fn": "cross_entropy"},
# ...     "approval": {"output_dim": 1, "loss_fn": "mse"}
# ... }
# >>> model = create_joint_vae_predictor(
# ...     modality_specs, predictor_specs
# ... )
# >>> model.fit(data, outcomes=outcomes)
#
# Memory-Efficient Training:
# >>> batch_size = model.recommend_batch_size(data, target_memory_gb=8.0)
# >>> model.fit(data, batch_size=batch_size, num_steps=10000)
#
# Modality-Specific Encoder/Decoder Configurations:
# >>> # Configure different architectures for different modalities
# >>> modality_specs = {
# ...     "text_embeddings": {"type": "gaussian", "dim": 768},
# ...     "voting_record": {"type": "bernoulli", "dim": 100}
# ... }
# >>> encoder_specs = {
# ...     "text_embeddings": {
# ...         "hidden_dims": (512, 256, 128),  # Deep encoder for text
# ...         "activation": "swish",
# ...         "dropout_rate": 0.2
# ...     },
# ...     "voting_record": {
# ...         "hidden_dims": (64,),  # Shallow encoder for votes
# ...         "activation": "relu",
# ...         "dropout_rate": 0.1
# ...     }
# ... }
# >>> decoder_specs = {
# ...     "text_embeddings": {
# ...         "hidden_dims": (128, 256, 512),  # Mirror encoder
# ...         "activation": "swish",
# ...         "dropout_rate": 0.2
# ...     },
# ...     "voting_record": {
# ...         "hidden_dims": (64,),
# ...         "activation": "relu",
# ...         "dropout_rate": 0.1
# ...     }
# ... }
# >>> model = DeepLatent(
# ...     modality_specs,
# ...     latent_dim=10,
# ...     encoder_specs=encoder_specs,
# ...     decoder_specs=decoder_specs
# ... )
# >>> model.fit(data)
#
# Reasonable Defaults (automatically applied if not specified):
# >>> # All encoders, decoders, and predictors default to:
# >>> # - hidden_dims: (128,) - one hidden layer with 128 neurons
# >>> # - activation: "relu"
# >>> # - dropout_rate: 0.1
# >>> model = DeepLatent(modality_specs, latent_dim=10)  # Uses defaults
#
# Asymmetric Encoding/Decoding (Cross-Modal Generation):
# >>> # Encode votes and tweets, decode only speeches
# >>> modality_specs = {
# ...     "votes": {"type": "bernoulli", "dim": 100},
# ...     "tweets": {"type": "gaussian", "dim": 384},
# ...     "speeches": {"type": "gaussian", "dim": 768}
# ... }
# >>> encoder_specs = {
# ...     "votes": {"hidden_dims": (64,)},
# ...     "tweets": {"hidden_dims": (128,)}
# ...     # speeches NOT in encoder_specs → won't be encoded
# ... }
# >>> decoder_specs = {
# ...     "speeches": {"hidden_dims": (256, 512)}
# ...     # votes and tweets NOT in decoder_specs → won't be decoded
# ... }
# >>> model = DeepLatent(
# ...     modality_specs,
# ...     latent_dim=10,
# ...     encoder_specs=encoder_specs,  # Encode votes + tweets
# ...     decoder_specs=decoder_specs    # Decode speeches only
# ... )
# >>> # Training: votes + tweets → z → speeches
# >>> model.fit(data)
# >>> 
# >>> # Inference: predict speeches from votes + tweets
# >>> test_data = {"votes": test_votes, "tweets": test_tweets}
# >>> generated = model.generate_samples(num_samples=100, data=test_data)
# >>> print(generated["speeches"])  # Generated speeches from votes+tweets
#
# MoE Weight Analysis:
# >>> # For multi-modal MoE models
# >>> modality_specs = {
# ...     "text": {"type": "gaussian", "dim": 768},
# ...     "votes": {"type": "bernoulli", "dim": 100}
# ... }
# >>> model = DeepLatent(modality_specs, fusion_strategy="moe_gating")
# >>> model.fit(data)
# >>> moe_analysis = model.analyze_moe_weights(data)
# >>> model.print_moe_analysis(moe_analysis)  # Human-readable summary
# >>>
# >>> # Individual-specific analysis
# >>> individual_analysis = model.get_individual_moe_weights(data, individual_indices=[0, 1, 2])
# >>> print("Politician 0 weights:", individual_analysis["mean_weights"][0])
#
# Bayesian Linear Regression for Causal Inference:
# >>> # For causal inference with uncertainty quantification
# >>> predictor_specs = {
# ...     "policy_support": {
# ...         "type": "linear_regression",      # Specify linear regression
# ...         "outcome_type": "gaussian",       # gaussian, bernoulli, categorical, poisson
# ...         "output_dim": 1,
# ...         "beta_prior_scale": 1.0,          # Prior on coefficients
# ...         "sigma_prior_scale": 1.0,         # Prior on noise (for Gaussian)
# ...         "intercept_prior_scale": 1.0,     # Prior on intercept
# ...         "feature_names": ["ideology", "party_dem", "party_rep", "region_south"]
# ...     }
# ... }
# >>> model = DeepLatent(
# ...     modality_specs=modality_specs,
# ...     predictor_specs=predictor_specs,
# ...     latent_dim=10,
# ...     predictor_covariate_dim=3  # Additional covariates
# ... )
# >>> model.fit(data, outcomes=outcomes, covariates_predictor=covariates)
# >>>
# >>> # Extract coefficients with 95% credible intervals
# >>> coef_results = model.get_linear_coefficients(
# ...     data=data,
# ...     outcomes=outcomes,
# ...     covariates_predictor=covariates,
# ...     num_samples=2000,
# ...     ci_level=0.95
# ... )
# >>> model.print_linear_coefficients(coef_results, top_k=10)
# >>>
# >>> # Access raw posterior samples for further analysis
# >>> beta_samples = coef_results["policy_support"]["beta_samples"]  # (2000, input_dim, 1)
# >>> # Compute custom statistics, diagnostics, etc.
#
# Neural Network Predictors for Prediction Tasks:
# >>> # For prediction with deep neural networks
# >>> predictor_specs = {
# ...     "vote_choice": {
# ...         "type": "neural_network",  # or omit (default)
# ...         "output_dim": 3,
# ...         "hidden_dims": (256, 128),
# ...         "activation": "relu",
# ...         "loss_fn": "cross_entropy"
# ...     }
# ... }
# >>> model = DeepLatent(modality_specs, predictor_specs=predictor_specs)
# >>> model.fit(data, outcomes=outcomes)
#
# Classes and Functions:
# =====================
# - DeepLatent: Main model class with full functionality
# - ModalityLikelihood: Base class for likelihood functions
# - DecoderMLP, PredictorMLP: Neural network architectures
# - create_*_vae(): Convenience functions for common configurations
# - Utility functions for batching, alignment, coverage analysis
#
# References:
# ==========
# - Deep Latent Factor Models for Text and Audiovisual Data, Gauthier, Ash, Widmer (2025)
# - The Neural Ideal Point Model, Gauthier, Subtil, Widmer (2025)
#
# ===============================================================================

from __future__ import annotations
from typing import Dict, Tuple, List, Optional, Any
from abc import ABC, abstractmethod

import jax
import jax.numpy as jnp
from jax import random
import jax.scipy as jsp

import numpy as np

import numpyro
from numpyro import sample, plate
import numpyro.distributions as dist
from numpyro.infer import SVI, Trace_ELBO, Predictive
from numpyro.infer.elbo import ELBO
import numpyro.optim as optim

from flax import linen as nn

# ===============================================================================
# Custom Weighted ELBO for Flexible Loss Weighting
# ===============================================================================

class WeightedELBO(ELBO):
    """
    Custom ELBO with flexible loss weighting for KL divergence, reconstruction, and prediction.
    
    This allows for:
    - KL annealing (gradually increasing KL weight from 0 to 1)
    - Per-modality reconstruction weighting (emphasize certain modalities)
    - Per-outcome prediction weighting (balance multiple prediction tasks)
    
    Args:
        kl_weight: Weight for KL divergence term (default: 1.0). 
                   Use < 1.0 for KL annealing (e.g., start at 0.0, gradually increase to 1.0)
        reconstruction_weights: Optional per-modality reconstruction weights.
                                Dict mapping modality names to floats. Default: all 1.0
        predictor_weights: Optional per-outcome prediction weights.
                          Dict mapping outcome names to floats. Default: all 1.0
    
    Example:
        >>> # KL annealing
        >>> elbo = WeightedELBO(kl_weight=0.5)
        >>> 
        >>> # Emphasize certain modalities
        >>> elbo = WeightedELBO(
        ...     kl_weight=1.0,
        ...     reconstruction_weights={"text": 1.0, "votes": 2.0}
        ... )
        >>> 
        >>> # Weight prediction tasks
        >>> elbo = WeightedELBO(
        ...     kl_weight=1.0,
        ...     predictor_weights={"vote_choice": 3.0, "approval": 1.0}
        ... )
    """
    
    def __init__(self, 
                 kl_weight: float = 1.0,
                 reconstruction_weights: Optional[Dict[str, float]] = None,
                 predictor_weights: Optional[Dict[str, float]] = None):
        super().__init__()
        self.kl_weight = kl_weight
        self.reconstruction_weights = reconstruction_weights or {}
        self.predictor_weights = predictor_weights or {}
    
    def loss(self, rng_key, param_map, model, guide, *args, **kwargs):
        """
        Compute weighted ELBO loss.
        
        The weighted ELBO is:
            L = kl_weight * KL[q(z|x) || p(z)] 
                - Σ_i w_i * E_q[log p(x_i|z)]  (reconstruction)
                - Σ_k w_k * E_q[log p(y_k|z)]  (prediction)
        
        Where w_i and w_k are the reconstruction and predictor weights respectively.
        """
        # Get model and guide traces
        model_trace = numpyro.handlers.trace(
            numpyro.handlers.substitute(
                numpyro.handlers.seed(model, rng_key), 
                param_map
            )
        ).get_trace(*args, **kwargs)
        
        guide_trace = numpyro.handlers.trace(
            numpyro.handlers.substitute(
                numpyro.handlers.seed(guide, rng_key), 
                param_map
            )
        ).get_trace(*args, **kwargs)
        
        elbo = 0.0
        
        # 1. Compute KL divergence term (weighted)
        for name in guide_trace:
            site = guide_trace[name]
            if site["type"] == "sample" and not site.get("is_observed", False):
                if name in model_trace:
                    model_site = model_trace[name]
                    # KL[q(z) || p(z)]
                    log_prob_guide = site["fn"].log_prob(site["value"])
                    log_prob_model = model_site["fn"].log_prob(site["value"])
                    kl = log_prob_guide - log_prob_model
                    # Apply KL weight
                    elbo -= self.kl_weight * kl.sum()
        
        # 2. Compute reconstruction and prediction likelihoods (per-site weights)
        for name in model_trace:
            site = model_trace[name]
            if site["type"] == "sample" and site.get("is_observed", False):
                log_prob = site["fn"].log_prob(site["value"])
                
                # Determine weight based on site name
                weight = 1.0
                
                # Check if it's a modality observation (starts with "obs_" but not "obs_outcome_")
                if name.startswith("obs_") and not name.startswith("obs_outcome_"):
                    mod_name = name[4:]  # Remove "obs_" prefix
                    weight = self.reconstruction_weights.get(mod_name, 1.0)
                
                # Check if it's a predictor outcome (starts with "obs_outcome_")
                elif name.startswith("obs_outcome_"):
                    pred_name = name[12:]  # Remove "obs_outcome_" prefix
                    weight = self.predictor_weights.get(pred_name, 1.0)
                
                # Apply weight to log-likelihood
                elbo += weight * log_prob.sum()
        
        # Return negative ELBO (for minimization)
        return -elbo

# ===============================================================================
# Inverse Autoregressive Flow implementation
# ===============================================================================

class MaskedLinear(nn.Module):
    """Masked linear layer for autoregressive flows."""
    features: int
    mask: jnp.ndarray

    @nn.compact
    def __call__(self, x):
        kernel = self.param('kernel', nn.initializers.normal(0.01), (x.shape[-1], self.features))
        bias = self.param('bias', nn.initializers.zeros, (self.features,))
        masked_kernel = kernel * self.mask
        return jnp.dot(x, masked_kernel) + bias

class IAFLayer(nn.Module):
    """Single Inverse Autoregressive Flow layer."""
    latent_dim: int
    hidden_dim: int = 64

    def setup(self):
        # Create autoregressive mask
        self.mask = self._create_mask()

    def _create_mask(self):
        """Create autoregressive mask for IAF."""
        latent_dim = self.latent_dim
        mask = jnp.tril(jnp.ones((latent_dim, latent_dim)), k=-1)
        return mask

    @nn.compact
    def __call__(self, z, context=None):
        # Autoregressive network for mean and log_scale
        h = z
        if context is not None:
            h = jnp.concatenate([h, context], axis=-1)

        # Hidden layers with masking
        h = MaskedLinear(features=self.hidden_dim, mask=jnp.ones((h.shape[-1], self.hidden_dim)))(h)
        h = nn.tanh(h)
        h = MaskedLinear(features=self.hidden_dim, mask=jnp.ones((self.hidden_dim, self.hidden_dim)))(h)
        h = nn.tanh(h)

        # Output layer for shift and log_scale
        out = MaskedLinear(features=2 * self.latent_dim, mask=jnp.ones((self.hidden_dim, 2 * self.latent_dim)))(h)
        shift, log_scale = jnp.split(out, 2, axis=-1)

        # Apply autoregressive mask to shift and log_scale
        shift = shift * self.mask
        log_scale = log_scale * self.mask

        # Transform: z_new = z * exp(log_scale) + shift
        z_new = z * jnp.exp(log_scale) + shift

        # Log determinant of Jacobian
        log_det_jacobian = jnp.sum(log_scale, axis=-1)

        return z_new, log_det_jacobian

class IAFFlow(nn.Module):
    """Inverse Autoregressive Flow with multiple layers."""
    latent_dim: int
    num_flows: int = 4
    hidden_dim: int = 64

    @nn.compact
    def __call__(self, z_0, context=None):
        z = z_0
        total_log_det = 0.0

        for i in range(self.num_flows):
            layer = IAFLayer(latent_dim=self.latent_dim, hidden_dim=self.hidden_dim, name=f'iaf_layer_{i}')
            z, log_det = layer(z, context)
            total_log_det += log_det

            # Permute dimensions (except for last layer)
            if i < self.num_flows - 1:
                perm = jnp.arange(self.latent_dim)[::-1]  # Simple reversal permutation
                z = z[..., perm]

        return z, total_log_det

# -----------------------------------------------------------------------------
# Batching utilities
# -----------------------------------------------------------------------------

def create_batches(data: Dict[str, jnp.ndarray],
                   covariates_prior: Optional[jnp.ndarray] = None,
                   covariates_decoder: Optional[jnp.ndarray] = None,
                   covariates_predictor: Optional[jnp.ndarray] = None,
                   outcomes: Optional[Dict[str, jnp.ndarray]] = None,
                   batch_size: int = 32,
                   shuffle: bool = True,
                   seed: int = 0) -> List[Tuple[Dict[str, jnp.ndarray], Optional[jnp.ndarray], Optional[jnp.ndarray], Optional[jnp.ndarray], Optional[Dict[str, jnp.ndarray]]]]:
    """
    Create batches from data with optional shuffling and covariate/outcome tracking.

    Args:
        data: Dictionary of modality data
        covariates_prior: Optional covariate matrix for structured prior
        covariates_decoder: Optional covariate matrix for decoder
        covariates_predictor: Optional covariate matrix for outcome prediction
        outcomes: Optional outcome dictionary
        batch_size: Size of each batch
        shuffle: Whether to shuffle data
        seed: Random seed for shuffling

    Returns:
        List of (batch_data, batch_covariates_prior, batch_covariates_decoder, batch_covariates_predictor, batch_outcomes) tuples
    """
    # Get number of samples from first modality
    n_samples = next(iter(data.values())).shape[0]

    # Create indices
    indices = jnp.arange(n_samples)
    if shuffle:
        rng_key = random.PRNGKey(seed)
        indices = random.permutation(rng_key, indices)

    # Create batches
    batches = []
    for start_idx in range(0, n_samples, batch_size):
        end_idx = min(start_idx + batch_size, n_samples)
        batch_indices = indices[start_idx:end_idx]

        # Create batch data
        batch_data = {}
        for mod_name, mod_data in data.items():
            batch_data[mod_name] = mod_data[batch_indices]

        # Create batch covariates if provided
        batch_covariates_prior = None
        if covariates_prior is not None:
            batch_covariates_prior = covariates_prior[batch_indices]
            
        batch_covariates_decoder = None
        if covariates_decoder is not None:
            batch_covariates_decoder = covariates_decoder[batch_indices]
            
        batch_covariates_predictor = None
        if covariates_predictor is not None:
            batch_covariates_predictor = covariates_predictor[batch_indices]

        # Create batch outcomes if provided
        batch_outcomes = None
        if outcomes is not None:
            batch_outcomes = {}
            for outcome_name, outcome_data in outcomes.items():
                batch_outcomes[outcome_name] = outcome_data[batch_indices]

        batches.append((batch_data, batch_covariates_prior, batch_covariates_decoder, batch_covariates_predictor, batch_outcomes))

    return batches


def get_batch_size_recommendation(data: Dict[str, jnp.ndarray],
                                  target_memory_gb: float = 8.0) -> int:
    """
    Recommend a batch size based on data size and target memory usage.

    Args:
        data: Dictionary of modality data
        target_memory_gb: Target memory usage in GB

    Returns:
        Recommended batch size
    """
    # Calculate approximate memory per sample (in bytes)
    memory_per_sample = 0
    n_samples = next(iter(data.values())).shape[0]

    for mod_data in data.values():
        # Assume float32 (4 bytes per element)
        memory_per_sample += mod_data.size / n_samples * 4

    # Add overhead for gradients, intermediate computations (factor of 4)
    memory_per_sample *= 4

    # Convert target memory to bytes
    target_memory_bytes = target_memory_gb * 1e9

    # Calculate recommended batch size
    recommended_batch_size = int(target_memory_bytes / memory_per_sample)

    # Ensure reasonable bounds
    recommended_batch_size = max(16, min(recommended_batch_size, 1024))

    return recommended_batch_size

# -----------------------------------------------------------------------------
# Base classes for modality-specific likelihoods
# -----------------------------------------------------------------------------

class ModalityLikelihood(ABC):
    """Abstract base class for modality-specific likelihood functions."""

    @abstractmethod
    def log_likelihood(self, predictions: jnp.ndarray, targets: jnp.ndarray, **kwargs) -> jnp.ndarray:
        """Compute log-likelihood of targets given predictions."""
        pass

    @abstractmethod
    def sample(self, predictions: jnp.ndarray, key: jax.random.PRNGKey, **kwargs) -> jnp.ndarray:
        """Sample from the likelihood given predictions."""
        pass


class GaussianLikelihood(ModalityLikelihood):
    """Gaussian likelihood for continuous embeddings."""

    def __init__(self, scale: float = 1.0, learnable_scale: bool = False):
        self.scale = scale
        self.learnable_scale = learnable_scale

    def log_likelihood(self, predictions: jnp.ndarray, targets: jnp.ndarray,
                       scale: Optional[jnp.ndarray] = None) -> jnp.ndarray:
        if scale is None:
            scale = self.scale
        return jnp.sum(dist.Normal(predictions, scale).log_prob(targets))

    def sample(self, predictions: jnp.ndarray, key: jax.random.PRNGKey,
               scale: Optional[jnp.ndarray] = None) -> jnp.ndarray:
        if scale is None:
            scale = self.scale
        return predictions + scale * random.normal(key, predictions.shape)


class CategoricalLikelihood(ModalityLikelihood):
    """Categorical likelihood for discrete choices."""

    def log_likelihood(self, logits: jnp.ndarray, targets: jnp.ndarray) -> jnp.ndarray:
        # targets should be one-hot encoded or class indices
        if targets.shape[-1] == logits.shape[-1]:  # one-hot
            return jnp.sum(targets * nn.log_softmax(logits))
        else:  # class indices
            return jnp.sum(dist.Categorical(logits=logits).log_prob(targets))

    def sample(self, logits: jnp.ndarray, key: jax.random.PRNGKey) -> jnp.ndarray:
        return dist.Categorical(logits=logits).sample(key)


class OrdinalLikelihood(ModalityLikelihood):
    """Ordinal likelihood for ordered discrete choices."""

    def __init__(self, num_classes: int):
        self.num_classes = num_classes

    def log_likelihood(self, eta: jnp.ndarray, targets: jnp.ndarray,
                       cutpoints: jnp.ndarray) -> jnp.ndarray:
        # Ordinal regression using cumulative probabilities
        # eta: linear predictor, cutpoints: ordered thresholds
        cum_probs = jax.nn.sigmoid(cutpoints[None, :] - eta[:, None])

        # Compute category probabilities
        probs = jnp.concatenate([
            cum_probs[:, :1],
            cum_probs[:, 1:] - cum_probs[:, :-1],
            1 - cum_probs[:, -1:]
        ], axis=1)

        # Clamp probabilities to avoid log(0)
        probs = jnp.clip(probs, 1e-10, 1.0)

        if targets.shape[-1] == self.num_classes:  # one-hot
            return jnp.sum(targets * jnp.log(probs))
        else:  # class indices
            return jnp.sum(jnp.log(probs[jnp.arange(len(targets)), targets]))

    def sample(self, eta: jnp.ndarray, key: jax.random.PRNGKey,
               cutpoints: jnp.ndarray) -> jnp.ndarray:
        cum_probs = jax.nn.sigmoid(cutpoints[None, :] - eta[:, None])
        probs = jnp.concatenate([
            cum_probs[:, :1],
            cum_probs[:, 1:] - cum_probs[:, :-1],
            1 - cum_probs[:, -1:]
        ], axis=1)
        return dist.Categorical(probs=probs).sample(key)


class PoissonLikelihood(ModalityLikelihood):
    """Poisson likelihood for count data."""

    def log_likelihood(self, log_rate: jnp.ndarray, targets: jnp.ndarray) -> jnp.ndarray:
        # Clip log_rate more aggressively to prevent overflow when exponentiated
        log_rate = jnp.clip(log_rate, -10.0, 10.0)
        rate = jnp.exp(log_rate)
        return jnp.sum(dist.Poisson(rate).log_prob(targets))

    def sample(self, log_rate: jnp.ndarray, key: jax.random.PRNGKey) -> jnp.ndarray:
        # Clip log_rate more aggressively to prevent overflow when exponentiated
        log_rate = jnp.clip(log_rate, -10.0, 10.0)
        rate = jnp.exp(log_rate)
        return dist.Poisson(rate).sample(key)


class NegativeBinomialLikelihood(ModalityLikelihood):
    """Negative Binomial likelihood for overdispersed count data."""

    def __init__(self, concentration_prior: float = 1.0):
        """
        Initialize Negative Binomial likelihood.

        Args:
            concentration_prior: Prior value for the concentration parameter.
                                Higher values lead to less overdispersion.
        """
        self.concentration_prior = concentration_prior

    def log_likelihood(self, predictions: jnp.ndarray, targets: jnp.ndarray,
                       concentration: Optional[jnp.ndarray] = None) -> jnp.ndarray:
        """
        Compute log-likelihood for negative binomial distribution.

        Args:
            predictions: Log mean parameters of shape (..., dim)
            targets: Count targets of shape (..., dim)
            concentration: Concentration parameter (inverse overdispersion).
                          If None, uses a learnable parameter.
        """
        # Clip log mean to prevent overflow
        log_rate = jnp.clip(predictions, -10.0, 10.0)
        rate = jnp.exp(log_rate)

        if concentration is None:
            # Use a fixed concentration parameter
            concentration = self.concentration_prior

        # Ensure concentration is positive
        concentration = jnp.abs(concentration) + 1e-6

        # Create negative binomial distribution
        # NegativeBinomial in numpyro is parameterized by total_count (concentration) and probs
        # We convert from rate parameterization: probs = concentration / (concentration + rate)
        probs = concentration / (concentration + rate + 1e-8)

        return jnp.sum(dist.NegativeBinomial(total_count=concentration, probs=probs).log_prob(targets))

    def sample(self, predictions: jnp.ndarray, key: jax.random.PRNGKey,
               concentration: Optional[jnp.ndarray] = None) -> jnp.ndarray:
        """Sample from negative binomial distribution."""
        # Clip log mean to prevent overflow
        log_rate = jnp.clip(predictions, -10.0, 10.0)
        rate = jnp.exp(log_rate)

        if concentration is None:
            concentration = self.concentration_prior

        # Ensure concentration is positive
        concentration = jnp.abs(concentration) + 1e-6

        # Convert to probs parameterization
        probs = concentration / (concentration + rate + 1e-8)

        return dist.NegativeBinomial(total_count=concentration, probs=probs).sample(key)


class CrossEntropyLikelihood(ModalityLikelihood):
    """
    Multinomial cross-entropy for count data (standard in neural topic models).
    
    This is the count-weighted cross-entropy used in AVITM, ProdLDA, Scholar, GSM, etc.
    The likelihood is: log p(x|z) = Σ_v c_v * log p_v
    where c_v are word counts and p_v are predicted probabilities.
    
    This is equivalent to the multinomial log-likelihood.
    """

    def __init__(self):
        """
        Initialize cross-entropy likelihood.

        """

    def log_likelihood(self, logits: jnp.ndarray, targets: jnp.ndarray) -> jnp.ndarray:
        """
        Compute count-weighted cross-entropy for bag-of-words data.

        Args:
            logits: Predicted logits of shape (..., dim)
            targets: Word counts of shape (..., dim) - NOT normalized!

        Returns:
            Log-likelihood: Σ_v c_v * log p_v (summed over all documents and words)
        """
        # Apply log-softmax to get log probabilities
        log_probs = jax.nn.log_softmax(logits, axis=-1)

        # Compute count-weighted cross-entropy: Σ c_v * log p_v
        # This is the standard multinomial log-likelihood for topic models
        return jnp.sum(targets * log_probs)  

    def sample(self, logits: jnp.ndarray, key: jax.random.PRNGKey) -> jnp.ndarray:
        """
        Sample from categorical distribution and return as counts.

        Note: This is somewhat artificial since cross-entropy typically doesn't
        define a proper generative model for counts.
        """
        # Sample from categorical distribution
        probs = jax.nn.softmax(logits, axis=-1)
        samples = dist.Categorical(probs=probs).sample(key)

        # Convert categorical samples to one-hot (pseudo-counts)
        return jax.nn.one_hot(samples, num_classes=logits.shape[-1])


class BernoulliLikelihood(ModalityLikelihood):
    """Bernoulli likelihood for binary data."""

    def log_likelihood(self, logits: jnp.ndarray, targets: jnp.ndarray) -> jnp.ndarray:
        return jnp.sum(dist.Bernoulli(logits=logits).log_prob(targets))

    def sample(self, logits: jnp.ndarray, key: jax.random.PRNGKey) -> jnp.ndarray:
        return dist.Bernoulli(logits=logits).sample(key)


# -----------------------------------------------------------------------------
# Enhanced neural network architectures
# -----------------------------------------------------------------------------

class EncoderMLP(nn.Module):
    """Multi-layer perceptron encoder for modality-specific inference."""
    output_dim: int  # latent_dim * 2 for mean and log-variance
    hidden_dims: Tuple[int, ...] = (128, 128)
    activation: str = "relu"
    dropout_rate: float = 0.0

    @nn.compact
    def __call__(self, x: jnp.ndarray, training: bool = True) -> jnp.ndarray:
        h = x
        for i, hidden_dim in enumerate(self.hidden_dims):
            h = nn.Dense(hidden_dim, name=f"hidden_{i}")(h)
            if self.activation == "relu":
                h = nn.relu(h)
            elif self.activation == "tanh":
                h = nn.tanh(h)
            elif self.activation == "swish":
                h = nn.swish(h)

            if self.dropout_rate > 0:
                h = nn.Dropout(rate=self.dropout_rate, deterministic=not training)(h)

        # Output mean and log-variance
        out = nn.Dense(self.output_dim, name="output")(h)
        return out


class GatingNetwork(nn.Module):
    """Gating network for Mixture of Experts."""
    num_experts: int
    hidden_dim: int
    activation: str = "relu"

    @nn.compact
    def __call__(self, x: jnp.ndarray) -> jnp.ndarray:
        h = nn.Dense(self.hidden_dim, name="hidden")(x)
        if self.activation == "relu":
            h = nn.relu(h)
        elif self.activation == "tanh":
            h = nn.tanh(h)
        elif self.activation == "swish":
            h = nn.swish(h)

        logits = nn.Dense(self.num_experts, name="output")(h)
        weights = nn.softmax(logits)
        return weights


class MultiModalGuide:
    """Custom guide for multi-modal VAE with different variational families."""

    def __init__(self, model_instance):
        self.model_instance = model_instance
        self.modality_specs = model_instance.modality_specs
        self.latent_dim = model_instance.latent_dim
        self.fusion_strategy = model_instance.fusion_strategy
        self.variational_family = model_instance.variational_family
        self.encoder_hidden_dims = model_instance.encoder_hidden_dims
        self.activation = model_instance.activation
        self.dropout_rate = model_instance.dropout_rate
        self.gating_hidden_dim = model_instance.gating_hidden_dim
        self.use_structured_prior = model_instance.use_structured_prior
        self.prior_covariate_dim = model_instance.prior_covariate_dim
        self.encode_modalities = model_instance.encode_modalities

        # Initialize encoders ONLY for modalities that should be encoded
        self.encoders = {}
        for mod_name in self.encode_modalities:
            spec = self.modality_specs[mod_name]
            if self.variational_family == "diagonal":
                output_dim = 2 * self.latent_dim  # mean and log-variance
            elif self.variational_family == "full_rank":
                # For full rank, we need mean + scale_tril (lower triangular)
                output_dim = self.latent_dim + (self.latent_dim * (self.latent_dim + 1)) // 2
            elif self.variational_family == "iaf":
                output_dim = 2 * self.latent_dim  # mean and log-variance for base distribution
            else:
                raise ValueError(f"Unsupported variational family: {self.variational_family}")

            # Get modality-specific encoder configuration
            encoder_config = model_instance._get_encoder_config(mod_name)

            self.encoders[mod_name] = EncoderMLP(
                output_dim=output_dim,
                hidden_dims=encoder_config["hidden_dims"],
                activation=encoder_config["activation"],
                dropout_rate=encoder_config["dropout_rate"]
            )

        # Initialize gating network if needed
        if self.fusion_strategy == "moe_gating":
            self.gating_network = GatingNetwork(
                num_experts=len(self.encode_modalities),  # Use encode_modalities count
                hidden_dim=self.gating_hidden_dim,
                activation=self.activation
            )

        # Initialize IAF flow if needed
        if self.variational_family == "iaf":
            self.iaf_flow = IAFFlow(
                latent_dim=self.latent_dim,
                num_flows=4,
                hidden_dim=max(32, self.latent_dim)
            )

    def __call__(self, data: Dict[str, jnp.ndarray] = None,
                 covariates_prior: Optional[jnp.ndarray] = None,
                 covariates_decoder: Optional[jnp.ndarray] = None,
                 covariates_predictor: Optional[jnp.ndarray] = None,
                 outcomes: Optional[Dict[str, jnp.ndarray]] = None):
        """Guide function for variational inference with MoE fusion."""

        if data is None:
            return

        # Get batch size
        N = next(iter(data.values())).shape[0]

        # Encode each modality separately - ONLY encode modalities in encode_modalities
        modality_distributions = []
        modality_samples = []

        for mod_name, mod_data in data.items():
            # Skip modalities not in encode_modalities list
            if mod_name not in self.encode_modalities:
                continue
                
            if mod_name not in self.modality_specs:
                continue

            # Create encoder for this modality
            encoder_def = self.encoders[mod_name]

            def encoder_init(rng_key, input_shape):
                dummy_input = jnp.zeros(input_shape)
                variables = encoder_def.init(rng_key, dummy_input, training=False)
                params = variables["params"]
                out = encoder_def.apply({"params": params}, dummy_input, training=False)
                return out, params  # <-- two outputs

            def encoder_apply(params, x_in):
                return encoder_def.apply({"params": params}, x_in, training=False)

            encoder_fn = numpyro.module(
                f"encoder_{mod_name}",
                (encoder_init, encoder_apply),
                input_shape=(mod_data.shape[-1],)
            )

            # Apply encoder
            encoder_out = jax.vmap(encoder_fn, in_axes=0)(mod_data)

            # Process encoder output based on variational family
            if self.variational_family == "diagonal":
                mu, log_sigma = jnp.split(encoder_out, 2, axis=-1)
                # Clip log_sigma to prevent numerical instability
                log_sigma = jnp.clip(log_sigma, -10.0, 10.0)
                modality_distributions.append((mu, log_sigma))
                # Sample from this modality's posterior
                z_sample = sample(f"z_{mod_name}",
                                dist.Normal(mu, jnp.exp(log_sigma)).to_event(1))
            elif self.variational_family == "full_rank":
                # Split output into mean and scale_tril parameters
                mu = encoder_out[:, :self.latent_dim]
                scale_tril_vec = encoder_out[:, self.latent_dim:]

                # Reconstruct lower triangular matrix from vector
                def vec_to_tril(vec, dim):
                    tril = jnp.zeros((dim, dim))
                    tril_indices = jnp.tril_indices(dim)
                    tril = tril.at[tril_indices].set(vec)
                    # Ensure positive diagonal elements
                    tril = tril.at[jnp.diag_indices(dim)].set(jnp.exp(jnp.diag(tril)))
                    return tril

                scale_tril = jax.vmap(lambda x: vec_to_tril(x, self.latent_dim))(scale_tril_vec)
                modality_distributions.append((mu, scale_tril))
                # Sample from multivariate normal
                z_sample = sample(f"z_{mod_name}",
                                dist.MultivariateNormal(mu, scale_tril=scale_tril))
            elif self.variational_family == "iaf":
                mu, log_sigma = jnp.split(encoder_out, 2, axis=-1)
                # Clip log_sigma to prevent numerical instability
                log_sigma = jnp.clip(log_sigma, -10.0, 10.0)
                modality_distributions.append((mu, log_sigma))

                # Sample from base distribution
                z_base = sample(f"z_base_{mod_name}",
                              dist.Normal(mu, jnp.exp(log_sigma)).to_event(1))

                # Apply IAF transformation
                iaf_def = self.iaf_flow

                def iaf_init(rng_key, input_shape):
                    dummy_input = jnp.zeros(input_shape)
                    variables = iaf_def.init(rng_key, dummy_input)
                    params = variables["params"]
                    out, log_det = iaf_def.apply({"params": params}, dummy_input)
                    return out, params

                def iaf_apply(params, z_in):
                    z_out, log_det = iaf_def.apply({"params": params}, z_in)
                    return z_out, log_det

                iaf_fn = numpyro.module(
                    f"iaf_{mod_name}",
                    (iaf_init, iaf_apply),
                    input_shape=(self.latent_dim,)
                )

                # Apply IAF transformation
                z_sample, log_det = jax.vmap(iaf_fn, in_axes=0)(z_base)
                # Store log determinant for ELBO computation
                numpyro.deterministic(f"iaf_log_det_{mod_name}", log_det)
            else:
                raise ValueError(f"Unsupported variational family: {self.variational_family}")

            modality_samples.append(z_sample)

        # Apply fusion strategy
        if len(modality_distributions) == 1:
            # Single modality case
            if self.variational_family == "diagonal":
                final_mu, final_log_sigma = modality_distributions[0]
                final_z = sample("z", dist.Normal(final_mu, jnp.exp(final_log_sigma)).to_event(1))
            elif self.variational_family == "full_rank":
                final_mu, final_scale_tril = modality_distributions[0]
                final_z = sample("z", dist.MultivariateNormal(final_mu, scale_tril=final_scale_tril))
            elif self.variational_family == "iaf":
                # For IAF, the final_z is already the transformed sample
                final_z = modality_samples[0]
            moe_weights = jnp.ones((N, 1))

        elif self.fusion_strategy == "concatenate":
            # Simple concatenation - use average of modality posteriors
            if self.variational_family == "diagonal":
                final_mu = jnp.mean(jnp.stack([mu for mu, _ in modality_distributions], axis=0), axis=0)
                final_log_sigma = jnp.mean(jnp.stack([log_sigma for _, log_sigma in modality_distributions], axis=0), axis=0)
                final_z = sample("z", dist.Normal(final_mu, jnp.exp(final_log_sigma)).to_event(1))
            elif self.variational_family in ["full_rank", "iaf"]:
                # For full_rank and IAF, use simple average of samples
                sample_stack = jnp.stack(modality_samples, axis=0)
                averaged_samples = jnp.mean(sample_stack, axis=0)
                final_z = sample("z", dist.Delta(averaged_samples).to_event(1))
            moe_weights = jnp.ones((N, len(modality_distributions))) / len(modality_distributions)

        elif self.fusion_strategy == "poe":
            # Product of Experts - only works well with diagonal Gaussian
            if self.variational_family == "diagonal":
                final_mu, final_log_sigma = self.model_instance._product_of_experts(modality_distributions)
                final_z = sample("z", dist.Normal(final_mu, jnp.exp(final_log_sigma)).to_event(1))
                # For PoE, weights are the normalized precisions
                variances = [jnp.exp(2 * log_sigma) for _, log_sigma in modality_distributions]
                precisions = [1.0 / (var + 1e-8) for var in variances]
                total_precision = sum(precisions)
                moe_weights = jnp.stack([prec / (total_precision + 1e-8) for prec in precisions], axis=-1)
            else:
                # For other variational families, fall back to sample averaging
                sample_stack = jnp.stack(modality_samples, axis=0)
                averaged_samples = jnp.mean(sample_stack, axis=0)
                final_z = sample("z", dist.Delta(averaged_samples).to_event(1))
                moe_weights = jnp.ones((N, len(modality_distributions))) / len(modality_distributions)

        elif self.fusion_strategy.startswith("moe"):
            # Mixture of Experts strategies
            strategy_type = self.fusion_strategy.split("_", 1)[1]  # "average", "gating", "learned"

            # Compute mixture weights
            if strategy_type == "gating":
                # Use gating network
                gating_def = self.gating_network

                def gating_init(rng_key, input_shape):
                    dummy_input = jnp.zeros(input_shape)
                    variables = gating_def.init(rng_key, dummy_input)
                    params = variables["params"]
                    out = gating_def.apply({"params": params}, dummy_input)
                    return out, params

                def gating_apply(params, x_in):
                    return gating_def.apply({"params": params}, x_in)

                # Prepare gating input based on variational family
                if self.variational_family == "diagonal":
                    gate_input = jnp.concatenate([jnp.concatenate([mu, log_sigma], axis=-1)
                                                for mu, log_sigma in modality_distributions], axis=-1)
                elif self.variational_family == "full_rank":
                    # For full rank, use mean and first few elements of scale_tril
                    gate_input_parts = []
                    for mu, scale_tril in modality_distributions:
                        # Use mean and diagonal elements of scale_tril
                        diag_elements = jnp.diagonal(scale_tril, axis1=-2, axis2=-1)
                        gate_input_parts.append(jnp.concatenate([mu, diag_elements], axis=-1))
                    gate_input = jnp.concatenate(gate_input_parts, axis=-1)
                elif self.variational_family == "iaf":
                    # For IAF, use the base distribution parameters
                    gate_input = jnp.concatenate([jnp.concatenate([mu, log_sigma], axis=-1)
                                                for mu, log_sigma in modality_distributions], axis=-1)
                else:
                    raise ValueError(f"Unsupported variational family for gating: {self.variational_family}")

                gating_fn = numpyro.module(
                    "gating_network",
                    (gating_init, gating_apply),
                    input_shape=(gate_input.shape[-1],)
                )

                moe_weights = jax.vmap(gating_fn, in_axes=0)(gate_input)

            elif strategy_type == "learned":
                # Learned fixed weights
                num_modalities = len(modality_distributions)
                learned_weights_raw = sample("moe_learned_weights",
                                           dist.Normal(0, 1).expand([num_modalities]).to_event(1))
                learned_weights = jax.nn.softmax(learned_weights_raw)
                moe_weights = jnp.tile(learned_weights[None, :], (N, 1))

            else:  # "average"
                # Simple averaging
                num_modalities = len(modality_distributions)
                moe_weights = jnp.ones((N, num_modalities)) / num_modalities

            # Compute weighted combination of samples (works for all variational families)
            sample_stack = jnp.stack(modality_samples, axis=1)  # [batch, num_modalities, latent_dim]
            weighted_samples = jnp.sum(sample_stack * moe_weights[..., None], axis=1)

            # For MoE, we need to properly sample from the mixture distribution
            # Use a delta distribution centered at the weighted combination
            final_z = sample("z", dist.Delta(weighted_samples).to_event(1))

        else:
            raise ValueError(f"Unknown fusion strategy: {self.fusion_strategy}")

        # Store MoE weights and expert information as deterministic variables
        numpyro.deterministic("moe_weights", moe_weights)
        if len(modality_distributions) > 1:
            # Store expert information based on variational family
            if self.variational_family == "diagonal":
                expert_means = jnp.stack([mu for mu, _ in modality_distributions], axis=1)
                expert_log_stds = jnp.stack([log_sigma for _, log_sigma in modality_distributions], axis=1)
                numpyro.deterministic("moe_expert_means", expert_means)
                numpyro.deterministic("moe_expert_log_stds", expert_log_stds)
                numpyro.deterministic("moe_expert_stds", jnp.exp(expert_log_stds))
            elif self.variational_family == "full_rank":
                expert_means = jnp.stack([mu for mu, _ in modality_distributions], axis=1)
                expert_scale_trils = jnp.stack([scale_tril for _, scale_tril in modality_distributions], axis=1)
                numpyro.deterministic("moe_expert_means", expert_means)
                numpyro.deterministic("moe_expert_scale_trils", expert_scale_trils)
            elif self.variational_family == "iaf":
                expert_means = jnp.stack([mu for mu, _ in modality_distributions], axis=1)
                expert_log_stds = jnp.stack([log_sigma for _, log_sigma in modality_distributions], axis=1)
                numpyro.deterministic("moe_expert_base_means", expert_means)
                numpyro.deterministic("moe_expert_base_log_stds", expert_log_stds)
                # Store transformed samples
                expert_samples = jnp.stack(modality_samples, axis=1)
                numpyro.deterministic("moe_expert_transformed_samples", expert_samples)

        # Handle structured prior parameters in the guide if needed
        if self.use_structured_prior:
            if covariates_prior is None:
                raise ValueError("covariates_prior must be provided when use_structured_prior=True")

            # Make the prior parameters variational
            from numpyro import param
            from numpyro.distributions import constraints as C

            # gamma: full-rank variational Normal
            gamma_loc = param("gamma_loc", jnp.zeros((self.prior_covariate_dim, self.latent_dim)))
            gamma_scale = param("gamma_scale", 0.1 * jnp.ones((self.prior_covariate_dim, self.latent_dim)),
                               constraint=C.positive)
            gamma = sample("gamma", dist.Normal(gamma_loc, gamma_scale).to_event(2))

            # prior_scale: point-estimated positive vector
            prior_scale_q = param("prior_scale_q", jnp.ones(self.latent_dim), constraint=C.positive)
            scale_prior = sample("prior_scale", dist.Delta(prior_scale_q).to_event(1))

            # Handle correlation matrix: LKJ requires at least 2 dimensions
            if self.latent_dim > 1:
                # prior_correlation (Cholesky): point-estimated lower-tri factor
                L0 = jnp.eye(self.latent_dim)
                L_q = param("prior_correlation_L", L0, constraint=C.lower_cholesky)
                correlation_matrix = sample("prior_correlation", dist.Delta(L_q).to_event(2))
            else:
                # For 1D latent space, no correlation matrix needed (just variance)
                pass

        return final_z


class DecoderMLP(nn.Module):
    """Multi-layer perceptron decoder with modality-specific outputs."""
    output_dim: int
    hidden_dims: Tuple[int, ...] = (128,)
    latent_dim: int = 1
    activation: str = "relu"
    dropout_rate: float = 0.0
    modality_type: str = "gaussian"  # gaussian, categorical, poisson, etc.

    @nn.compact
    def __call__(self, theta: jnp.ndarray, training: bool = True) -> jnp.ndarray:
        h = theta
            
        for i, hidden_dim in enumerate(self.hidden_dims):
            h = nn.Dense(hidden_dim, name=f"hidden_{i}")(h)
            if self.activation == "relu":
                h = nn.relu(h)
            elif self.activation == "tanh":
                h = nn.tanh(h)
            elif self.activation == "swish":
                h = nn.swish(h)

            if self.dropout_rate > 0:
                h = nn.Dropout(rate=self.dropout_rate, deterministic=not training)(h)

        # Modality-specific output layers
        if self.modality_type == "gaussian":
            # Output mean and log-variance for Gaussian
            out = nn.Dense(2 * self.output_dim, name="output")(h)
            return out
        elif self.modality_type == "categorical":
            # Output logits for categorical
            logits = nn.Dense(self.output_dim, name="output")(h)
            return logits
        elif self.modality_type == "categorical_multi":
            # Output logits for multiple categorical questions
            # output_dim should be the total number of logits across all questions
            logits = nn.Dense(self.output_dim, name="output")(h)
            return logits
        elif self.modality_type == "poisson":
            # Output log-rate for Poisson
            log_rate = nn.Dense(self.output_dim, name="output")(h)
            return log_rate
        elif self.modality_type == "bernoulli":
            # Output logits for Bernoulli
            logits = nn.Dense(self.output_dim, name="output")(h)
            return logits
        elif self.modality_type == "ordinal":
            # Output linear predictor for ordinal regression
            eta = nn.Dense(self.output_dim, name="output")(h)
            return eta
        elif self.modality_type == "ordinal_multi":
            # Output linear predictor for multiple ordinal questions
            eta = nn.Dense(self.output_dim, name="output")(h)
            return eta
        else:
            # Default: simple linear output
            out = nn.Dense(self.output_dim, name="output")(h)
            return out


class PredictorMLP(nn.Module):
    """Multi-layer perceptron predictor with outcome-specific outputs (same pattern as DecoderMLP)."""
    output_dim: int
    hidden_dims: Tuple[int, ...] = (128, 64)
    latent_dim: int = 1  # input_dim for consistency with DecoderMLP
    activation: str = "relu"
    dropout_rate: float = 0.0
    modality_type: str = "gaussian"  # Same types as DecoderMLP

    @nn.compact
    def __call__(self, x: jnp.ndarray, training: bool = True) -> jnp.ndarray:
        """
        Forward pass of the predictor.

        Args:
            x: Input tensor of shape (batch_size, latent_dim) containing
               concatenated ideal points and covariates
            training: Whether the model is in training mode

        Returns:
            Predictions of shape (batch_size, output_dim)
        """
        h = x
        
        for i, hidden_dim in enumerate(self.hidden_dims):
            h = nn.Dense(hidden_dim, name=f"hidden_{i}")(h)
            if self.activation == "relu":
                h = nn.relu(h)
            elif self.activation == "tanh":
                h = nn.tanh(h)
            elif self.activation == "swish":
                h = nn.swish(h)

            if self.dropout_rate > 0:
                h = nn.Dropout(rate=self.dropout_rate, deterministic=not training)(h)

        # Outcome-specific output layers (same logic as DecoderMLP)
        if self.modality_type == "gaussian":
            # Output mean and log-variance for Gaussian
            out = nn.Dense(2 * self.output_dim, name="output")(h)
            return out
        elif self.modality_type == "categorical":
            # Output logits for categorical
            logits = nn.Dense(self.output_dim, name="output")(h)
            return logits
        elif self.modality_type == "categorical_multi":
            # Output logits for multiple categorical questions
            logits = nn.Dense(self.output_dim, name="output")(h)
            return logits
        elif self.modality_type == "poisson":
            # Output log-rate for Poisson
            log_rate = nn.Dense(self.output_dim, name="output")(h)
            return log_rate
        elif self.modality_type == "bernoulli":
            # Output logits for Bernoulli
            logits = nn.Dense(self.output_dim, name="output")(h)
            return logits
        elif self.modality_type == "ordinal":
            # Output linear predictor for ordinal regression
            eta = nn.Dense(self.output_dim, name="output")(h)
            return eta
        elif self.modality_type == "ordinal_multi":
            # Output linear predictor for multiple ordinal questions
            eta = nn.Dense(self.output_dim, name="output")(h)
            return eta
        elif self.modality_type == "negative_binomial":
            # Output log-rate for negative binomial
            log_rate = nn.Dense(self.output_dim, name="output")(h)
            return log_rate
        elif self.modality_type == "cross_entropy":
            # Output logits for cross-entropy
            logits = nn.Dense(self.output_dim, name="output")(h)
            return logits
        else:
            # Default: simple linear output
            out = nn.Dense(self.output_dim, name="output")(h)
            return out

# -----------------------------------------------------------------------------
# Enhanced Multi-Modal VAE with flexible likelihood support
# -----------------------------------------------------------------------------

class DeepLatent:
    """
    Multi-modal Variational Autoencoder supporting various data types and likelihoods.
    """

    def __init__(self,
                 modality_specs: Dict[str, Dict[str, Any]],
                 latent_dim: int = 10,
                 encoder_hidden_dims: Tuple[int, ...] = (128,),
                 decoder_hidden_dims: Tuple[int, ...] = (128,),
                 activation: str = "relu",
                 dropout_rate: float = 0.1,
                 encoder_specs: Optional[Dict[str, Dict[str, Any]]] = None,
                 decoder_specs: Optional[Dict[str, Dict[str, Any]]] = None,
                 fusion_strategy: str = "concatenate",
                 gating_hidden_dim: Optional[int] = None,
                 variational_family: str = "diagonal",
                 prior_type: str = "gaussian",
                 use_structured_prior: bool = False,
                 prior_covariate_dim: Optional[int] = None,
                 decoder_covariate_dim: int = 0,
                 predictor_covariate_dim: int = 0,
                 predictor_specs: Optional[Dict[str, Dict[str, Any]]] = None,
                 kl_weight: float = 1.0,
                 reconstruction_weights: Optional[Dict[str, float]] = None,
                 predictor_weights: Optional[Dict[str, float]] = None):
        """
        Initialize multi-modal VAE.

        Args:
            modality_specs: Dictionary specifying each modality. Format:
                {
                    "modality_name": {
                        "type": "gaussian" | "categorical" | "poisson" | etc.,
                        "dim": int,  # dimension of this modality
                        "likelihood_params": {...}  # modality-specific parameters
                    }
                }
            latent_dim: Dimension of shared latent space
            encoder_hidden_dims: Default hidden layer sizes for encoders (default: (128,))
            decoder_hidden_dims: Default hidden layer sizes for decoders (default: (128,))
            activation: Default activation function ("relu", "tanh", "swish") (default: "relu")
            dropout_rate: Default dropout rate for regularization (default: 0.1)
            encoder_specs: Optional dictionary for modality-specific encoder configurations. Format:
                {
                    "modality_name": {
                        "hidden_dims": Tuple[int, ...],  # Optional, overrides encoder_hidden_dims
                        "activation": str,  # Optional, overrides activation
                        "dropout_rate": float  # Optional, overrides dropout_rate
                    }
                }
                **IMPORTANT**: Only modalities listed in encoder_specs will be encoded (used for inference).
                If encoder_specs is None or empty, ALL modalities will be encoded with default configs.
            decoder_specs: Optional dictionary for modality-specific decoder configurations. Format:
                {
                    "modality_name": {
                        "hidden_dims": Tuple[int, ...],  # Optional, overrides decoder_hidden_dims
                        "activation": str,  # Optional, overrides activation
                        "dropout_rate": float  # Optional, overrides dropout_rate
                    }
                }
                **IMPORTANT**: Only modalities listed in decoder_specs will be decoded (reconstructed).
                If decoder_specs is None or empty, ALL modalities will be decoded with default configs.
            fusion_strategy: How to combine multiple modalities:
                - "concatenate": Concatenate all modalities and use shared encoder
                - "poe": Product of Experts - combine modality-specific posteriors
                - "moe_average": Mixture of Experts with simple averaging
                - "moe_gating": Mixture of Experts with learned gating network
                - "moe_learned": Mixture of Experts with learned fixed weights
            gating_hidden_dim: Hidden dimension for gating network (if using moe_gating)
            variational_family: Type of variational posterior family:
                - "diagonal": Diagonal Gaussian (default)
                - "full_rank": Full covariance Gaussian
                - "iaf": Inverse Autoregressive Flow
            prior_type: Type of prior distribution:
                - "gaussian": Standard or structured Gaussian prior
                - "logistic_normal": Standard or structured Logistic-Normal prior
            use_structured_prior: Whether to use structured prior with covariates
            prior_covariate_dim: Dimension of covariates for structured prior (required if use_structured_prior=True)
            decoder_covariate_dim: Dimension of covariates for decoder (concatenated with z, default 0)
            predictor_covariate_dim: Dimension of covariates for outcome predictors (concatenated with z, default 0)
            predictor_specs: Optional dictionary specifying outcome predictors for joint training. Format:
                {
                    "outcome_name": {
                        "output_dim": int,  # Number of outputs for this outcome
                        "hidden_dims": Tuple[int, ...],  # Optional, defaults to (128,)
                        "activation": str,  # Optional, defaults to "relu"
                        "dropout_rate": float,  # Optional, defaults to 0.1
                        "output_activation": str,  # Optional, e.g., "sigmoid", "softmax"
                        "loss_fn": str,  # "mse", "cross_entropy", "binary_cross_entropy"
                        "loss_weight": float  # Weight for this loss term (default: 1.0)
                    }
                }
            kl_weight: Weight for KL divergence term (default: 1.0). 
                       Use < 1.0 for KL annealing (e.g., start at 0.0, gradually increase to 1.0).
                       When 1.0 (default), standard Trace_ELBO is used for efficiency.
            reconstruction_weights: Optional per-modality reconstruction weights.
                                   Dict mapping modality names to floats. 
                                   If specified, triggers use of WeightedELBO.
                                   Example: {"text": 1.0, "votes": 2.0} emphasizes voting reconstruction.
            predictor_weights: Optional per-outcome prediction weights.
                              Dict mapping outcome names to floats.
                              If specified, triggers use of WeightedELBO.
                              Example: {"vote_choice": 3.0, "approval": 1.0} emphasizes classification.
        """
        self.modality_specs = modality_specs
        self.latent_dim = latent_dim
        self.encoder_hidden_dims = encoder_hidden_dims
        self.decoder_hidden_dims = decoder_hidden_dims
        self.activation = activation
        self.dropout_rate = dropout_rate
        self.encoder_specs = encoder_specs or {}
        self.decoder_specs = decoder_specs or {}
        
        # Determine which modalities to encode and decode based on specs
        # If specs are empty, encode/decode all modalities
        if self.encoder_specs:
            self.encode_modalities = list(self.encoder_specs.keys())
        else:
            self.encode_modalities = list(modality_specs.keys())
            
        if self.decoder_specs:
            self.decode_modalities = list(self.decoder_specs.keys())
        else:
            self.decode_modalities = list(modality_specs.keys())
        
        # Validate that specified modalities exist in modality_specs
        for mod in self.encode_modalities:
            if mod not in modality_specs:
                raise ValueError(f"Encoder modality '{mod}' not found in modality_specs")
        for mod in self.decode_modalities:
            if mod not in modality_specs:
                raise ValueError(f"Decoder modality '{mod}' not found in modality_specs")
        
        self.fusion_strategy = fusion_strategy
        self.gating_hidden_dim = gating_hidden_dim or len(self.encode_modalities)  # Use encode count
        self.variational_family = variational_family
        self.prior_type = prior_type
        self.use_structured_prior = use_structured_prior
        self.prior_covariate_dim = prior_covariate_dim
        self.decoder_covariate_dim = decoder_covariate_dim
        self.predictor_covariate_dim = predictor_covariate_dim
        self.predictor_specs = predictor_specs or {}
        
        # Loss weighting parameters
        self.kl_weight = kl_weight
        self.reconstruction_weights = reconstruction_weights or {}
        self.predictor_weights = predictor_weights or {}
        
        # Set default weights for modalities not specified
        for mod_name in modality_specs.keys():
            if mod_name not in self.reconstruction_weights:
                self.reconstruction_weights[mod_name] = 1.0
                
        # Set default weights for predictors not specified
        if self.predictor_specs:
            for pred_name in self.predictor_specs.keys():
                if pred_name not in self.predictor_weights:
                    self.predictor_weights[pred_name] = 1.0

        if use_structured_prior and prior_covariate_dim is None:
            raise ValueError("prior_covariate_dim must be specified when use_structured_prior=True")

        valid_strategies = ["concatenate", "poe", "moe_average", "moe_gating", "moe_learned"]
        if fusion_strategy not in valid_strategies:
            raise ValueError(f"fusion_strategy must be one of {valid_strategies}")

        valid_families = ["diagonal", "full_rank", "iaf"]
        if variational_family not in valid_families:
            raise ValueError(f"variational_family must be one of {valid_families}")

        valid_priors = ["gaussian", "logistic_normal"]
        if prior_type not in valid_priors:
            raise ValueError(f"prior_type must be one of {valid_priors}")

        # Initialize likelihood functions
        self.likelihoods = {}
        for mod_name, spec in modality_specs.items():
            self.likelihoods[mod_name] = self._create_likelihood(spec)

        # Initialize predictor MLPs if specified
        self.predictor_mlps = {}
        if self.predictor_specs:
            predictor_input_dim = latent_dim + predictor_covariate_dim
            for pred_name, pred_spec in self.predictor_specs.items():
                # Determine outcome type (with backward compatibility for loss_fn)
                outcome_type = pred_spec.get("type", None)
                if outcome_type is None and "loss_fn" in pred_spec:
                    loss_fn = pred_spec["loss_fn"]
                    if loss_fn == "mse":
                        outcome_type = "gaussian"
                    elif loss_fn == "cross_entropy":
                        outcome_type = "categorical"
                    elif loss_fn == "binary_cross_entropy":
                        outcome_type = "bernoulli"
                elif outcome_type is None:
                    outcome_type = "gaussian"  # Default
                    
                self.predictor_mlps[pred_name] = PredictorMLP(
                    output_dim=pred_spec["output_dim"],
                    hidden_dims=pred_spec.get("hidden_dims", (128,)),
                    latent_dim=predictor_input_dim,
                    activation=pred_spec.get("activation", "relu"),
                    dropout_rate=pred_spec.get("dropout_rate", 0.1),
                    modality_type=outcome_type
                )

        # Placeholder for learned weights if needed
        self._mixture_weights = None

        # Placeholder for guide
        self.auto_guide = None

    def _get_encoder_config(self, modality_name: str) -> Dict[str, Any]:
        """
        Get encoder configuration for a specific modality.
        
        Args:
            modality_name: Name of the modality
            
        Returns:
            Dictionary with 'hidden_dims', 'activation', 'dropout_rate'
        """
        if modality_name in self.encoder_specs:
            spec = self.encoder_specs[modality_name]
            return {
                "hidden_dims": spec.get("hidden_dims", self.encoder_hidden_dims),
                "activation": spec.get("activation", self.activation),
                "dropout_rate": spec.get("dropout_rate", self.dropout_rate)
            }
        else:
            return {
                "hidden_dims": self.encoder_hidden_dims,
                "activation": self.activation,
                "dropout_rate": self.dropout_rate
            }

    def _get_decoder_config(self, modality_name: str) -> Dict[str, Any]:
        """
        Get decoder configuration for a specific modality.
        
        Args:
            modality_name: Name of the modality
            
        Returns:
            Dictionary with 'hidden_dims', 'activation', 'dropout_rate'
        """
        if modality_name in self.decoder_specs:
            spec = self.decoder_specs[modality_name]
            return {
                "hidden_dims": spec.get("hidden_dims", self.decoder_hidden_dims),
                "activation": spec.get("activation", self.activation),
                "dropout_rate": spec.get("dropout_rate", self.dropout_rate)
            }
        else:
            return {
                "hidden_dims": self.decoder_hidden_dims,
                "activation": self.activation,
                "dropout_rate": self.dropout_rate
            }

    def _create_likelihood(self, spec: Dict[str, Any]) -> ModalityLikelihood:
        """Create appropriate likelihood function for a modality."""
        mod_type = spec["type"]
        params = spec.get("likelihood_params", {})

        if mod_type == "gaussian":
            return GaussianLikelihood(**params)
        elif mod_type == "categorical":
            return CategoricalLikelihood()
        elif mod_type == "categorical_multi":
            return CategoricalLikelihood()
        elif mod_type == "ordinal":
            return OrdinalLikelihood(num_classes=spec["dim"], **params)
        elif mod_type == "ordinal_multi":
            # For ordinal_multi, we don't pass categories_per_question to OrdinalLikelihood
            # since it's handled in the model sampling logic
            filtered_params = {k: v for k, v in params.items() if k != "categories_per_question"}
            return OrdinalLikelihood(num_classes=5, **filtered_params)  # Placeholder, actual num_classes handled per question
        elif mod_type == "poisson":
            return PoissonLikelihood(**params)
        elif mod_type == "negative_binomial":
            return NegativeBinomialLikelihood(**params)
        elif mod_type == "cross_entropy":
            return CrossEntropyLikelihood(**params)
        elif mod_type == "bernoulli":
            return BernoulliLikelihood(**params)
        else:
            raise ValueError(f"Unsupported modality type: {mod_type}")

    def _create_auto_guide(self):
        """Create the custom MultiModalGuide for variational inference."""
        # Always use custom guide for all cases
        return MultiModalGuide(self)

    def _product_of_experts(self, distributions: List[Tuple[jnp.ndarray, jnp.ndarray]]) -> Tuple[jnp.ndarray, jnp.ndarray]:
        """
        Combine multiple Gaussian distributions using Product of Experts.

        Args:
            distributions: List of (mu, log_sigma) tuples

        Returns:
            Combined (mu, log_sigma) tuple
        """
        if len(distributions) == 1:
            return distributions[0]

        eps = 1e-8
        precisions = []
        precision_weighted_means = []

        for mu, log_sigma in distributions:
            sigma = jnp.exp(log_sigma)
            precision = 1.0 / (sigma ** 2 + eps)
            precisions.append(precision)
            precision_weighted_means.append(mu * precision)

        combined_precision = sum(precisions)
        combined_mu = sum(precision_weighted_means) / (combined_precision + eps)
        combined_sigma = jnp.sqrt(1.0 / (combined_precision + eps))
        combined_log_sigma = jnp.log(combined_sigma)

        return combined_mu, combined_log_sigma

    def _mixture_of_experts_weights(self, distributions: List[Tuple[jnp.ndarray, jnp.ndarray]],
                                    strategy: str,
                                    learned_weights: Optional[jnp.ndarray] = None,
                                    gating_fn: Optional[callable] = None) -> jnp.ndarray:
        """
        Compute mixture weights for Mixture of Experts.

        Args:
            distributions: List of (mu, log_sigma) tuples
            strategy: "average", "gating", or "learned"
            learned_weights: optional fixed weights [num_modalities]
            gating_fn: optional gating function for learned gating

        Returns:
            Mixture weights of shape [batch_size, num_modalities]
        """
        batch_size = distributions[0][0].shape[0]
        num_modalities = len(distributions)

        if strategy == "average":
            return jnp.ones((batch_size, num_modalities)) / num_modalities

        elif strategy == "gating":
            # Concatenate modality-specific distribution parameters as input to gating network
            gate_input = jnp.concatenate([jnp.concatenate([mu, log_sigma], axis=-1)
                                          for mu, log_sigma in distributions], axis=-1)

            if gating_fn is not None:
                # Use the learned gating network
                weights = jax.vmap(gating_fn, in_axes=0)(gate_input)
                return weights
            else:
                # Fallback to simple heuristic-based gating if no gating function provided
                # Use inverse-variance weighting as a reasonable heuristic
                variances = [jnp.exp(2 * log_sigma) for _, log_sigma in distributions]
                inv_variances = [1.0 / (var + 1e-8) for var in variances]
                weights = jnp.stack(inv_variances, axis=-1)
                weights = weights / (jnp.sum(weights, axis=-1, keepdims=True) + 1e-8)
                return weights

        elif strategy == "learned":
            if learned_weights is not None:
                weights = jnp.tile(learned_weights[None, :], (batch_size, 1))
                return weights
            # fallback: inverse-variance weighting
            variances = [jnp.exp(2 * log_sigma) for _, log_sigma in distributions]
            inv_variances = [1.0 / (var + 1e-8) for var in variances]
            weights = jnp.stack(inv_variances, axis=-1)
            weights = weights / (jnp.sum(weights, axis=-1, keepdims=True) + 1e-8)
            return weights

        else:
            raise ValueError(f"Unknown MoE strategy: {strategy}")

    def _mixture_of_experts(self, distributions: List[Tuple[jnp.ndarray, jnp.ndarray]],
                            samples: List[jnp.ndarray],
                            strategy: str,
                            learned_weights: Optional[jnp.ndarray] = None,
                            gating_fn: Optional[callable] = None
                            ) -> Tuple[jnp.ndarray, jnp.ndarray, jnp.ndarray]:
        """
        Combine samples using Mixture of Experts.

        Args:
            distributions: List of (mu, log_sigma) tuples
            samples: List of sampled latent variables
            strategy: "average", "gating", or "learned"
            learned_weights: optional fixed weights [num_modalities]
            gating_fn: optional gating function for learned gating

        Returns:
            (combined_sample, combined_mu, combined_log_sigma)
        """
        weights = self._mixture_of_experts_weights(distributions, strategy, learned_weights, gating_fn)

        # Weighted combination of samples
        sample_stack = jnp.stack(samples, axis=1)  # [batch, num_modalities, latent_dim]
        combined_sample = jnp.sum(sample_stack * weights[..., None], axis=1)

        # Weighted combination of means
        mu_stack = jnp.stack([mu for mu, _ in distributions], axis=1)
        combined_mu = jnp.sum(mu_stack * weights[..., None], axis=1)

        # Weighted combination of log-std via arithmetic mean in log-space
        log_sigma_stack = jnp.stack([log_sigma for _, log_sigma in distributions], axis=1)
        combined_log_sigma = jnp.sum(log_sigma_stack * weights[..., None], axis=1)

        return combined_sample, combined_mu, combined_log_sigma

    def model(self, data: Dict[str, jnp.ndarray] = None,
              covariates_prior: Optional[jnp.ndarray] = None,
              covariates_decoder: Optional[jnp.ndarray] = None,
              covariates_predictor: Optional[jnp.ndarray] = None,
              outcomes: Optional[Dict[str, jnp.ndarray]] = None):
        """Generative model for multi-modal VAE with optional structured prior and outcome prediction."""
        # Get dimensions
        if data is not None and len(data) > 0:
            N = next(iter(data.values())).shape[0]
        else:
            N = 1  # For sampling

        # Structured or standard latent prior
        if self.use_structured_prior:
            if covariates_prior is None:
                raise ValueError("covariates_prior must be provided when use_structured_prior=True")

            # Prior parameters for structured prior (sampled globally, not per observation)
            # Regression coefficients gamma: shape (prior_covariate_dim, latent_dim)
            gamma = sample("gamma",
                          dist.Normal(0, 1).expand([self.prior_covariate_dim, self.latent_dim]).to_event(2))

            # Full covariance matrix sigma: shape (latent_dim, latent_dim)
            # Use LKJ prior for correlation matrix and separate scale parameters
            scale_prior = sample("prior_scale",
                               dist.HalfNormal(1).expand([self.latent_dim]).to_event(1))

            # Handle correlation matrix: LKJ requires at least 2 dimensions
            if self.latent_dim > 1:
                correlation_matrix = sample("prior_correlation",
                                          dist.LKJCholesky(self.latent_dim, concentration=2.0))
                # Construct covariance matrix from scale and correlation
                scale_diag = jnp.diag(scale_prior)
                covariance_chol = scale_diag @ correlation_matrix
            else:
                # For 1D latent space, covariance is just the variance (scalar)
                covariance_chol = jnp.diag(scale_prior)  # This creates a 1x1 matrix

            # Compute structured prior mean: X @ gamma (where X is covariates_prior)
            # covariates_prior: (N, prior_covariate_dim), gamma: (prior_covariate_dim, latent_dim)
            prior_mean = covariates_prior @ gamma  # (N, latent_dim)

            # Store deterministic quantities for monitoring
            numpyro.deterministic("prior_mean", prior_mean)
            numpyro.deterministic("covariance_matrix", covariance_chol @ covariance_chol.T)

            # Sample from structured prior
            with plate("samples", N):
                if self.prior_type == "gaussian":
                    # Use MultivariateNormal for both 1D and multi-dimensional cases
                    z = sample("z", dist.MultivariateNormal(prior_mean, scale_tril=covariance_chol))
                elif self.prior_type == "logistic_normal":
                    # Sample from multivariate normal in unconstrained space
                    z_unconstrained = sample("z",
                                           dist.MultivariateNormal(prior_mean, scale_tril=covariance_chol))
                    # Transform to simplex via softmax
                    z = jax.nn.softmax(z_unconstrained, axis=-1)
                    numpyro.deterministic("z_simplex", z)
                else:
                    raise ValueError(f"Unsupported prior_type: {self.prior_type}")

        else:
            # Standard isotropic prior
            with plate("samples", N):
                if self.prior_type == "gaussian":
                    z = sample("z", dist.Normal(0, 1).expand([self.latent_dim]).to_event(1))
                elif self.prior_type == "logistic_normal":
                    # Sample from standard multivariate normal and transform to simplex
                    z_unconstrained = sample("z",
                                           dist.Normal(0, 1).expand([self.latent_dim]).to_event(1))
                    # Transform to simplex via softmax
                    z = jax.nn.softmax(z_unconstrained, axis=-1)
                    numpyro.deterministic("z_simplex", z)
                else:
                    raise ValueError(f"Unsupported prior_type: {self.prior_type}")

        if data is not None:
            # Prepare decoder input: concatenate z with decoder covariates if provided
            if self.decoder_covariate_dim > 0 and covariates_decoder is not None:
                decoder_input = jnp.concatenate([z, covariates_decoder], axis=1)
                decoder_input_dim = self.latent_dim + self.decoder_covariate_dim
            else:
                decoder_input = z
                decoder_input_dim = self.latent_dim
            
            # Register decoders ONLY for modalities in decode_modalities
            for mod_name in self.decode_modalities:
                spec = self.modality_specs[mod_name]
                
                # Get modality-specific decoder configuration
                decoder_config = self._get_decoder_config(mod_name)
                
                decoder_def = DecoderMLP(
                    output_dim=spec["dim"],
                    hidden_dims=decoder_config["hidden_dims"],
                    latent_dim=decoder_input_dim,
                    activation=decoder_config["activation"],
                    dropout_rate=decoder_config["dropout_rate"],
                    modality_type=spec["type"]
                )

                # replace your decoder_init with this
                def decoder_init(rng_key, input_shape):
                    dummy_input = jnp.zeros(input_shape)
                    variables = decoder_def.init(rng_key, dummy_input, training=False)
                    params = variables["params"]
                    out = decoder_def.apply({"params": params}, dummy_input, training=False)
                    return out, params  # <-- two output

                def decoder_apply(params, input_in):
                    return decoder_def.apply({"params": params}, input_in, training=False)

                decoder_fn = numpyro.module(
                    f"decoder_{mod_name}",
                    (decoder_init, decoder_apply),
                    input_shape=(decoder_input_dim,),
                )

                # Apply decoder per item in batch
                decoder_out = jax.vmap(decoder_fn, in_axes=0)(decoder_input)
                numpyro.deterministic(f"decoder_{mod_name}_output", decoder_out)

                # Sample observations based on modality type
                self._sample_modality_observations(mod_name, spec, decoder_out, data.get(mod_name))

        # Add outcome predictors if specified
        if self.predictor_specs and outcomes is not None:
            # Prepare predictor input: concatenate z with covariates_predictor
            if covariates_predictor is not None:
                predictor_input = jnp.concatenate([z, covariates_predictor], axis=1)
            else:
                predictor_input = z

            # Register and apply predictors for each outcome
            for pred_name, pred_spec in self.predictor_specs.items():
                if pred_name not in outcomes:
                    continue  # Skip if outcome data not provided

                predictor_type = pred_spec.get("type", "neural_network")  # Default to NN

                if predictor_type == "linear_regression":
                    # BAYESIAN LINEAR REGRESSION PATH
                    input_dim = predictor_input.shape[-1]
                    output_dim = pred_spec["output_dim"]
                    
                    # Priors on regression coefficients
                    beta_prior_scale = pred_spec.get("beta_prior_scale", 1.0)
                    beta = sample(f"beta_{pred_name}",
                                dist.Normal(0, beta_prior_scale).expand([input_dim, output_dim]).to_event(2))
                    
                    # Prior on intercept
                    intercept_prior_scale = pred_spec.get("intercept_prior_scale", 1.0)
                    intercept = sample(f"intercept_{pred_name}",
                                     dist.Normal(0, intercept_prior_scale).expand([output_dim]).to_event(1))
                    
                    # Linear prediction: X @ beta + intercept
                    mu = predictor_input @ beta + intercept  # Shape: (N, output_dim)
                    
                    # Store predictions as deterministic for monitoring
                    numpyro.deterministic(f"predictor_{pred_name}_mu", mu)
                    
                    # Outcome-specific likelihood
                    outcome_type = pred_spec.get("outcome_type", "gaussian")
                    
                    if outcome_type == "gaussian":
                        # Prior on observation noise
                        sigma_prior_scale = pred_spec.get("sigma_prior_scale", 1.0)
                        sigma = sample(f"sigma_{pred_name}",
                                     dist.HalfNormal(sigma_prior_scale).expand([output_dim]).to_event(1))
                        sample(f"obs_outcome_{pred_name}",
                              dist.Normal(mu, sigma).to_event(1),
                              obs=outcomes[pred_name])
                    
                    elif outcome_type == "bernoulli":
                        # Logistic regression (mu are logits)
                        sample(f"obs_outcome_{pred_name}",
                              dist.Bernoulli(logits=mu).to_event(1),
                              obs=outcomes[pred_name])
                    
                    elif outcome_type == "categorical":
                        # Multinomial logistic regression
                        sample(f"obs_outcome_{pred_name}",
                              dist.Categorical(logits=mu),
                              obs=outcomes[pred_name])
                    
                    elif outcome_type == "poisson":
                        # Poisson regression (mu are log rates)
                        mu_clipped = jnp.clip(mu, -10.0, 10.0)
                        rate = jnp.exp(mu_clipped)
                        sample(f"obs_outcome_{pred_name}",
                              dist.Poisson(rate).to_event(1),
                              obs=outcomes[pred_name])
                    
                    else:
                        raise ValueError(f"Unsupported outcome_type for linear_regression: {outcome_type}")

                else:
                    # NEURAL NETWORK PATH (existing code)
                    predictor_def = self.predictor_mlps[pred_name]

                    def predictor_init(rng_key, input_shape):
                        dummy_input = jnp.zeros(input_shape)
                        variables = predictor_def.init(rng_key, dummy_input, training=False)
                        params = variables["params"]
                        out = predictor_def.apply({"params": params}, dummy_input, training=False)
                        return out, params

                    def predictor_apply(params, x_in):
                        return predictor_def.apply({"params": params}, x_in, training=False)

                    predictor_fn = numpyro.module(
                        f"predictor_{pred_name}",
                        (predictor_init, predictor_apply),
                        input_shape=(predictor_input.shape[-1],)
                    )

                    # Apply predictor
                    pred_out = jax.vmap(predictor_fn, in_axes=0)(predictor_input)
                    numpyro.deterministic(f"predictor_{pred_name}_output", pred_out)

                    # Use the same sampling method as modalities!
                    self._sample_modality_observations(f"outcome_{pred_name}", pred_spec, pred_out, outcomes[pred_name])

    def _sample_modality_observations(self, mod_name: str, spec: Dict[str, Any],
                                      decoder_out: jnp.ndarray, obs_data: jnp.ndarray = None):
        """Sample observations for a specific modality."""
        # Handle both "type" (for modalities) and "modality_type" (for predictors)
        mod_type = spec.get("type") or spec.get("modality_type")
        if mod_type is None:
            raise ValueError(f"Spec must contain either 'type' or 'modality_type' key. Got: {spec.keys()}")

        if mod_type == "gaussian":
            mu, log_sigma = jnp.split(decoder_out, 2, axis=-1)
            sigma = jnp.exp(log_sigma)
            sample(f"obs_{mod_name}", dist.Normal(mu, sigma).to_event(1), obs=obs_data)

        elif mod_type == "categorical":
            logits = decoder_out
            sample(f"obs_{mod_name}", dist.Categorical(logits=logits), obs=obs_data)

        elif mod_type == "categorical_multi":
            # Handle multiple categorical questions
            categories_per_question = spec.get("likelihood_params", {}).get("categories_per_question", [])
            if not categories_per_question:
                raise ValueError("categorical_multi requires 'categories_per_question' in likelihood_params")

            # Split decoder output into logits for each question
            logits_split = []
            start_idx = 0
            for num_cats in categories_per_question:
                end_idx = start_idx + num_cats
                question_logits = decoder_out[:, start_idx:end_idx]
                logits_split.append(question_logits)
                start_idx = end_idx

            # Sample each question separately using categorical distribution
            if obs_data is not None:
                # obs_data should be shape (batch_size, num_questions) with integer responses
                for q, question_logits in enumerate(logits_split):
                    question_obs = obs_data[:, q] if obs_data.ndim > 1 else None
                    sample(f"obs_{mod_name}_q{q}", dist.Categorical(logits=question_logits), obs=question_obs)
            else:
                # For generation, sample all questions
                for q, question_logits in enumerate(logits_split):
                    sample(f"obs_{mod_name}_q{q}", dist.Categorical(logits=question_logits))

        elif mod_type == "ordinal":
            eta = decoder_out
            # For ordinal, we need cutpoints as additional parameters
            num_classes = spec["dim"]
            cutpoints = sample(f"cutpoints_{mod_name}",
                               dist.Normal(0, 1).expand([num_classes - 1]).to_event(1))
            cutpoints = jnp.cumsum(jnp.abs(cutpoints))  # Ensure ordering
            numpyro.deterministic(f"cutpoints_{mod_name}_ordered", cutpoints)

            # Compute ordinal probabilities
            cum_probs = jax.nn.sigmoid(cutpoints[None, :] - eta[:, None])
            probs = jnp.concatenate([
                cum_probs[:, :1],
                cum_probs[:, 1:] - cum_probs[:, :-1],
                1 - cum_probs[:, -1:]
            ], axis=1)
            sample(f"obs_{mod_name}", dist.Categorical(probs=probs), obs=obs_data)

        elif mod_type == "ordinal_multi":
            # Handle multiple ordinal questions
            categories_per_question = spec.get("likelihood_params", {}).get("categories_per_question", [])
            if not categories_per_question:
                raise ValueError("ordinal_multi requires 'categories_per_question' in likelihood_params")

            # Split decoder output into eta for each question
            eta_split = []
            start_idx = 0
            for num_cats in categories_per_question:
                end_idx = start_idx + 1  # Each ordinal question needs 1 eta value
                question_eta = decoder_out[:, start_idx:end_idx]
                # Clip eta to prevent numerical instability
                question_eta = jnp.clip(question_eta.squeeze(-1), -10.0, 10.0)
                eta_split.append(question_eta)
                start_idx = end_idx

            # Sample each question separately using ordinal distribution
            if obs_data is not None:
                # obs_data should be shape (batch_size, num_questions) with integer responses
                for q, (question_eta, num_cats) in enumerate(zip(eta_split, categories_per_question)):
                    # Create cutpoints for this question
                    cutpoints_q = sample(f"cutpoints_{mod_name}_q{q}",
                                       dist.Normal(0, 0.5).expand([num_cats - 1]).to_event(1))
                    cutpoints_q = jnp.cumsum(jnp.abs(cutpoints_q))  # Ensure ordering
                    # Clip cutpoints to reasonable range
                    cutpoints_q = jnp.clip(cutpoints_q, -10.0, 10.0)
                    numpyro.deterministic(f"cutpoints_{mod_name}_q{q}_ordered", cutpoints_q)

                    # Compute ordinal probabilities for this question
                    # Use more stable computation
                    logits = cutpoints_q[None, :] - question_eta[:, None]
                    cum_probs = jax.nn.sigmoid(logits)

                    # Compute category probabilities with numerical stability
                    probs_list = []
                    # First category
                    probs_list.append(cum_probs[:, 0:1])
                    # Middle categories
                    for k in range(1, num_cats - 1):
                        probs_list.append(cum_probs[:, k:k+1] - cum_probs[:, k-1:k])
                    # Last category
                    probs_list.append(1.0 - cum_probs[:, -1:])

                    probs = jnp.concatenate(probs_list, axis=1)
                    # Ensure probabilities are valid
                    probs = jnp.clip(probs, 1e-10, 1.0)
                    probs = probs / jnp.sum(probs, axis=1, keepdims=True)

                    question_obs = obs_data[:, q] if obs_data.ndim > 1 else None
                    sample(f"obs_{mod_name}_q{q}", dist.Categorical(probs=probs), obs=question_obs)
            else:
                # For generation, sample all questions
                for q, (question_eta, num_cats) in enumerate(zip(eta_split, categories_per_question)):
                    # Create cutpoints for this question
                    cutpoints_q = sample(f"cutpoints_{mod_name}_q{q}",
                                       dist.Normal(0, 0.5).expand([num_cats - 1]).to_event(1))
                    cutpoints_q = jnp.cumsum(jnp.abs(cutpoints_q))  # Ensure ordering
                    # Clip cutpoints to reasonable range
                    cutpoints_q = jnp.clip(cutpoints_q, -10.0, 10.0)
                    numpyro.deterministic(f"cutpoints_{mod_name}_q{q}_ordered", cutpoints_q)

                    # Compute ordinal probabilities for this question
                    # Use more stable computation
                    logits = cutpoints_q[None, :] - question_eta[:, None]
                    cum_probs = jax.nn.sigmoid(logits)

                    # Compute category probabilities with numerical stability
                    probs_list = []
                    # First category
                    probs_list.append(cum_probs[:, 0:1])
                    # Middle categories
                    for k in range(1, num_cats - 1):
                        probs_list.append(cum_probs[:, k:k+1] - cum_probs[:, k-1:k])
                    # Last category
                    probs_list.append(1.0 - cum_probs[:, -1:])

                    probs = jnp.concatenate(probs_list, axis=1)
                    # Ensure probabilities are valid
                    probs = jnp.clip(probs, 1e-10, 1.0)
                    probs = probs / jnp.sum(probs, axis=1, keepdims=True)

                    sample(f"obs_{mod_name}_q{q}", dist.Categorical(probs=probs))

        elif mod_type == "poisson":
            log_rate = decoder_out
            # Clip log_rate more aggressively to prevent overflow when exponentiated
            log_rate = jnp.clip(log_rate, -10.0, 10.0)
            rate = jnp.exp(log_rate)
            sample(f"obs_{mod_name}", dist.Poisson(rate).to_event(1), obs=obs_data)

        elif mod_type == "negative_binomial":
            log_rate = decoder_out
            # Clip log_rate more aggressively to prevent overflow when exponentiated
            log_rate = jnp.clip(log_rate, -10.0, 10.0)
            rate = jnp.exp(log_rate)

            # Sample concentration parameter if not provided
            concentration_prior = spec.get("likelihood_params", {}).get("concentration_prior", 1.0)
            concentration = sample(f"concentration_{mod_name}",
                                 dist.Exponential(1.0 / concentration_prior))

            # Convert to probs parameterization for NegativeBinomial
            probs = concentration / (concentration + rate + 1e-8)
            sample(f"obs_{mod_name}",
                   dist.NegativeBinomial(total_count=concentration, probs=probs).to_event(1),
                   obs=obs_data)

        elif mod_type == "cross_entropy":
            logits = decoder_out
            # For cross-entropy with count data, we sample from categorical and convert to counts
            if obs_data is not None:
                # During training/inference with observed data
                # Treat this as a categorical likelihood where obs_data are normalized counts
                probs = jax.nn.softmax(logits, axis=-1)
                # Normalize observed counts to probabilities for likelihood computation
                epsilon = spec.get("likelihood_params", {}).get("epsilon", 1e-8)
                obs_normalized = obs_data / (jnp.sum(obs_data, axis=-1, keepdims=True) + epsilon)

                # Use Dirichlet-Multinomial or just Categorical for simplicity
                sample(f"obs_{mod_name}", dist.Categorical(probs=probs),
                       obs=jnp.argmax(obs_normalized, axis=-1))
            else:
                # During generation
                sample(f"obs_{mod_name}", dist.Categorical(logits=logits))

        elif mod_type == "bernoulli":
            logits = decoder_out
            sample(f"obs_{mod_name}", dist.Bernoulli(logits=logits).to_event(1), obs=obs_data)

        else:
            raise ValueError(f"Unsupported modality type: {mod_type}")

    def fit(self, data: Dict[str, jnp.ndarray], *,
            covariates_prior: Optional[jnp.ndarray] = None,
            covariates_decoder: Optional[jnp.ndarray] = None,
            covariates_predictor: Optional[jnp.ndarray] = None,
            outcomes: Optional[Dict[str, jnp.ndarray]] = None,
            batch_size: Optional[int] = None,
            shuffle: bool = True,
            num_steps: int = 10000, lr: float = 1e-3,
            seed: int = 0, progress_bar: bool = True, track_params: bool = False):
        """
        Fit the multi-modal VAE using stochastic variational inference.

        Args:
            data: Dictionary of modality data
            covariates_prior: Optional covariate matrix for structured prior
            covariates_decoder: Optional covariate matrix for decoder
            covariates_predictor: Optional covariate matrix for outcome prediction
            outcomes: Optional outcome dictionary
            batch_size: Mini-batch size (None for full-batch training)
            shuffle: Whether to shuffle data between epochs
            num_steps: Number of training iterations
            lr: Learning rate
            seed: Random seed
            progress_bar: Whether to show progress bar
            track_params: Whether to track parameter evolution (structured prior only)
        """

        if self.use_structured_prior and covariates_prior is None:
            raise ValueError("covariates_prior must be provided when use_structured_prior=True")

        if self.predictor_specs and outcomes is None:
            raise ValueError("Outcomes must be provided when predictor_specs is specified")

        rng_key = random.PRNGKey(seed)

        # Get dataset size
        n_samples = next(iter(data.values())).shape[0]

        # Handle batch size
        if batch_size is None:
            # Full-batch training (backward compatibility)
            batch_size = n_samples
            print(f"Using full-batch training with {n_samples} samples")
        else:
            print(f"Using mini-batch training with batch_size={batch_size}, dataset_size={n_samples}")
            if batch_size > n_samples:
                print(f"Warning: batch_size ({batch_size}) > dataset_size ({n_samples}), using full-batch")
                batch_size = n_samples

        # Create custom MultiModalGuide
        self.auto_guide = self._create_auto_guide()
        guide = self.auto_guide

        # Print guide information
        print(f"Using custom MultiModalGuide with {self.fusion_strategy} fusion strategy")
        
        # Determine whether to use WeightedELBO or standard Trace_ELBO
        use_weighted_elbo = (
            self.kl_weight != 1.0 or  # Non-default KL weight
            any(w != 1.0 for w in self.reconstruction_weights.values()) or  # Custom reconstruction weights
            any(w != 1.0 for w in self.predictor_weights.values())  # Custom predictor weights
        )
        
        if use_weighted_elbo:
            # Use WeightedELBO for flexible loss weighting
            elbo_loss = WeightedELBO(
                kl_weight=self.kl_weight,
                reconstruction_weights=self.reconstruction_weights,
                predictor_weights=self.predictor_weights
            )
            print(f"Using WeightedELBO (kl_weight={self.kl_weight:.3f}, "
                  f"custom_reconstruction_weights={any(w != 1.0 for w in self.reconstruction_weights.values())}, "
                  f"custom_predictor_weights={any(w != 1.0 for w in self.predictor_weights.values())})")
        else:
            # Use standard Trace_ELBO for efficiency (default behavior)
            elbo_loss = Trace_ELBO()
            print("Using standard Trace_ELBO (no custom loss weighting)")

        # Initialize SVI with the appropriate batch size
        # The guide needs to be initialized with the same batch size used during training
        if batch_size < n_samples:
            # For mini-batch training, initialize with a representative batch
            init_batches = create_batches(data, covariates_prior, covariates_decoder, covariates_predictor, outcomes, batch_size, shuffle=False, seed=seed)
            init_data, init_covariates_prior, init_covariates_decoder, init_covariates_predictor, init_outcomes = init_batches[0]
            print(f"Initializing guide with batch_size={batch_size}")
        else:
            # For full-batch training, use the full dataset
            init_data, init_covariates_prior, init_covariates_decoder, init_covariates_predictor, init_outcomes = data, covariates_prior, covariates_decoder, covariates_predictor, outcomes
            print(f"Initializing guide with full dataset (size={n_samples})")

        svi = SVI(self.model, guide, optim.Adam(lr), elbo_loss)
        svi_state = svi.init(rng_key, data=init_data, covariates_prior=init_covariates_prior, covariates_decoder=init_covariates_decoder, covariates_predictor=init_covariates_predictor, outcomes=init_outcomes)

        # Training tracking
        losses = []
        param_history = [] if track_params and self.use_structured_prior else None

        # JIT compiled update function for batch training
        @jax.jit
        def svi_update_batch(state, batch_data, batch_cov_prior, batch_cov_decoder, batch_cov_predictor, batch_outcomes):
            return svi.update(state, data=batch_data, covariates_prior=batch_cov_prior, covariates_decoder=batch_cov_decoder, covariates_predictor=batch_cov_predictor, outcomes=batch_outcomes)

        if progress_bar:
            try:
                import tqdm
                pbar = tqdm.tqdm(total=num_steps, desc="Training")
            except ImportError:
                print("tqdm not available, training without progress bar...")
                pbar = None
        else:
            pbar = None

        step = 0
        epoch = 0

        while step < num_steps:
            # Create batches for this epoch
            epoch_seed = seed + epoch if shuffle else seed
            batches = create_batches(data, covariates_prior, covariates_decoder, covariates_predictor, outcomes, batch_size, shuffle, epoch_seed)

            for batch_data, batch_cov_prior, batch_cov_decoder, batch_cov_predictor, batch_outcomes in batches:
                if step >= num_steps:
                    break

                # Skip batches that don't match the initialized batch size
                # This handles the last batch which might be smaller
                current_batch_size = next(iter(batch_data.values())).shape[0]
                if batch_size < n_samples and current_batch_size != batch_size:
                    if pbar is not None:
                        pbar.set_postfix({'status': f'skipping batch (size {current_batch_size})'})
                    continue

                # Update with current batch
                svi_state, batch_loss = svi_update_batch(svi_state, batch_data, batch_cov_prior, batch_cov_decoder, batch_cov_predictor, batch_outcomes)
                losses.append(float(batch_loss))

                # Track parameters if requested
                if track_params and self.use_structured_prior and step % 100 == 0:
                    current_params = svi.get_params(svi_state)
                    param_snapshot = {
                        'step': step,
                        'epoch': epoch,
                        'loss': float(batch_loss),
                        'gamma_loc': current_params.get('gamma_loc', None),
                        'gamma_scale': current_params.get('gamma_scale', None),
                        'prior_scale_q': current_params.get('prior_scale_q', None),
                        'prior_correlation_L': current_params.get('prior_correlation_L', None)
                    }
                    param_history.append(param_snapshot)

                # Update progress
                if pbar is not None:
                    postfix = {'loss': f'{batch_loss:.6f}', 'epoch': epoch}
                    pbar.set_postfix(postfix)
                    pbar.update(1)
                elif step % 100 == 0:
                    print(f"Step {step}/{num_steps}, Epoch {epoch}, Loss: {batch_loss:.6f}")

                step += 1

            epoch += 1

        if pbar is not None:
            pbar.close()

        # Store results
        self.losses = losses
        self.param_history = param_history
        self.params = svi.get_params(svi_state)
        self.svi = svi

        print(f"Training completed. Final loss: {losses[-1]:.6f}")

    def sample_posterior(self, data: Dict[str, jnp.ndarray], *,
                         covariates_prior: Optional[jnp.ndarray] = None,
                         covariates_decoder: Optional[jnp.ndarray] = None,
                         covariates_predictor: Optional[jnp.ndarray] = None,
                         outcomes: Optional[Dict[str, jnp.ndarray]] = None,
                         batch_size: Optional[int] = None,
                         num_samples: int = 1000,
                         num_samples_per_chunk: Optional[int] = None,
                         return_sites: Optional[List[str]] = None,
                         seed: int = 1) -> Dict[str, jnp.ndarray]:
        """
        Sample from the posterior distribution, including latent z and decoder outputs.

        Args:
            data: Dictionary of modality data
            covariates_prior: Optional covariate matrix for structured prior
            covariates_decoder: Optional covariate matrix for decoder
            covariates_predictor: Optional covariate matrix for outcome prediction
            outcomes: Optional outcome dictionary
            batch_size: Mini-batch size for processing data samples (None for full-batch, recommended)
            num_samples: Total number of posterior samples to generate
            num_samples_per_chunk: If specified, chunk posterior sampling to avoid OOM 
                                   (None = no chunking, recommended for speed)
            return_sites: Optional list of sites to return. If None, uses intelligent defaults:
                         - Always includes "z" (latent variables)
                         - Includes "z_simplex" if using logistic_normal prior
                         - Includes predictor parameters for linear_regression predictors
                         - Includes structured prior parameters if use_structured_prior=True
                         - Does NOT include decoder outputs by default (to save memory)
                         Example: ["z", "z_simplex", "beta_outcome", "gamma"]
            seed: Random seed

        Returns:
            Dictionary of sampled values
            
        Note:
            By default (batch_size=None, num_samples_per_chunk=None), this uses the fastest
            approach: single call to Predictive with all data and all posterior samples.
            Only use chunking if you encounter GPU OOM errors.
        """
        if not hasattr(self, 'params'):
            raise ValueError("Model must be fitted before sampling. Call fit() first.")

        if self.use_structured_prior and covariates_prior is None:
            raise ValueError("covariates_prior must be provided when use_structured_prior=True")

        rng_key = random.PRNGKey(seed)
        n_data_samples = next(iter(data.values())).shape[0]

        # Check if we need to chunk over posterior samples (only if explicitly requested)
        if num_samples_per_chunk is not None and num_samples > num_samples_per_chunk:
            print(f"Chunking posterior sampling: {num_samples} samples in chunks of {num_samples_per_chunk}")
            return self._sample_posterior_chunked(
                data, covariates_prior, covariates_decoder, covariates_predictor, outcomes,
                batch_size, num_samples, num_samples_per_chunk, return_sites, rng_key
            )

        # Handle batch size (original behavior if num_samples <= num_samples_per_chunk)
        if batch_size is None or batch_size >= n_data_samples:
            # Full-batch sampling (original behavior)
            return self._sample_posterior_batch(data, covariates_prior, covariates_decoder, covariates_predictor, outcomes, num_samples, return_sites, rng_key)

        # Mini-batch sampling
        print(f"Using batched posterior sampling with batch_size={batch_size}")

        # Create batches (no shuffling for inference)
        batches = create_batches(data, covariates_prior, covariates_decoder, covariates_predictor, outcomes, batch_size, shuffle=False, seed=seed)

        # Collect results from all batches
        all_results = []

        for i, (batch_data, batch_cov_prior, batch_cov_decoder, batch_cov_predictor, batch_outcomes) in enumerate(batches):
            rng_key, batch_key = random.split(rng_key)
            batch_results = self._sample_posterior_batch(
                batch_data, batch_cov_prior, batch_cov_decoder, batch_cov_predictor, batch_outcomes, num_samples, return_sites, batch_key
            )
            all_results.append(batch_results)

        # Concatenate results across batches
        final_results = {}
        for key in all_results[0].keys():
            # Concatenate along the data dimension (axis=1 for most arrays with shape (num_samples, n_data, ...))
            if all_results[0][key].ndim >= 2:
                final_results[key] = jnp.concatenate([r[key] for r in all_results], axis=1)
            else:
                # For scalars or 1D arrays, take from first batch (should be the same across batches)
                final_results[key] = all_results[0][key]

        return final_results

    def _sample_posterior_chunked(self, data: Dict[str, jnp.ndarray],
                                  covariates_prior: Optional[jnp.ndarray],
                                  covariates_decoder: Optional[jnp.ndarray],
                                  covariates_predictor: Optional[jnp.ndarray],
                                  outcomes: Optional[Dict[str, jnp.ndarray]],
                                  batch_size: Optional[int],
                                  num_samples: int,
                                  num_samples_per_chunk: int,
                                  return_sites: Optional[List[str]],
                                  rng_key: jax.random.PRNGKey) -> Dict[str, jnp.ndarray]:
        """
        Sample posterior in chunks over the posterior samples dimension to avoid OOM.
        
        This method chunks the posterior sampling over num_samples (axis=0) rather than
        over the data samples (axis=1). This is necessary because the Predictive class
        generates all num_samples at once, which can cause GPU OOM for large num_samples.
        
        Args:
            data: Dictionary of observed data arrays
            covariates_prior: Optional covariate matrix for structured prior
            covariates_decoder: Optional covariate matrix for decoder
            covariates_predictor: Optional covariate matrix for outcome prediction
            outcomes: Optional outcome dictionary
            batch_size: Optional batch size for data dimension (passed through to underlying method)
            num_samples: Total number of posterior samples to draw
            num_samples_per_chunk: Number of posterior samples to generate per chunk
            rng_key: PRNG key
            
        Returns:
            Dictionary of posterior samples with shape (num_samples, n_data, ...)
        """
        all_chunks = []
        remaining_samples = num_samples
        
        while remaining_samples > 0:
            current_chunk_size = min(num_samples_per_chunk, remaining_samples)
            
            # Split RNG key for this chunk
            rng_key, chunk_key = random.split(rng_key)
            
            # Sample this chunk using the standard sample_posterior path (but with reduced num_samples)
            # We need to call the internal logic directly to avoid infinite recursion
            n_data_samples = next(iter(data.values())).shape[0]
            
            if batch_size is None or batch_size >= n_data_samples:
                # Full-batch sampling for this chunk
                chunk_results = self._sample_posterior_batch(
                    data, covariates_prior, covariates_decoder, covariates_predictor, outcomes,
                    current_chunk_size, return_sites, chunk_key
                )
            else:
                # Mini-batch sampling for this chunk
                batches = create_batches(data, covariates_prior, covariates_decoder, covariates_predictor, 
                                       outcomes, batch_size, shuffle=False, seed=int(chunk_key[0]))
                
                batch_results = []
                for batch_data, batch_cov_prior, batch_cov_decoder, batch_cov_predictor, batch_outcomes in batches:
                    chunk_key, batch_key = random.split(chunk_key)
                    batch_result = self._sample_posterior_batch(
                        batch_data, batch_cov_prior, batch_cov_decoder, batch_cov_predictor, batch_outcomes,
                        current_chunk_size, return_sites, batch_key
                    )
                    batch_results.append(batch_result)
                
                # Concatenate across batches (data dimension, axis=1)
                chunk_results = {}
                for key in batch_results[0].keys():
                    if batch_results[0][key].ndim >= 2:
                        chunk_results[key] = jnp.concatenate([r[key] for r in batch_results], axis=1)
                    else:
                        chunk_results[key] = batch_results[0][key]
            
            all_chunks.append(chunk_results)
            remaining_samples -= current_chunk_size
        
        # Concatenate chunks along the posterior samples dimension (axis=0)
        final_results = {}
        for key in all_chunks[0].keys():
            # Concatenate along axis=0 (posterior samples dimension)
            final_results[key] = jnp.concatenate([chunk[key] for chunk in all_chunks], axis=0)
        
        return final_results

    def _sample_posterior_batch(self, data: Dict[str, jnp.ndarray],
                                covariates_prior: Optional[jnp.ndarray],
                                covariates_decoder: Optional[jnp.ndarray],
                                covariates_predictor: Optional[jnp.ndarray],
                                outcomes: Optional[Dict[str, jnp.ndarray]],
                                num_samples: int,
                                return_sites: Optional[List[str]],
                                rng_key: jax.random.PRNGKey) -> Dict[str, jnp.ndarray]:
        """Helper method to sample posterior for a single batch."""
        
        # Build intelligent default return_sites if not provided
        if return_sites is None:
            return_sites = ["z"]  # Always include latent variables
            
            # Add logistic normal sites if used
            if self.prior_type == "logistic_normal":
                return_sites.append("z_simplex")
            
            # Add predictor outputs/parameters if specified
            if self.predictor_specs:
                for pred_name, pred_spec in self.predictor_specs.items():
                    predictor_type = pred_spec.get("type", "neural_network")
                    if predictor_type == "linear_regression":
                        # Add linear regression parameters (what user wants for uncertainty)
                        return_sites.extend([
                            f"beta_{pred_name}",
                            f"intercept_{pred_name}",
                            f"predictor_{pred_name}_mu"
                        ])
                        # Add sigma for Gaussian outcomes
                        if pred_spec.get("modality_type", "gaussian") == "gaussian":
                            return_sites.append(f"sigma_{pred_name}")
                    else:
                        # Neural network predictors - include outputs
                        return_sites.append(f"predictor_{pred_name}_output")

            # Add structured prior sites if used
            if self.use_structured_prior:
                return_sites.extend(["gamma", "prior_scale", "prior_correlation", "prior_mean", "covariance_matrix"])

            # Add MoE-specific sites when using custom guide with multi-modal fusion
            if isinstance(self.auto_guide, MultiModalGuide):
                return_sites.extend(["moe_weights"])
                if len(self.encode_modalities) > 1:
                    return_sites.extend(["moe_expert_means", "moe_expert_log_stds", "moe_expert_stds"])

                # Add modality-specific latent samples (only for encoded modalities)
                for mod_name in self.encode_modalities:
                    return_sites.append(f"z_{mod_name}")

                # Add learned weights for moe_learned strategy
                if self.fusion_strategy == "moe_learned":
                    return_sites.append("moe_learned_weights")
            
            # NOTE: We do NOT include decoder outputs by default to save memory
            # Users can explicitly request them by passing return_sites=["z", "decoder_word_counts_output", ...]

        predictive = Predictive(self.model, guide=self.auto_guide, params=self.params,
                                num_samples=num_samples, return_sites=return_sites)

        return predictive(rng_key, data=data, covariates_prior=covariates_prior, covariates_decoder=covariates_decoder, covariates_predictor=covariates_predictor, outcomes=outcomes)

    def generate_samples(self, num_samples: int = 100,
                         covariates_prior: Optional[jnp.ndarray] = None,
                         covariates_decoder: Optional[jnp.ndarray] = None,
                         covariates_predictor: Optional[jnp.ndarray] = None,
                         outcomes: Optional[Dict[str, jnp.ndarray]] = None,
                         batch_size: Optional[int] = None,
                         seed: int = 2) -> Dict[str, jnp.ndarray]:
        """
        Generate new samples from the trained model.

        Args:
            num_samples: Number of samples to generate
            covariates_prior: Optional covariate matrix for structured prior
            covariates_decoder: Optional covariate matrix for decoder
            covariates_predictor: Optional covariate matrix for outcome prediction
            outcomes: Optional outcome dictionary (if using predictors)
            batch_size: Mini-batch size for generation (None for full-batch)
            seed: Random seed

        Returns:
            Dictionary of generated samples
        """
        if not hasattr(self, 'params'):
            raise ValueError("Model must be fitted before generating. Call fit() first.")

        if self.use_structured_prior and covariates_prior is None:
            raise ValueError("covariates_prior must be provided when use_structured_prior=True")

        # Handle batch size for memory efficiency
        if batch_size is None or batch_size >= num_samples:
            # Generate all samples at once
            return self._generate_samples_batch(num_samples, covariates_prior, covariates_decoder, covariates_predictor, outcomes, seed)

        # Generate in batches
        print(f"Generating {num_samples} samples in batches of {batch_size}")

        rng_key = random.PRNGKey(seed)
        all_generated = []

        remaining_samples = num_samples
        while remaining_samples > 0:
            current_batch_size = min(batch_size, remaining_samples)

            # Handle covariates_prior for current batch
            batch_covariates_prior = None
            if covariates_prior is not None:
                if covariates_prior.shape[0] == num_samples:
                    # Covariates provided for each sample
                    start_idx = num_samples - remaining_samples
                    end_idx = start_idx + current_batch_size
                    batch_covariates_prior = covariates_prior[start_idx:end_idx]
                else:
                    # Single covariate pattern, broadcast
                    batch_covariates_prior = jnp.tile(covariates_prior, (current_batch_size, 1))
                    
            # Handle covariates_decoder for current batch
            batch_covariates_decoder = None
            if covariates_decoder is not None:
                if covariates_decoder.shape[0] == num_samples:
                    start_idx = num_samples - remaining_samples
                    end_idx = start_idx + current_batch_size
                    batch_covariates_decoder = covariates_decoder[start_idx:end_idx]
                else:
                    batch_covariates_decoder = jnp.tile(covariates_decoder, (current_batch_size, 1))
                    
            # Handle covariates_predictor for current batch
            batch_covariates_predictor = None
            if covariates_predictor is not None:
                if covariates_predictor.shape[0] == num_samples:
                    start_idx = num_samples - remaining_samples
                    end_idx = start_idx + current_batch_size
                    batch_covariates_predictor = covariates_predictor[start_idx:end_idx]
                else:
                    batch_covariates_predictor = jnp.tile(covariates_predictor, (current_batch_size, 1))

            # Handle outcomes for current batch (similar logic)
            batch_outcomes = None
            if outcomes is not None:
                batch_outcomes = {}
                for outcome_name, outcome_data in outcomes.items():
                    if outcome_data.shape[0] == num_samples:
                        start_idx = num_samples - remaining_samples
                        end_idx = start_idx + current_batch_size
                        batch_outcomes[outcome_name] = outcome_data[start_idx:end_idx]
                    else:
                        batch_outcomes[outcome_name] = jnp.tile(outcome_data, (current_batch_size, 1))

            # Generate batch
            rng_key, batch_key = random.split(rng_key)
            batch_generated = self._generate_samples_batch(
                current_batch_size, batch_covariates_prior, batch_covariates_decoder, batch_covariates_predictor, batch_outcomes, int(batch_key[0])
            )
            all_generated.append(batch_generated)

            remaining_samples -= current_batch_size

        # Concatenate all batches
        final_generated = {}
        for key in all_generated[0].keys():
            final_generated[key] = jnp.concatenate([batch[key] for batch in all_generated], axis=0)

        return final_generated

    def _generate_samples_batch(self, num_samples: int,
                                covariates_prior: Optional[jnp.ndarray],
                                covariates_decoder: Optional[jnp.ndarray],
                                covariates_predictor: Optional[jnp.ndarray],
                                outcomes: Optional[Dict[str, jnp.ndarray]],
                                seed: int) -> Dict[str, jnp.ndarray]:
        """Helper method to generate samples for a single batch."""
        rng_key = random.PRNGKey(seed)

        # Create dummy data for the model signature
        dummy_data = {}
        for mod_name, spec in self.modality_specs.items():
            dummy_data[mod_name] = jnp.zeros((num_samples, spec["dim"]))

        # Create dummy outcomes if predictors are specified
        dummy_outcomes = None
        if self.predictor_specs:
            dummy_outcomes = {}
            for pred_name, pred_spec in self.predictor_specs.items():
                dummy_outcomes[pred_name] = jnp.zeros((num_samples, pred_spec["output_dim"]))

        predictive = Predictive(self.model, params=self.params, num_samples=1)
        samples = predictive(rng_key, data=dummy_data, covariates_prior=covariates_prior, covariates_decoder=covariates_decoder, covariates_predictor=covariates_predictor, outcomes=dummy_outcomes)

        # Extract generated observations
        generated = {}
        for mod_name in self.modality_specs.keys():
            obs_key = f"obs_{mod_name}"
            if obs_key in samples:
                generated[mod_name] = samples[obs_key][0]  # Remove sample dimension

        # Extract generated outcomes
        if self.predictor_specs:
            for pred_name in self.predictor_specs.keys():
                outcome_key = f"outcome_{pred_name}"
                if outcome_key in samples:
                    generated[f"outcome_{pred_name}"] = samples[outcome_key][0]

        return generated

    def get_variational_params(self) -> Dict[str, jnp.ndarray]:
        """Get the learned variational parameters (for structured prior models)."""
        if not hasattr(self, 'params'):
            raise ValueError("Model must be fitted before accessing parameters. Call fit() first.")

        if not self.use_structured_prior:
            raise ValueError("This method is only available for structured prior models.")

        # Extract the variational parameters
        variational_params = {}
        if 'gamma_loc' in self.params:
            variational_params['gamma_loc'] = self.params['gamma_loc']
        if 'gamma_scale' in self.params:
            variational_params['gamma_scale'] = self.params['gamma_scale']
        if 'prior_scale_q' in self.params:
            variational_params['prior_scale_q'] = self.params['prior_scale_q']
        if 'prior_correlation_L' in self.params:
            variational_params['prior_correlation_L'] = self.params['prior_correlation_L']

        return variational_params

    def predict_outcomes(self, data: Dict[str, jnp.ndarray],
                        covariates: Optional[jnp.ndarray] = None,
                        predictor_mlps: Dict[str, PredictorMLP] = None,
                        predictor_params: Dict[str, Dict] = None,
                        batch_size: Optional[int] = None,
                        num_samples: int = 100,
                        seed: int = 42) -> Dict[str, jnp.ndarray]:
        """
        Predict K outcomes using ideal points and additional covariates.

        Args:
            data: Dictionary of modality data used to infer ideal points
            covariates: Additional covariates for prediction (N, covariate_dim)
            predictor_mlps: Dictionary mapping outcome names to PredictorMLP instances
            predictor_params: Dictionary mapping outcome names to their trained parameters
            batch_size: Mini-batch size for processing (None for full-batch)
            num_samples: Number of posterior samples to use for prediction
            seed: Random seed for sampling

        Returns:
            Dictionary mapping outcome names to predicted values (N, K) arrays

        Example:
            >>> # Create predictor MLPs for each outcome
            >>> predictor_mlps = {
            ...     "voting_outcome_1": PredictorMLP(output_dim=1, input_dim=latent_dim + covariate_dim),
            ...     "voting_outcome_2": PredictorMLP(output_dim=1, input_dim=latent_dim + covariate_dim)
            ... }
            >>> predictions = model.predict_outcomes(
            ...     data=test_data,
            ...     covariates=test_covariates,
            ...     predictor_mlps=predictor_mlps,
            ...     predictor_params=trained_predictor_params,
            ...     batch_size=256  # For memory efficiency
            ... )
        """
        if not hasattr(self, 'params'):
            raise ValueError("Model must be fitted before prediction. Call fit() first.")

        if predictor_mlps is None or predictor_params is None:
            raise ValueError("Both predictor_mlps and predictor_params must be provided for prediction.")

        n_data_samples = next(iter(data.values())).shape[0]

        # Handle batch size
        if batch_size is None or batch_size >= n_data_samples:
            # Full-batch prediction
            return self._predict_outcomes_batch(data, covariates, predictor_mlps,
                                               predictor_params, num_samples, seed)

        # Mini-batch prediction
        print(f"Using batched prediction with batch_size={batch_size}")

        # Create batches (no shuffling for inference)
        batches = create_batches(data, covariates, outcomes=None, batch_size=batch_size,
                                shuffle=False, seed=seed)

        # Process each batch
        all_predictions = []
        for i, (batch_data, batch_covariates, _) in enumerate(batches):
            batch_predictions = self._predict_outcomes_batch(
                batch_data, batch_covariates, predictor_mlps,
                predictor_params, num_samples, seed + i
            )
            all_predictions.append(batch_predictions)

        # Concatenate results across batches
        final_predictions = {}
        for outcome_name in predictor_mlps.keys():
            final_predictions[outcome_name] = jnp.concatenate(
                [batch[outcome_name] for batch in all_predictions], axis=0
            )

        return final_predictions

    def _predict_outcomes_batch(self, data: Dict[str, jnp.ndarray],
                               covariates: Optional[jnp.ndarray],
                               predictor_mlps: Dict[str, PredictorMLP],
                               predictor_params: Dict[str, Dict],
                               num_samples: int,
                               seed: int) -> Dict[str, jnp.ndarray]:
        """Helper method to predict outcomes for a single batch."""
        # Sample posterior ideal points
        posterior_samples = self.sample_posterior(
            data=data,
            covariates=covariates,
            num_samples=num_samples,
            seed=seed
        )

        # Extract ideal points (latent variables z)
        z_samples = posterior_samples["z"]  # Shape: (num_samples, N, latent_dim)

        # Average over posterior samples to get point estimates of ideal points
        z_mean = jnp.mean(z_samples, axis=0)  # Shape: (N, latent_dim)

        # Prepare input for predictor: concatenate ideal points with covariates
        if covariates is not None:
            predictor_input = jnp.concatenate([z_mean, covariates], axis=1)  # (N, latent_dim + covariate_dim)
        else:
            predictor_input = z_mean  # (N, latent_dim)

        # Make predictions for each outcome
        predictions = {}
        for outcome_name, predictor_mlp in predictor_mlps.items():
            if outcome_name not in predictor_params:
                raise ValueError(f"Parameters for predictor '{outcome_name}' not found in predictor_params.")

            params = predictor_params[outcome_name]
            # Apply predictor
            pred = predictor_mlp.apply({"params": params}, predictor_input, training=False)
            predictions[outcome_name] = pred

        return predictions

    @staticmethod
    def summary(samples: Dict[str, jnp.ndarray], *, hdi_prob: float = 0.95) -> Dict[str, Dict[str, jnp.ndarray]]:
        """Return per-parameter mean, SD, and equal-tailed HDI."""
        alpha = 1.0 - hdi_prob
        low_q, high_q = alpha / 2, 1.0 - alpha / 2

        stats = {}
        for name, arr in samples.items():
            if name.startswith("obs_"):
                continue  # Skip observation samples
            stats[name] = {
                "mean": jnp.mean(arr, axis=0),
                "sd": jnp.std(arr, axis=0),
                "low": jnp.quantile(arr, low_q, axis=0),
                "high": jnp.quantile(arr, high_q, axis=0),
            }
        return stats

    def recommend_batch_size(self, data: Dict[str, jnp.ndarray],
                            target_memory_gb: float = 8.0,
                            safety_factor: float = 0.7) -> int:
        """
        Recommend a batch size based on data size and target memory usage.

        Args:
            data: Dictionary of modality data
            target_memory_gb: Target memory usage in GB
            safety_factor: Safety factor to account for additional memory overhead

        Returns:
            Recommended batch size
        """
        return get_batch_size_recommendation(data, target_memory_gb * safety_factor)

    def get_moe_weights_summary(self, posterior_samples: Dict[str, jnp.ndarray]) -> Dict[str, jnp.ndarray]:
        """
        Extract and summarize MoE weights from posterior samples.

        Args:
            posterior_samples: Dictionary of posterior samples from sample_posterior()

        Returns:
            Dictionary with MoE weight statistics
        """
        if not isinstance(self.auto_guide, MultiModalGuide):
            raise ValueError("MoE weights are only available when using custom guide with multi-modal fusion strategies")

        if "moe_weights" not in posterior_samples:
            raise ValueError("MoE weights not found in posterior samples")

        moe_weights = posterior_samples["moe_weights"]  # Shape: (num_samples, N, num_modalities)
        modality_names = list(self.modality_specs.keys())

        # Basic statistics
        summary = {
            "mean_weights": jnp.mean(moe_weights, axis=0),  # (N, num_modalities)
            "std_weights": jnp.std(moe_weights, axis=0),   # (N, num_modalities)
            "median_weights": jnp.median(moe_weights, axis=0),  # (N, num_modalities)
            "weight_samples": moe_weights,  # Full samples for detailed analysis
            "modality_names": modality_names,
        }

        # Population-level statistics
        summary["global_mean_weights"] = jnp.mean(summary["mean_weights"], axis=0)  # (num_modalities,)
        summary["global_std_weights"] = jnp.std(summary["mean_weights"], axis=0)   # (num_modalities,)

        # Add expert-specific information if available
        if "moe_expert_means" in posterior_samples:
            summary["expert_means"] = posterior_samples["moe_expert_means"]
            summary["expert_stds"] = posterior_samples["moe_expert_stds"]

        return summary

    def analyze_moe_weights(self, data: Dict[str, jnp.ndarray],
                           covariates: Optional[jnp.ndarray] = None,
                           num_samples: int = 100,
                           seed: int = 42) -> Dict[str, jnp.ndarray]:
        """
        Convenience method to sample posterior and analyze MoE weights.

        Args:
            data: Dictionary of modality data
            covariates: Optional covariate matrix
            num_samples: Number of posterior samples
            seed: Random seed

        Returns:
            Dictionary with MoE weight analysis
        """
        if not isinstance(self.auto_guide, MultiModalGuide):
            raise ValueError("MoE analysis is only available when using custom guide with multi-modal fusion strategies")

        # Sample posterior
        posterior_samples = self.sample_posterior(
            data=data,
            covariates=covariates,
            num_samples=num_samples,
            seed=seed
        )

        # Analyze MoE weights
        return self.get_moe_weights_summary(posterior_samples)

    def get_linear_coefficients(self, data: Dict[str, jnp.ndarray],
                               outcomes: Dict[str, jnp.ndarray],
                               covariates_prior: Optional[jnp.ndarray] = None,
                               covariates_decoder: Optional[jnp.ndarray] = None,
                               covariates_predictor: Optional[jnp.ndarray] = None,
                               predictor_name: str = None,
                               num_samples: int = 1000,
                               ci_level: float = 0.95,
                               num_samples_per_chunk: int = 100,
                               seed: int = 42) -> Dict[str, Any]:
        """
        Extract coefficient estimates with credible intervals for linear regression predictors.
        
        This method is specifically for causal inference with Bayesian linear regression predictors.
        It provides full posterior distributions over coefficients, properly accounting for
        uncertainty in the latent variables.
        
        Args:
            data: Dictionary of modality data (to infer z)
            outcomes: Dictionary of outcome data
            covariates_prior: Optional covariate matrix for structured prior
            covariates_decoder: Optional covariate matrix for decoder
            covariates_predictor: Optional covariates for prediction
            predictor_name: Name of the predictor (if None, returns all linear predictors)
            num_samples: Number of posterior samples (default: 1000)
            ci_level: Credible interval level (default: 0.95)
            num_samples_per_chunk: Chunk size for posterior sampling to avoid OOM (default: 100)
            seed: Random seed
            
        Returns:
            Dictionary with coefficient statistics for each linear predictor:
            {
                "predictor_name": {
                    "beta_mean": (input_dim, output_dim),
                    "beta_std": (input_dim, output_dim),
                    "beta_lower": (input_dim, output_dim),
                    "beta_upper": (input_dim, output_dim),
                    "intercept_mean": (output_dim,),
                    "intercept_std": (output_dim,),
                    "intercept_lower": (output_dim,),
                    "intercept_upper": (output_dim,),
                    "sigma_mean": (output_dim,) [if Gaussian],
                    "beta_samples": (num_samples, input_dim, output_dim),
                    "feature_names": List[str]
                }
            }
            
        Example:
            >>> # For causal inference
            >>> predictor_specs = {
            ...     "policy_support": {
            ...         "type": "linear_regression",
            ...         "outcome_type": "gaussian",
            ...         "output_dim": 1,
            ...         "feature_names": ["ideology", "party", "region"]
            ...     }
            ... }
            >>> model = DeepLatent(modality_specs, predictor_specs=predictor_specs)
            >>> model.fit(data, outcomes=outcomes, covariates_predictor=covariates)
            >>> coef_results = model.get_linear_coefficients(
            ...     data, outcomes, covariates_predictor, num_samples=2000
            ... )
            >>> model.print_linear_coefficients(coef_results)
        """
        if not hasattr(self, 'params'):
            raise ValueError("Model must be fitted before extracting coefficients. Call fit() first.")
        
        # Sample posterior
        posterior_samples = self.sample_posterior(
            data=data,
            covariates_prior=covariates_prior,
            covariates_decoder=covariates_decoder,
            covariates_predictor=covariates_predictor,
            outcomes=outcomes,
            num_samples=num_samples,
            num_samples_per_chunk=num_samples_per_chunk,
            seed=seed
        )
        
        # Determine which predictors to extract
        if predictor_name is not None:
            if predictor_name not in self.predictor_specs:
                raise ValueError(f"Predictor '{predictor_name}' not found")
            if self.predictor_specs[predictor_name].get("type") != "linear_regression":
                raise ValueError(f"Predictor '{predictor_name}' is not a linear_regression type")
            predictor_names = [predictor_name]
        else:
            # Extract all linear regression predictors
            predictor_names = [name for name, spec in self.predictor_specs.items() 
                             if spec.get("type") == "linear_regression"]
        
        if not predictor_names:
            raise ValueError("No linear regression predictors found")
        
        # Extract coefficient statistics
        alpha = 1 - ci_level
        results = {}
        
        for pred_name in predictor_names:
            pred_spec = self.predictor_specs[pred_name]
            
            # Extract beta samples
            beta_key = f"beta_{pred_name}"
            if beta_key not in posterior_samples:
                raise ValueError(f"Beta coefficients not found for '{pred_name}'. "
                               f"Make sure the predictor type is 'linear_regression'.")
            
            beta_samples = posterior_samples[beta_key]  # (num_samples, input_dim, output_dim)
            intercept_samples = posterior_samples[f"intercept_{pred_name}"]  # (num_samples, output_dim)
            
            # Compute statistics
            pred_results = {
                "beta_mean": jnp.mean(beta_samples, axis=0),
                "beta_std": jnp.std(beta_samples, axis=0),
                "beta_lower": jnp.quantile(beta_samples, alpha/2, axis=0),
                "beta_upper": jnp.quantile(beta_samples, 1-alpha/2, axis=0),
                "intercept_mean": jnp.mean(intercept_samples, axis=0),
                "intercept_std": jnp.std(intercept_samples, axis=0),
                "intercept_lower": jnp.quantile(intercept_samples, alpha/2, axis=0),
                "intercept_upper": jnp.quantile(intercept_samples, 1-alpha/2, axis=0),
                "beta_samples": beta_samples,
                "intercept_samples": intercept_samples,
            }
            
            # Add sigma if Gaussian outcome
            if pred_spec.get("outcome_type", "gaussian") == "gaussian":
                sigma_samples = posterior_samples[f"sigma_{pred_name}"]
                pred_results["sigma_mean"] = jnp.mean(sigma_samples, axis=0)
                pred_results["sigma_std"] = jnp.std(sigma_samples, axis=0)
                pred_results["sigma_lower"] = jnp.quantile(sigma_samples, alpha/2, axis=0)
                pred_results["sigma_upper"] = jnp.quantile(sigma_samples, 1-alpha/2, axis=0)
            
            # Add feature names if provided
            if "feature_names" in pred_spec:
                pred_results["feature_names"] = pred_spec["feature_names"]
            else:
                # Generate default names
                input_dim = beta_samples.shape[1]
                z_dim = self.latent_dim
                feature_names = [f"z_{i}" for i in range(z_dim)]
                if covariates_predictor is not None:
                    cov_dim = input_dim - z_dim
                    feature_names.extend([f"covariate_{i}" for i in range(cov_dim)])
                pred_results["feature_names"] = feature_names
            
            results[pred_name] = pred_results
        
        return results
    
    def print_linear_coefficients(self, coefficient_results: Dict[str, Any],
                                  predictor_name: str = None,
                                  top_k: Optional[int] = None):
        """
        Pretty-print coefficient estimates with credible intervals.
        
        Args:
            coefficient_results: Results from get_linear_coefficients()
            predictor_name: Name of predictor to print (if None, prints all)
            top_k: Only print top k coefficients by absolute magnitude
            
        Example:
            >>> coef_results = model.get_linear_coefficients(data, outcomes, covariates)
            >>> model.print_linear_coefficients(coef_results, top_k=10)
        """
        if predictor_name is not None:
            if predictor_name not in coefficient_results:
                raise ValueError(f"Predictor '{predictor_name}' not found in results")
            predictors = [predictor_name]
        else:
            predictors = list(coefficient_results.keys())
        
        for pred_name in predictors:
            results = coefficient_results[pred_name]
            print(f"\n{'='*80}")
            print(f"Linear Regression Coefficients: {pred_name}")
            print(f"{'='*80}\n")
            
            beta_mean = results["beta_mean"]
            beta_std = results["beta_std"]
            beta_lower = results["beta_lower"]
            beta_upper = results["beta_upper"]
            feature_names = results["feature_names"]
            
            output_dim = beta_mean.shape[1] if beta_mean.ndim > 1 else 1
            
            for j in range(output_dim):
                if output_dim > 1:
                    print(f"\nOutcome dimension {j}:")
                    print("-" * 80)
                
                # Get coefficients for this output dimension
                if beta_mean.ndim > 1:
                    coefs = [(feature_names[i], beta_mean[i, j], beta_std[i, j], 
                             beta_lower[i, j], beta_upper[i, j]) 
                            for i in range(len(feature_names))]
                else:
                    coefs = [(feature_names[i], beta_mean[i], beta_std[i], 
                             beta_lower[i], beta_upper[i]) 
                            for i in range(len(feature_names))]
                
                # Sort by absolute magnitude
                coefs.sort(key=lambda x: abs(x[1]), reverse=True)
                
                # Print top_k or all
                coefs_to_print = coefs[:top_k] if top_k else coefs
                
                print(f"{'Feature':<20} {'Mean':>10} {'Std':>10} {'95% CI':>25}")
                print("-" * 80)
                
                for name, mean, std, lower, upper in coefs_to_print:
                    ci_str = f"[{lower:7.3f}, {upper:7.3f}]"
                    print(f"{name:<20} {mean:10.3f} {std:10.3f} {ci_str:>25}")
                
                # Print intercept
                print("-" * 80)
                if output_dim > 1:
                    intercept_mean = results["intercept_mean"][j]
                    intercept_std = results["intercept_std"][j]
                    intercept_lower = results["intercept_lower"][j]
                    intercept_upper = results["intercept_upper"][j]
                else:
                    intercept_mean = float(results["intercept_mean"])
                    intercept_std = float(results["intercept_std"])
                    intercept_lower = float(results["intercept_lower"])
                    intercept_upper = float(results["intercept_upper"])
                ci_str = f"[{intercept_lower:7.3f}, {intercept_upper:7.3f}]"
                print(f"{'(Intercept)':<20} {intercept_mean:10.3f} {intercept_std:10.3f} {ci_str:>25}")
                
                # Print sigma if available
                if "sigma_mean" in results:
                    if output_dim > 1:
                        sigma_mean = results["sigma_mean"][j]
                        sigma_std = results["sigma_std"][j]
                        sigma_lower = results["sigma_lower"][j]
                        sigma_upper = results["sigma_upper"][j]
                    else:
                        sigma_mean = float(results["sigma_mean"])
                        sigma_std = float(results["sigma_std"])
                        sigma_lower = float(results["sigma_lower"])
                        sigma_upper = float(results["sigma_upper"])
                    ci_str = f"[{sigma_lower:7.3f}, {sigma_upper:7.3f}]"
                    print(f"{'(Residual SD)':<20} {sigma_mean:10.3f} {sigma_std:10.3f} {ci_str:>25}")

# -----------------------------------------------------------------------------
# Alignment and coverage utilities
# -----------------------------------------------------------------------------

def _linear_align(source: np.ndarray, target: np.ndarray) -> Tuple[np.ndarray, float, float]:
    """Affine-align a 1-D source vector onto a target via least-squares."""
    x = source.ravel()
    y = target.ravel()
    x_mean, y_mean = x.mean(), y.mean()
    cov = ((x - x_mean) * (y - y_mean)).mean()
    var = ((x - x_mean) ** 2).mean()
    slope = cov / var if var > 0 else 0.0
    intercept = y_mean - slope * x_mean
    return slope * x + intercept, slope, intercept

def align_theta_samples(theta_samples: np.ndarray, ideal_points: np.ndarray) -> np.ndarray:
    """Vectorised alignment of every posterior draw to the fixed `ideal_points`."""
    if theta_samples.ndim == 3:
        theta_samples = theta_samples[..., 0]

    S, D = theta_samples.shape
    aligned = np.empty_like(theta_samples)

    for s in range(S):
        aligned[s], _, _ = _linear_align(theta_samples[s], ideal_points)
    return aligned

def coverage_rate(theta_samples: np.ndarray,
                  ideal_points:  np.ndarray,
                  *, ci_level: float = 0.95
                 ) -> Tuple[float, int]:
    """Coverage rate after global affine alignment."""
    theta_samples = np.asarray(theta_samples.squeeze())
    ideal_points  = np.asarray(ideal_points).ravel()

    theta_mean = theta_samples.mean(axis=0)
    _, slope, intercept = _linear_align(theta_mean, ideal_points)

    aligned = theta_samples * slope + intercept

    mean = aligned.mean(axis=0)
    sd   = aligned.std(axis=0, ddof=0)

    from scipy.stats import norm
    z = norm.ppf(0.5 + ci_level / 2.0)

    lower, upper = mean - z * sd, mean + z * sd
    within = (ideal_points >= lower) & (ideal_points <= upper)

    count    = int(within.sum())
    coverage = 100.0 * count / ideal_points.size
    return coverage, count