"""
Comprehensive Test Script for Ideal Point Model
================================================

This script tests a 1D ideal point model with:
- Three modalities: votes (Bernoulli), speeches (Gaussian), surveys (categorical)
- Three types of covariates: prior, decoder, and predictor
- Two outcome predictions related to ideal points (linear regression)

The test includes:
1. Data generation with known parameters
2. Model estimation with PoE fusion and KL annealing
3. Confidence intervals for estimated ideal points
4. Coefficient extraction and comparison for structured prior
5. Coefficient extraction and comparison for outcome predictors
"""

import sys
import os
import numpy as np
import jax
import jax.numpy as jnp
from jax import random
import matplotlib.pyplot as plt
from scipy.stats import linregress
from typing import Dict, Tuple, Any
import warnings
warnings.filterwarnings("ignore")

# Import model and utilities
sys.path.insert(0, os.path.join(os.path.dirname(__file__), '..'))
from models import DeepLatent, coverage_rate, _linear_align


def generate_ideal_point_data(
    n_individuals: int = 500,
    n_votes: int = 80,
    n_speeches: int = 50,
    n_survey_questions: int = 20,
    n_prior_covariates: int = 3,
    n_decoder_covariates: int = 3,
    n_predictor_covariates: int = 3,
    seed: int = 42
) -> Dict[str, Any]:
    """
    Generate synthetic data for ideal point estimation.
    
    Returns:
        Dictionary containing:
        - data: modality data (votes, speeches, surveys)
        - covariates_prior: covariates for structured prior
        - covariates_decoder: covariates for decoder
        - covariates_predictor: covariates for predictor
        - outcomes: two outcomes for linear regression
        - true_params: ground truth parameters
    """
    np.random.seed(seed)
    rng = random.PRNGKey(seed)
    
    print("=" * 80)
    print("DATA GENERATION")
    print("=" * 80)
    
    # ============================================
    # 1. Generate True Ideal Points (1D)
    # ============================================
    true_ideal_points = np.random.normal(0, 1.0, n_individuals)
    print(f"\n✓ Generated {n_individuals} ideal points")
    print(f"  Range: [{true_ideal_points.min():.2f}, {true_ideal_points.max():.2f}]")
    print(f"  Mean: {true_ideal_points.mean():.2f}, SD: {true_ideal_points.std():.2f}")
    
    # ============================================
    # 2. Generate Covariates
    # ============================================
    # Prior covariates (e.g., party affiliation, region)
    covariates_prior = np.random.randn(n_individuals, n_prior_covariates).astype(np.float32)
    # Standardize
    covariates_prior = (covariates_prior - covariates_prior.mean(axis=0)) / covariates_prior.std(axis=0)
    
    # Decoder covariates (e.g., session, time period)
    covariates_decoder = np.random.randn(n_individuals, n_decoder_covariates).astype(np.float32)
    covariates_decoder = (covariates_decoder - covariates_decoder.mean(axis=0)) / covariates_decoder.std(axis=0)
    
    # Predictor covariates (e.g., demographics, district characteristics)
    covariates_predictor = np.random.randn(n_individuals, n_predictor_covariates).astype(np.float32)
    covariates_predictor = (covariates_predictor - covariates_predictor.mean(axis=0)) / covariates_predictor.std(axis=0)
    
    print(f"\n✓ Generated covariates:")
    print(f"  Prior covariates: {covariates_prior.shape}")
    print(f"  Decoder covariates: {covariates_decoder.shape}")
    print(f"  Predictor covariates: {covariates_predictor.shape}")
    
    # ============================================
    # 3. Structured Prior: theta ~ N(X @ gamma, sigma^2)
    # ============================================
    # True gamma coefficients (effect of covariates on ideal points)
    true_gamma_prior = np.array([0.5, -0.3, 0.4])[:n_prior_covariates].reshape(-1, 1)
    prior_mean = covariates_prior @ true_gamma_prior
    prior_mean = prior_mean.squeeze()
    
    # Add noise to create actual ideal points
    prior_sigma = 0.8
    true_ideal_points = prior_mean + np.random.normal(0, prior_sigma, n_individuals)
    
    print(f"\n✓ Structured prior parameters:")
    print(f"  True gamma (prior): {true_gamma_prior.squeeze()}")
    print(f"  Prior sigma: {prior_sigma}")
    
    # ============================================
    # 4. Generate Voting Records (Bernoulli)
    # ============================================
    # Bill parameters: difficulty and discrimination
    bill_difficulty = np.random.normal(0, 0.6, n_votes)
    bill_discrimination = np.random.uniform(0.8, 1.5, n_votes)
    
    # 2PL IRT model: P(vote=1) = sigmoid(discrimination * theta - difficulty)
    vote_logits = (true_ideal_points[:, None] * bill_discrimination[None, :] - 
                   bill_difficulty[None, :])
    vote_probs = 1 / (1 + np.exp(-vote_logits))
    votes = np.random.binomial(1, vote_probs).astype(np.float32)
    
    print(f"\n✓ Generated voting records: {votes.shape}")
    print(f"  Average vote rate: {votes.mean():.2%}")
    
    # ============================================
    # 5. Generate Speech Embeddings (Gaussian)
    # ============================================
    # Speeches are influenced by ideal points and decoder covariates
    speech_loading = np.random.normal(0, 0.5, (1, n_speeches))
    speech_decoder_coef = np.random.normal(0, 0.3, (n_decoder_covariates, n_speeches))
    
    speeches = (true_ideal_points[:, None] @ speech_loading + 
                covariates_decoder @ speech_decoder_coef +
                np.random.normal(0, 0.2, (n_individuals, n_speeches)))
    speeches = speeches.astype(np.float32)
    
    print(f"\n✓ Generated speech embeddings: {speeches.shape}")
    
    # ============================================
    # 6. Generate Survey Responses (Categorical)
    # ============================================
    n_categories = 5  # 5-point Likert scale
    survey_responses = []
    
    for q in range(n_survey_questions):
        # Question parameters
        question_discrimination = np.random.uniform(0.5, 1.2)
        question_difficulty = np.random.normal(0, 0.5)
        
        # Ordered thresholds
        raw_thresholds = np.random.normal(0, 0.4, n_categories - 1)
        thresholds = np.sort(raw_thresholds)
        
        # Ordinal model
        eta = true_ideal_points * question_discrimination - question_difficulty
        cum_probs = 1 / (1 + np.exp(thresholds[None, :] - eta[:, None]))
        
        # Category probabilities
        probs = np.zeros((n_individuals, n_categories))
        probs[:, 0] = cum_probs[:, 0]
        for k in range(1, n_categories - 1):
            probs[:, k] = cum_probs[:, k] - cum_probs[:, k-1]
        probs[:, -1] = 1 - cum_probs[:, -1]
        
        # Ensure valid probabilities
        probs = np.maximum(probs, 1e-8)
        probs = probs / probs.sum(axis=1, keepdims=True)
        
        # Sample responses
        responses = np.array([np.random.choice(n_categories, p=probs[i]) 
                             for i in range(n_individuals)])
        survey_responses.append(responses)
    
    surveys = np.array(survey_responses).T.astype(np.int32)
    categories_per_question = [n_categories] * n_survey_questions
    
    print(f"\n✓ Generated survey responses: {surveys.shape}")
    print(f"  Categories per question: {n_categories}")
    
    # ============================================
    # 7. Generate Outcomes (Linear Regression)
    # ============================================
    # Outcome 1: Policy preference (continuous)
    # y1 = beta1_0 + beta1_theta * theta + beta1_cov @ covariates + noise
    true_beta1 = np.array([1.2, -0.5, 0.3])[:n_predictor_covariates]
    true_beta1_theta = 0.8
    true_intercept1 = 0.2
    true_sigma1 = 0.4
    
    outcome1 = (true_intercept1 + 
                true_beta1_theta * true_ideal_points + 
                covariates_predictor @ true_beta1 + 
                np.random.normal(0, true_sigma1, n_individuals))
    
    # Outcome 2: Vote share (continuous)
    # y2 = beta2_0 + beta2_theta * theta + beta2_cov @ covariates + noise
    true_beta2 = np.array([0.6, 0.4, -0.7])[:n_predictor_covariates]
    true_beta2_theta = -0.9
    true_intercept2 = 0.5
    true_sigma2 = 0.5
    
    outcome2 = (true_intercept2 + 
                true_beta2_theta * true_ideal_points + 
                covariates_predictor @ true_beta2 + 
                np.random.normal(0, true_sigma2, n_individuals))
    
    print(f"\n✓ Generated outcomes:")
    print(f"  Outcome 1 (policy preference): mean={outcome1.mean():.2f}, sd={outcome1.std():.2f}")
    print(f"  Outcome 2 (vote share): mean={outcome2.mean():.2f}, sd={outcome2.std():.2f}")
    print(f"\n  True coefficients (Outcome 1):")
    print(f"    Intercept: {true_intercept1:.3f}")
    print(f"    Beta_theta: {true_beta1_theta:.3f}")
    print(f"    Beta_cov: {true_beta1}")
    print(f"    Sigma: {true_sigma1:.3f}")
    print(f"\n  True coefficients (Outcome 2):")
    print(f"    Intercept: {true_intercept2:.3f}")
    print(f"    Beta_theta: {true_beta2_theta:.3f}")
    print(f"    Beta_cov: {true_beta2}")
    print(f"    Sigma: {true_sigma2:.3f}")
    
    # ============================================
    # Return data dictionary
    # ============================================
    return {
        "data": {
            "votes": jnp.array(votes),
            "speeches": jnp.array(speeches),
            "surveys": jnp.array(surveys),
            "categories_per_question": categories_per_question
        },
        "covariates_prior": jnp.array(covariates_prior),
        "covariates_decoder": jnp.array(covariates_decoder),
        "covariates_predictor": jnp.array(covariates_predictor),
        "outcomes": {
            "policy_preference": jnp.array(outcome1.astype(np.float32)),
            "vote_share": jnp.array(outcome2.astype(np.float32))
        },
        "true_params": {
            "ideal_points": true_ideal_points,
            "gamma_prior": true_gamma_prior.squeeze(),
            "prior_sigma": prior_sigma,
            "beta1": true_beta1,
            "beta1_theta": true_beta1_theta,
            "intercept1": true_intercept1,
            "sigma1": true_sigma1,
            "beta2": true_beta2,
            "beta2_theta": true_beta2_theta,
            "intercept2": true_intercept2,
            "sigma2": true_sigma2
        }
    }


def estimate_model(dataset: Dict[str, Any], num_steps: int = 20000) -> Tuple[DeepLatent, Dict]:
    """
    Estimate the ideal point model with PoE fusion and KL annealing.
    """
    print("\n" + "=" * 80)
    print("MODEL ESTIMATION")
    print("=" * 80)
    
    # Extract data
    data = dataset["data"]
    covariates_prior = dataset["covariates_prior"]
    covariates_decoder = dataset["covariates_decoder"]
    covariates_predictor = dataset["covariates_predictor"]
    outcomes = dataset["outcomes"]
    
    # Get dimensions
    n_individuals = data["votes"].shape[0]
    n_votes = data["votes"].shape[1]
    n_speeches = data["speeches"].shape[1]
    n_survey_questions = len(data["categories_per_question"])
    n_prior_cov = covariates_prior.shape[1]
    n_decoder_cov = covariates_decoder.shape[1]
    n_predictor_cov = covariates_predictor.shape[1]
    
    print(f"\n✓ Data dimensions:")
    print(f"  Individuals: {n_individuals}")
    print(f"  Votes: {n_votes}")
    print(f"  Speech embeddings: {n_speeches}")
    print(f"  Survey questions: {n_survey_questions}")
    
    # ============================================
    # Define Modality Specifications
    # ============================================
    total_survey_logits = sum(data["categories_per_question"])
    
    modality_specs = {
        "votes": {
            "type": "bernoulli",
            "dim": n_votes
        },
        "speeches": {
            "type": "gaussian",
            "dim": n_speeches
        },
        "surveys": {
            "type": "categorical_multi",
            "dim": total_survey_logits,
            "likelihood_params": {
                "categories_per_question": data["categories_per_question"]
            }
        }
    }
    
    # ============================================
    # Define Predictor Specifications (Linear Regression)
    # ============================================
    predictor_specs = {
        "policy_preference": {
            "type": "linear_regression",
            "outcome_type": "gaussian",
            "output_dim": 1,
            "beta_prior_scale": 1.0,
            "sigma_prior_scale": 1.0,
            "intercept_prior_scale": 1.0,
            "feature_names": ["theta"] + [f"cov_pred_{i+1}" for i in range(n_predictor_cov)]
        },
        "vote_share": {
            "type": "linear_regression",
            "outcome_type": "gaussian",
            "output_dim": 1,
            "beta_prior_scale": 1.0,
            "sigma_prior_scale": 1.0,
            "intercept_prior_scale": 1.0,
            "feature_names": ["theta"] + [f"cov_pred_{i+1}" for i in range(n_predictor_cov)]
        }
    }
    
    print(f"\n✓ Model specifications:")
    print(f"  Fusion strategy: Product of Experts (PoE)")
    print(f"  KL annealing: Linear schedule over {num_steps} steps")
    print(f"  Latent dimension: 1 (ideal points)")
    print(f"  Structured prior: Yes ({n_prior_cov} covariates)")
    print(f"  Outcome predictors: 2 (linear regression)")
    
    # ============================================
    # Create KL Annealing Schedule
    # ============================================
    warmup_steps = 1000
    kl_weights_used = []
    
    def kl_schedule_fn(step):
        """Linear KL annealing schedule."""
        if step < warmup_steps:
            weight = step / warmup_steps
        else:
            weight = 1.0
        return weight
    
    # Pre-compute all KL weights for tracking
    for step in range(num_steps):
        kl_weights_used.append(kl_schedule_fn(step))
    
    # ============================================
    # Create Model with PoE Fusion
    # ============================================
    # We'll use a simpler approach: train with different KL weights in phases
    # Phase 1: warmup with gradually increasing KL weight
    # Phase 2: full training with KL weight = 1.0
    
    model = DeepLatent(
        modality_specs=modality_specs,
        latent_dim=1,
        encoder_hidden_dims=(64, 32),
        decoder_hidden_dims=(32, 64),
        activation="relu",
        dropout_rate=0.1,
        fusion_strategy="poe",  # Product of Experts
        use_structured_prior=True,
        prior_covariate_dim=n_prior_cov,
        decoder_covariate_dim=n_decoder_cov,
        predictor_covariate_dim=n_predictor_cov,
        predictor_specs=predictor_specs,
        kl_weight=0.01  # Start with low KL weight for warmup
    )
    
    print(f"\n✓ Model created successfully")
    
    # ============================================
    # Train with KL Annealing (in phases)
    # ============================================
    print(f"\n✓ Training model with KL annealing...")
    print(f"  Total steps: {num_steps}")
    print(f"  Warmup steps: {warmup_steps} (KL weight: 0 → 1)")
    print(f"  Full training steps: {num_steps - warmup_steps} (KL weight: 1)")
    print(f"  Learning rate: 1e-3")
    
    # Prepare data for training (remove metadata)
    train_data = {
        "votes": data["votes"],
        "speeches": data["speeches"],
        "surveys": data["surveys"]
    }
    
    all_losses = []
    
    # Phase 1: Warmup with gradually increasing KL weight
    print(f"\n  Phase 1: Warmup training ({warmup_steps} steps)...")
    num_warmup_phases = 10
    steps_per_phase = warmup_steps // num_warmup_phases
    
    for phase in range(num_warmup_phases):
        kl_weight = (phase + 1) / num_warmup_phases
        model.kl_weight = kl_weight
        
        model.fit(
            data=train_data,
            covariates_prior=covariates_prior,
            covariates_decoder=covariates_decoder,
            covariates_predictor=covariates_predictor,
            outcomes=outcomes,
            num_steps=steps_per_phase,
            lr=1e-3,
            seed=42 + phase,
            progress_bar=False
        )
        
        all_losses.extend(model.losses)
        
        if (phase + 1) % 3 == 0:
            print(f"    Phase {phase + 1}/{num_warmup_phases}, KL weight: {kl_weight:.2f}, Loss: {model.losses[-1]:.2f}")
    
    # Phase 2: Full training with KL weight = 1.0
    print(f"\n  Phase 2: Full training ({num_steps - warmup_steps} steps)...")
    model.kl_weight = 1.0
    
    model.fit(
        data=train_data,
        covariates_prior=covariates_prior,
        covariates_decoder=covariates_decoder,
        covariates_predictor=covariates_predictor,
        outcomes=outcomes,
        num_steps=num_steps - warmup_steps,
        lr=1e-3,
        seed=42 + num_warmup_phases,
        progress_bar=True
    )
    
    all_losses.extend(model.losses)
    
    print(f"\n✓ Training completed!")
    print(f"  Final loss: {all_losses[-1]:.2f}")
    
    # Store losses for plotting
    model.losses = all_losses
    
    return model, {"losses": all_losses, "kl_weights": kl_weights_used}


def analyze_ideal_points(model: DeepLatent, dataset: Dict[str, Any], num_samples: int = 2000, num_samples_per_chunk: int = 100, batch_size: int = 250):
    """
    Extract and analyze estimated ideal points with confidence intervals.
    """
    print("\n" + "=" * 80)
    print("IDEAL POINT ANALYSIS")
    print("=" * 80)
    
    # Extract data
    data = {
        "votes": dataset["data"]["votes"],
        "speeches": dataset["data"]["speeches"],
        "surveys": dataset["data"]["surveys"]
    }
    true_ideal_points = dataset["true_params"]["ideal_points"]
    
    # Sample posterior
    print(f"\n✓ Sampling posterior ({num_samples} samples with chunks of {num_samples_per_chunk}, batch_size={batch_size})...")
    posterior_samples = model.sample_posterior(
        data=data,
        covariates_prior=dataset["covariates_prior"],
        covariates_decoder=dataset["covariates_decoder"],
        covariates_predictor=dataset["covariates_predictor"],
        outcomes=dataset["outcomes"],
        num_samples=num_samples,
        num_samples_per_chunk=num_samples_per_chunk,
        batch_size=batch_size,
        seed=123
    )
    
    # Extract ideal point samples
    theta_samples = np.asarray(posterior_samples["z"])
    if theta_samples.ndim == 3:
        theta_samples = theta_samples.squeeze(-1)
    
    print(f"  Theta samples shape: {theta_samples.shape}")
    
    # Compute posterior statistics
    theta_mean = theta_samples.mean(axis=0)
    theta_std = theta_samples.std(axis=0)
    
    # Align to true ideal points
    theta_aligned, slope, intercept = _linear_align(theta_mean, true_ideal_points)
    
    # Align all samples for CI calculation
    aligned_samples = theta_samples * slope + intercept
    
    # Compute 95% credible intervals
    theta_lower = np.percentile(aligned_samples, 2.5, axis=0)
    theta_upper = np.percentile(aligned_samples, 97.5, axis=0)
    
    # Compute metrics
    correlation = np.corrcoef(true_ideal_points, theta_aligned)[0, 1]
    fit = linregress(true_ideal_points, theta_aligned)
    r_squared = fit.rvalue ** 2
    rmse = np.sqrt(np.mean((theta_aligned - true_ideal_points) ** 2))
    
    # Coverage rate
    coverage, n_in = coverage_rate(theta_samples, true_ideal_points, ci_level=0.95)
    
    print(f"\n✓ Alignment results:")
    print(f"  Slope: {slope:.3f}")
    print(f"  Intercept: {intercept:.3f}")
    
    print(f"\n✓ Estimation quality:")
    print(f"  Correlation: {correlation:.3f}")
    print(f"  R²: {r_squared:.3f}")
    print(f"  RMSE: {rmse:.3f}")
    print(f"  95% Coverage: {coverage:.1f}% ({n_in}/{len(true_ideal_points)} in CI)")
    
    # Create visualization
    plot_ideal_points(true_ideal_points, theta_aligned, theta_lower, theta_upper, 
                     correlation, r_squared, coverage)
    
    return {
        "theta_mean": theta_mean,
        "theta_aligned": theta_aligned,
        "theta_std": theta_std,
        "theta_lower": theta_lower,
        "theta_upper": theta_upper,
        "correlation": correlation,
        "r_squared": r_squared,
        "rmse": rmse,
        "coverage": coverage,
        "slope": slope,
        "intercept": intercept
    }


def plot_ideal_points(true_theta, estimated_theta, lower_ci, upper_ci, 
                     correlation, r_squared, coverage):
    """Create visualization of ideal point estimation."""
    fig, axes = plt.subplots(1, 3, figsize=(18, 5))
    
    # Plot 1: Scatter plot with regression line
    axes[0].scatter(true_theta, estimated_theta, alpha=0.6, s=30, color='steelblue')
    
    # Add regression line
    fit = linregress(true_theta, estimated_theta)
    x_line = np.array([true_theta.min(), true_theta.max()])
    y_line = fit.slope * x_line + fit.intercept
    axes[0].plot(x_line, y_line, 'r--', linewidth=2, label='Fitted line')
    
    # Add identity line
    axes[0].plot(x_line, x_line, 'k--', alpha=0.3, linewidth=1, label='Identity')
    
    axes[0].set_xlabel('True Ideal Points', fontsize=12)
    axes[0].set_ylabel('Estimated Ideal Points', fontsize=12)
    axes[0].set_title(f'Ideal Point Recovery\nCorr = {correlation:.3f}, R² = {r_squared:.3f}', 
                     fontsize=12)
    axes[0].legend()
    axes[0].grid(True, alpha=0.3)
    
    # Plot 2: Confidence intervals (sorted by true values)
    sort_idx = np.argsort(true_theta)
    true_sorted = true_theta[sort_idx]
    est_sorted = estimated_theta[sort_idx]
    lower_sorted = lower_ci[sort_idx]
    upper_sorted = upper_ci[sort_idx]
    
    x_pos = np.arange(len(true_theta))
    axes[1].fill_between(x_pos, lower_sorted, upper_sorted, 
                         alpha=0.3, color='steelblue', label='95% CI')
    axes[1].plot(x_pos, est_sorted, 'o', markersize=3, color='steelblue', 
                label='Posterior mean', alpha=0.7)
    axes[1].plot(x_pos, true_sorted, 's', markersize=3, color='red', 
                label='True values', alpha=0.7)
    
    axes[1].set_xlabel('Individuals (sorted by true ideal point)', fontsize=12)
    axes[1].set_ylabel('Ideal Points', fontsize=12)
    axes[1].set_title(f'Confidence Intervals\nCoverage = {coverage:.1f}%', fontsize=12)
    axes[1].legend()
    axes[1].grid(True, alpha=0.3)
    
    # Plot 3: Error distribution
    errors = estimated_theta - true_theta
    axes[2].hist(errors, bins=40, alpha=0.7, color='steelblue', edgecolor='black')
    axes[2].axvline(0, color='red', linestyle='--', linewidth=2, label='Zero error')
    axes[2].axvline(errors.mean(), color='orange', linestyle='--', linewidth=2, 
                   label=f'Mean error = {errors.mean():.3f}')
    
    axes[2].set_xlabel('Estimation Error', fontsize=12)
    axes[2].set_ylabel('Frequency', fontsize=12)
    axes[2].set_title(f'Error Distribution\nSD = {errors.std():.3f}', fontsize=12)
    axes[2].legend()
    axes[2].grid(True, alpha=0.3)
    
    plt.tight_layout()
    plt.savefig('ideal_point_estimation.png', dpi=150, bbox_inches='tight')
    print(f"\n✓ Plot saved: ideal_point_estimation.png")
    plt.show()


def analyze_coefficients(model: DeepLatent, dataset: Dict[str, Any], num_samples: int = 2000, num_samples_per_chunk: int = 100, batch_size: int = 250):
    """
    Analyze and compare estimated coefficients to true values.
    """
    print("\n" + "=" * 80)
    print("COEFFICIENT ANALYSIS")
    print("=" * 80)
    
    # Extract data
    data = {
        "votes": dataset["data"]["votes"],
        "speeches": dataset["data"]["speeches"],
        "surveys": dataset["data"]["surveys"]
    }
    true_params = dataset["true_params"]
    
    # ============================================
    # 1. Analyze Structured Prior Coefficients (gamma)
    # ============================================
    print("\n" + "-" * 80)
    print("1. STRUCTURED PRIOR COEFFICIENTS")
    print("-" * 80)
    
    # Sample posterior
    print(f"\n✓ Sampling posterior ({num_samples} samples with chunks of {num_samples_per_chunk}, batch_size={batch_size})...")
    posterior_samples = model.sample_posterior(
        data=data,
        covariates_prior=dataset["covariates_prior"],
        covariates_decoder=dataset["covariates_decoder"],
        covariates_predictor=dataset["covariates_predictor"],
        outcomes=dataset["outcomes"],
        num_samples=num_samples,
        num_samples_per_chunk=num_samples_per_chunk,
        batch_size=batch_size,
        seed=456
    )
    
    # Extract gamma samples
    if "gamma" in posterior_samples:
        gamma_samples = np.asarray(posterior_samples["gamma"])  # (num_samples, n_cov, latent_dim)
        gamma_samples = gamma_samples.squeeze(-1)  # Remove latent_dim dimension (it's 1)
        
        gamma_mean = gamma_samples.mean(axis=0)
        gamma_std = gamma_samples.std(axis=0)
        gamma_lower = np.percentile(gamma_samples, 2.5, axis=0)
        gamma_upper = np.percentile(gamma_samples, 97.5, axis=0)
        
        true_gamma = true_params["gamma_prior"]
        
        print(f"\nGamma coefficients (effect of covariates on ideal points):")
        print(f"{'Covariate':<15} {'True':>10} {'Estimated':>10} {'Std':>10} {'95% CI':>25}")
        print("-" * 80)
        
        for i in range(len(gamma_mean)):
            ci_str = f"[{gamma_lower[i]:7.3f}, {gamma_upper[i]:7.3f}]"
            in_ci = "✓" if gamma_lower[i] <= true_gamma[i] <= gamma_upper[i] else "✗"
            print(f"cov_prior_{i+1:<7} {true_gamma[i]:10.3f} {gamma_mean[i]:10.3f} "
                  f"{gamma_std[i]:10.3f} {ci_str:>25} {in_ci}")
        
        # Calculate coverage
        gamma_coverage = np.mean([gamma_lower[i] <= true_gamma[i] <= gamma_upper[i] 
                                 for i in range(len(true_gamma))]) * 100
        print(f"\nGamma coefficient coverage: {gamma_coverage:.1f}%")
    else:
        print("Warning: Gamma samples not found in posterior")
    
    # ============================================
    # 2. Analyze Outcome Predictor Coefficients
    # ============================================
    print("\n" + "-" * 80)
    print("2. OUTCOME PREDICTOR COEFFICIENTS")
    print("-" * 80)
    
    # Get linear coefficients - must include all covariates since model has structured prior
    print(f"\n✓ Extracting linear coefficients ({num_samples} samples with chunks of {num_samples_per_chunk}, batch_size={batch_size})...")
    coef_results = model.get_linear_coefficients(
        data=data,
        covariates_prior=dataset["covariates_prior"],
        covariates_decoder=dataset["covariates_decoder"],
        covariates_predictor=dataset["covariates_predictor"],
        outcomes=dataset["outcomes"],
        num_samples=num_samples,
        ci_level=0.95,
        num_samples_per_chunk=num_samples_per_chunk,
        batch_size=batch_size,
        seed=789
    )
    
    # Analyze Outcome 1: Policy Preference
    print("\n" + "=" * 80)
    print("OUTCOME 1: Policy Preference")
    print("=" * 80)
    
    analyze_outcome_coefficients(
        coef_results["policy_preference"],
        true_beta=np.concatenate([[true_params["beta1_theta"]], true_params["beta1"]]),
        true_intercept=true_params["intercept1"],
        true_sigma=true_params["sigma1"]
    )
    
    # Analyze Outcome 2: Vote Share
    print("\n" + "=" * 80)
    print("OUTCOME 2: Vote Share")
    print("=" * 80)
    
    analyze_outcome_coefficients(
        coef_results["vote_share"],
        true_beta=np.concatenate([[true_params["beta2_theta"]], true_params["beta2"]]),
        true_intercept=true_params["intercept2"],
        true_sigma=true_params["sigma2"]
    )
    
    return coef_results


def analyze_outcome_coefficients(results: Dict, true_beta: np.ndarray, 
                                 true_intercept: float, true_sigma: float):
    """Analyze and print outcome coefficient results."""
    beta_mean = np.asarray(results["beta_mean"]).squeeze()
    beta_std = np.asarray(results["beta_std"]).squeeze()
    beta_lower = np.asarray(results["beta_lower"]).squeeze()
    beta_upper = np.asarray(results["beta_upper"]).squeeze()
    feature_names = results["feature_names"]
    
    # Handle intercept (could be scalar or 1D array with single element)
    intercept_mean = float(np.asarray(results["intercept_mean"]).squeeze())
    intercept_std = float(np.asarray(results["intercept_std"]).squeeze())
    intercept_lower = float(np.asarray(results["intercept_lower"]).squeeze())
    intercept_upper = float(np.asarray(results["intercept_upper"]).squeeze())
    
    # Handle sigma (could be scalar or 1D array with single element)
    sigma_mean = float(np.asarray(results["sigma_mean"]).squeeze())
    sigma_std = float(np.asarray(results["sigma_std"]).squeeze())
    sigma_lower = float(np.asarray(results["sigma_lower"]).squeeze())
    sigma_upper = float(np.asarray(results["sigma_upper"]).squeeze())
    
    print(f"\n{'Feature':<15} {'True':>10} {'Estimated':>10} {'Std':>10} {'95% CI':>25} {'In CI'}")
    print("-" * 90)
    
    # Beta coefficients
    for i, name in enumerate(feature_names):
        ci_str = f"[{beta_lower[i]:7.3f}, {beta_upper[i]:7.3f}]"
        in_ci = "✓" if beta_lower[i] <= true_beta[i] <= beta_upper[i] else "✗"
        print(f"{name:<15} {true_beta[i]:10.3f} {beta_mean[i]:10.3f} "
              f"{beta_std[i]:10.3f} {ci_str:>25} {in_ci}")
    
    # Intercept
    print("-" * 90)
    ci_str = f"[{intercept_lower:7.3f}, {intercept_upper:7.3f}]"
    in_ci = "✓" if intercept_lower <= true_intercept <= intercept_upper else "✗"
    print(f"{'(Intercept)':<15} {true_intercept:10.3f} {intercept_mean:10.3f} "
          f"{intercept_std:10.3f} {ci_str:>25} {in_ci}")
    
    # Sigma
    ci_str = f"[{sigma_lower:7.3f}, {sigma_upper:7.3f}]"
    in_ci = "✓" if sigma_lower <= true_sigma <= sigma_upper else "✗"
    print(f"{'(Residual SD)':<15} {true_sigma:10.3f} {sigma_mean:10.3f} "
          f"{sigma_std:10.3f} {ci_str:>25} {in_ci}")
    
    # Calculate coverage
    all_params = list(true_beta) + [true_intercept, true_sigma]
    all_lowers = list(beta_lower) + [intercept_lower, sigma_lower]
    all_uppers = list(beta_upper) + [intercept_upper, sigma_upper]
    
    coverage = np.mean([all_lowers[i] <= all_params[i] <= all_uppers[i] 
                       for i in range(len(all_params))]) * 100
    
    print(f"\nOverall parameter coverage: {coverage:.1f}%")


def main():
    """Main test script."""
    print("=" * 80)
    print("COMPREHENSIVE IDEAL POINT MODEL TEST")
    print("=" * 80)
    print("\nThis test evaluates:")
    print("  • 1D ideal point estimation")
    print("  • Three modalities: votes, speeches, surveys")
    print("  • Product of Experts (PoE) fusion")
    print("  • Linear KL annealing schedule")
    print("  • Structured prior with covariates")
    print("  • Two linear regression outcomes")
    print("  • Confidence interval coverage")
    print("  • Coefficient recovery")
    
    # Generate data
    dataset = generate_ideal_point_data(
        n_individuals=2500,
        n_votes=500,
        n_speeches=500,
        n_survey_questions=100,
        n_prior_covariates=3,
        n_decoder_covariates=3,
        n_predictor_covariates=3,
        seed=42
    )
    
    # Estimate model
    model, training_info = estimate_model(dataset, num_steps=5000)
    
    # Analyze ideal points (2000 samples in chunks of 100, 250 individuals per batch to avoid GPU OOM)
    ideal_point_results = analyze_ideal_points(model, dataset, num_samples=2000, num_samples_per_chunk=100, batch_size=250)
    
    # Analyze coefficients (2000 samples in chunks of 100, 250 individuals per batch to avoid GPU OOM)
    coefficient_results = analyze_coefficients(model, dataset, num_samples=2000, num_samples_per_chunk=100, batch_size=250)
    
    # Summary
    print("\n" + "=" * 80)
    print("TEST SUMMARY")
    print("=" * 80)
    
    print(f"\n✓ Ideal Point Estimation:")
    print(f"  • Correlation: {ideal_point_results['correlation']:.3f}")
    print(f"  • R²: {ideal_point_results['r_squared']:.3f}")
    print(f"  • RMSE: {ideal_point_results['rmse']:.3f}")
    print(f"  • 95% Coverage: {ideal_point_results['coverage']:.1f}%")
    
    print(f"\n✓ Model converged successfully")
    print(f"  • Final loss: {training_info['losses'][-1]:.2f}")
    print(f"  • KL annealing: Linear schedule (0 → 1 over 1000 steps)")
    
    print(f"\n✓ All analyses completed")
    print(f"  • Ideal points extracted with confidence intervals")
    print(f"  • Prior coefficients (gamma) compared to true values")
    print(f"  • Outcome coefficients compared to true values")
    
    print("\n" + "=" * 80)
    print("TEST COMPLETED SUCCESSFULLY!")
    print("=" * 80)
    
    return model, dataset, ideal_point_results, coefficient_results


if __name__ == "__main__":
    model, dataset, ideal_point_results, coefficient_results = main()
