# Mixture of Experts Test Script: Votes + Surveys
# ===============================================
# Test script for Mixture of Experts fusion strategies combining voting records
# and survey responses to estimate political ideal points.

import numpy as np
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")

def create_votes_and_surveys_data(
    num_politicians: int = 500, 
    num_bills: int = 100,
    num_survey_questions: int = 30,
    seed: int = 42
) -> Tuple[Dict[str, jnp.ndarray], jnp.ndarray]:
    """
    Create synthetic data with both voting records and survey responses.
    
    Args:
        num_politicians: Number of politicians/individuals
        num_bills: Number of bills/votes
        num_survey_questions: Number of survey questions
        seed: Random seed for reproducibility
        
    Returns:
        data: Dictionary containing votes and survey data
        true_ideal_points: Ground truth 1D ideal points
    """
    np.random.seed(seed)
    
    # Generate true 1D ideal points
    true_ideal_points = np.random.normal(0, 1, num_politicians)
    
    # ===========================================
    # 1. Generate Voting Records (Bernoulli)
    # ===========================================
    # Generate bill parameters using 2PL IRT model
    bill_difficulty = np.random.normal(0, 0.8, num_bills)      # δ (difficulty)
    bill_discrimination = np.random.normal(0, 1.2, num_bills)  # α (discrimination)
    
    # Voting probability: P(vote=1) = sigmoid(α * θ - δ)
    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(float)
    
    # ===========================================
    # 2. Generate Survey Responses (Categorical)
    # ===========================================
    # Each survey question has 3-7 response categories
    survey_responses = []
    categories_per_question = []
    
    for q in range(num_survey_questions):
        # Random number of categories for this question (3-7)
        num_categories = np.random.randint(3, 8)
        categories_per_question.append(num_categories)
        
        # Question parameters
        question_difficulty = np.random.normal(0, 0.5)
        question_discrimination = np.random.normal(0, 1.0)
        
        # Generate ordered thresholds for ordinal responses
        raw_thresholds = np.random.normal(0, 0.3, num_categories - 1)
        thresholds = np.sort(raw_thresholds)
        
        # Compute response probabilities using cumulative logistic model
        eta = true_ideal_points * question_discrimination - question_difficulty
        
        # Cumulative probabilities
        cum_probs = 1 / (1 + np.exp(thresholds[None, :] - eta[:, None]))
        
        # Category probabilities
        probs = np.zeros((num_politicians, num_categories))
        probs[:, 0] = cum_probs[:, 0] if num_categories > 1 else 1.0
        for k in range(1, num_categories - 1):
            probs[:, k] = cum_probs[:, k] - cum_probs[:, k-1]
        if num_categories > 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(num_categories, p=probs[i]) 
                             for i in range(num_politicians)])
        survey_responses.append(responses)
    
    # Stack survey responses: shape (num_politicians, num_questions)
    survey_responses = np.array(survey_responses).T
    
    # Prepare data dictionary
    data = {
        "votes": jnp.array(votes, dtype=jnp.float32),
        "surveys": jnp.array(survey_responses, dtype=jnp.int32),
        "categories_per_question": categories_per_question
    }
    
    return data, jnp.array(true_ideal_points, dtype=jnp.float32)


def create_modality_specs(num_bills: int, categories_per_question: list) -> Dict[str, Dict[str, Any]]:
    """Create modality specifications for votes and surveys."""
    
    # Calculate total logits needed for categorical_multi
    total_logits = sum(categories_per_question)
    
    return {
        "votes": {
            "type": "bernoulli",
            "dim": num_bills,
            "likelihood_params": {}
        },
        "surveys": {
            "type": "categorical_multi", 
            "dim": total_logits,
            "likelihood_params": {
                "categories_per_question": categories_per_question
            }
        }
    }


def test_mixture_of_experts_votes_surveys():
    """
    Test different Mixture of Experts fusion strategies with votes and surveys.
    """
    print("Mixture of Experts Test: Votes + Surveys")
    print("=" * 50)
    
    # Generate data
    print("Generating synthetic data...")
    data, true_ideal_points = create_votes_and_surveys_data(
        num_politicians=800,
        num_bills=200, 
        num_survey_questions=40,
        seed=42
    )
    
    print(f"Data shapes:")
    print(f"  Votes: {data['votes'].shape}")
    print(f"  Surveys: {data['surveys'].shape}")
    print(f"  True ideal points: {true_ideal_points.shape}")
    print(f"  Categories per question (first 10): {data['categories_per_question'][:10]}")
    
    # Create modality specifications
    modality_specs = create_modality_specs(
        num_bills=data['votes'].shape[1],
        categories_per_question=data['categories_per_question']
    )
    
    # Remove metadata from data dict for model training
    model_data = {
        "votes": data["votes"],
        "surveys": data["surveys"]
    }
    
    print(f"\nModel specifications:")
    print(f"  Votes: {modality_specs['votes']}")
    print(f"  Surveys: {modality_specs['surveys']}")
    
    # Test different MoE fusion strategies and baseline
    fusion_strategies = ["concatenate", "moe_average", "moe_gating", "moe_learned"]
    all_results = {}
    
    for fusion_strategy in fusion_strategies:
        print(f"Testing {fusion_strategy}...")
        
        try:
            # Create model with appropriate parameters for MoE strategies
            model_kwargs = {
                "modality_specs": modality_specs,
                "fusion_strategy": fusion_strategy,
                "latent_dim": 1,
                "encoder_hidden_dims": (64, 32),
                "decoder_hidden_dims": (),
                "activation": "relu",
                "dropout_rate": 0.0
            }
            
            # Create and train model
            model = create_multimodal_vae(**model_kwargs)
            
            # Train model with progress bar enabled
            model.fit(
                data=model_data,
                num_steps=6000,
                lr=1e-3,
                seed=42,
                progress_bar=True
            )
            
            # Sample from posterior
            posterior = model.sample_posterior(model_data, num_samples=1000, seed=42)
            
            if "z" not in posterior:
                print(f"Warning: No 'z' samples found for {fusion_strategy}")
                print(f"Available keys: {list(posterior.keys())}")
                continue
                
            # Extract and process latent samples
            theta_samples = np.asarray(posterior["z"])
            if theta_samples.ndim == 3 and theta_samples.shape[-1] == 1:
                theta_samples = theta_samples.squeeze(-1)
            
            # Compute posterior mean and align to truth
            theta_mean = theta_samples.mean(axis=0)
            theta_aligned, slope, intercept = _linear_align(theta_mean, true_ideal_points)
            
            # Compute metrics
            correlation = np.corrcoef(true_ideal_points, theta_aligned)[0, 1]
            fit = linregress(true_ideal_points, theta_aligned)
            r_squared = fit.rvalue**2
            cov, n_in = coverage_rate(theta_samples, true_ideal_points)
            
            # Store results
            all_results[fusion_strategy] = {
                "correlation": correlation,
                "r_squared": r_squared,
                "coverage": cov,
                "n_in_ci": n_in,
                "slope": slope,
                "intercept": intercept,
                "theta_samples": theta_samples,
                "theta_aligned": theta_aligned,
                "theta_mean": theta_mean
            }
            
            # Print results for this strategy
            print(f"  Correlation: {correlation:.3f}, R²: {r_squared:.3f}, Coverage: {cov:.1f}%")
            
        except Exception as e:
            print(f"Error training {fusion_strategy}: {str(e)}")
            continue
    
    if not all_results:
        print("No strategies completed successfully!")
        return {}, data, true_ideal_points
    
    # Print detailed results
    print(f"\n{'-'*60}")
    print("DETAILED RESULTS:")
    print(f"{'-'*60}")
    for strategy, results in all_results.items():
        print(f"{strategy.upper()}:")
        print(f"  Correlation: {results['correlation']:.3f}")
        print(f"  R²: {results['r_squared']:.3f}")
        print(f"  Coverage: {results['coverage']:.1f}% ({results['n_in_ci']}/{len(true_ideal_points)})")
        print(f"  Slope: {results['slope']:.3f}, Intercept: {results['intercept']:.3f}")
        print()
    
    # Create visualization
    print("Creating visualizations...")
    create_moe_comparison_plots(all_results, true_ideal_points, fusion_strategies)
    
    # Print summary comparison
    print_moe_summary_comparison(all_results, fusion_strategies)
    
    return all_results, data, true_ideal_points


def create_moe_comparison_plots(all_results: Dict[str, Dict], 
                               true_ideal_points: np.ndarray, 
                               fusion_strategies: list):
    """Create comparison plots for different MoE fusion strategies."""
    
    # Filter to only successful strategies
    successful_strategies = [s for s in fusion_strategies if s in all_results]
    num_strategies = len(successful_strategies)
    
    if num_strategies == 0:
        print("No successful strategies to plot!")
        return
    
    # Determine plot layout - accommodate all strategies
    cols = min(4, num_strategies)
    rows = 2
    
    fig, axes = plt.subplots(rows, cols, figsize=(5*cols, 12))
    
    # Ensure axes is always a 2D array for consistent indexing
    if cols == 1 and rows == 1:
        # Single subplot case
        axes = np.array([[axes]])
    elif cols == 1:
        # Single column, multiple rows
        axes = axes.reshape(-1, 1)
    elif rows == 1:
        # Single row, multiple columns 
        axes = axes.reshape(1, -1)
    # If both cols > 1 and rows > 1, axes is already 2D
    
    # Color scheme for different MoE strategies
    colors = {
        "concatenate": "blue", 
        "moe_average": "green", 
        "moe_gating": "red", 
        "moe_learned": "orange"
    }
    strategy_names = {
        "concatenate": "Concatenation", 
        "moe_average": "MoE Average", 
        "moe_gating": "MoE Gating", 
        "moe_learned": "MoE Learned"
    }
    
    for i, strategy in enumerate(successful_strategies[:cols]):
        results = all_results[strategy]
        color = colors.get(strategy, 'gray')
        name = strategy_names.get(strategy, strategy.replace('_', ' ').title())
        
        # Plot 1: Scatter plot of estimated vs true
        axes[0, i].scatter(true_ideal_points, results["theta_aligned"], 
                          alpha=0.6, s=15, color=color)
        axes[0, i].plot([true_ideal_points.min(), true_ideal_points.max()], 
                       [true_ideal_points.min(), true_ideal_points.max()], 
                       'k--', alpha=0.5)
        axes[0, i].set_xlabel("True Ideal Points")
        axes[0, i].set_ylabel("Estimated Ideal Points")
        axes[0, i].set_title(f"{name}\nR² = {results['r_squared']:.3f}")
        axes[0, i].grid(True, alpha=0.3)
        
        # Plot 2: Error distribution
        errors = results["theta_aligned"] - true_ideal_points
        axes[1, i].hist(errors, bins=30, alpha=0.7, color=color, density=True)
        axes[1, i].set_xlabel("Estimation Error")
        axes[1, i].set_ylabel("Density")
        axes[1, i].set_title(f"Error Distribution\nStd = {np.std(errors):.3f}")
        axes[1, i].grid(True, alpha=0.3)
        axes[1, i].axvline(0, color='k', linestyle='--', alpha=0.5)
    
    # Hide unused subplots
    for i in range(len(successful_strategies), cols):
        axes[0, i].set_visible(False)
        axes[1, i].set_visible(False)
    
    plt.tight_layout()
    plt.show()
    
    # Additional comparison plots if multiple strategies
    if len(successful_strategies) > 1:
        create_moe_metric_comparison(all_results, successful_strategies)


def create_moe_metric_comparison(all_results: Dict[str, Dict], strategies: list):
    """Create metric comparison plots for MoE strategies."""
    
    fig, axes = plt.subplots(1, 3, figsize=(18, 6))
    
    # Colors
    colors = {
        "concatenate": "blue", 
        "moe_average": "green", 
        "moe_gating": "red", 
        "moe_learned": "orange"
    }
    
    # Extract metrics
    coverages = [all_results[s]["coverage"] for s in strategies]
    correlations = [all_results[s]["correlation"] for s in strategies]
    r_squareds = [all_results[s]["r_squared"] for s in strategies]
    
    # Plot 1: Coverage comparison
    bars1 = axes[0].bar(strategies, coverages, 
                       color=[colors.get(s, 'gray') for s in strategies],
                       alpha=0.7)
    axes[0].axhline(y=95, color='black', linestyle='--', alpha=0.7, label='Target (95%)')
    axes[0].set_ylabel("Coverage (%)")
    axes[0].set_title("Coverage Comparison")
    axes[0].legend()
    axes[0].tick_params(axis='x', rotation=45)
    
    # Add value labels
    for bar, cov in zip(bars1, coverages):
        axes[0].text(bar.get_x() + bar.get_width()/2, bar.get_height() + 1,
                   f'{cov:.1f}%', ha='center', va='bottom')
    
    # Plot 2: Correlation comparison
    bars2 = axes[1].bar(strategies, correlations,
                       color=[colors.get(s, 'gray') for s in strategies],
                       alpha=0.7)
    axes[1].set_ylabel("Correlation")
    axes[1].set_title("Correlation Comparison")
    axes[1].set_ylim(0, 1)
    axes[1].tick_params(axis='x', rotation=45)
    
    # Add value labels
    for bar, corr in zip(bars2, correlations):
        axes[1].text(bar.get_x() + bar.get_width()/2, bar.get_height() + 0.02,
                   f'{corr:.3f}', ha='center', va='bottom')
    
    # Plot 3: R² comparison
    bars3 = axes[2].bar(strategies, r_squareds,
                       color=[colors.get(s, 'gray') for s in strategies],
                       alpha=0.7)
    axes[2].set_ylabel("R²")
    axes[2].set_title("R² Comparison")
    axes[2].set_ylim(0, 1)
    axes[2].tick_params(axis='x', rotation=45)
    
    # Add value labels
    for bar, r2 in zip(bars3, r_squareds):
        axes[2].text(bar.get_x() + bar.get_width()/2, bar.get_height() + 0.02,
                   f'{r2:.3f}', ha='center', va='bottom')
    
    plt.tight_layout()
    plt.show()


def plot_moe_coverage_visualization(results: Dict, true_ideal_points: np.ndarray, title: str):
    """Create detailed coverage visualization for MoE."""
    
    fig, axes = plt.subplots(1, 2, figsize=(15, 6))
    
    # Sort by true ideal points for better visualization
    sort_idx = np.argsort(true_ideal_points)
    true_sorted = true_ideal_points[sort_idx]
    
    # Align all posterior samples
    theta_samples = results["theta_samples"]
    slope, intercept = results["slope"], results["intercept"]
    aligned_samples = theta_samples * slope + intercept
    
    mean_sorted = aligned_samples.mean(axis=0)[sort_idx]
    std_sorted = aligned_samples.std(axis=0)[sort_idx]
    
    # Plot 1: Confidence intervals
    x_pos = np.arange(len(true_sorted))
    axes[0].fill_between(x_pos,
                        mean_sorted - 1.96 * std_sorted,
                        mean_sorted + 1.96 * std_sorted,
                        alpha=0.3, color='green', label="95% CI")
    axes[0].scatter(x_pos, mean_sorted, alpha=0.7, s=8, color='green', label="Posterior Mean")
    axes[0].scatter(x_pos, true_sorted, alpha=0.7, s=8, color='red', label="True Values")
    
    axes[0].set_xlabel("Politicians (sorted by true ideal point)")
    axes[0].set_ylabel("Ideal Point")
    axes[0].set_title(f"{title}: Coverage Visualization\n({results['coverage']:.1f}% coverage)")
    axes[0].legend()
    axes[0].grid(True, alpha=0.3)
    
    # Plot 2: Residuals vs fitted
    fitted_values = results["theta_aligned"]
    residuals = fitted_values - true_ideal_points
    
    axes[1].scatter(fitted_values, residuals, alpha=0.6, s=15)
    axes[1].axhline(y=0, color='red', linestyle='--', alpha=0.7)
    axes[1].set_xlabel("Fitted Values")
    axes[1].set_ylabel("Residuals")
    axes[1].set_title("Residuals vs Fitted")
    axes[1].grid(True, alpha=0.3)
    
    plt.tight_layout()
    plt.show()


def print_moe_summary_comparison(all_results: Dict[str, Dict], fusion_strategies: list):
    """Print a summary comparison table for MoE strategies."""
    
    print(f"\n{'='*75}")
    print("SUMMARY COMPARISON: MIXTURE OF EXPERTS (VOTES + SURVEYS)")
    print(f"{'='*75}")
    print(f"{'Strategy':<15} {'Coverage':<10} {'Correlation':<12} {'R²':<8} {'Comments'}")
    print("-" * 75)
    
    successful_strategies = []
    
    for strategy in fusion_strategies:
        if strategy not in all_results:
            print(f"{strategy.upper():<15} {'ERROR':<10} {'ERROR':<12} {'ERROR':<8} Failed to train")
            continue
            
        successful_strategies.append(strategy)
        results = all_results[strategy]
        coverage = results["coverage"]
        correlation = results["correlation"]
        r_squared = results["r_squared"]
        
        # Generate comments based on performance
        comments = []
        if coverage >= 93:
            comments.append("Good coverage")
        elif coverage >= 90:
            comments.append("OK coverage")
        else:
            comments.append("Low coverage")
            
        if correlation >= 0.8:
            comments.append("High correlation")
        elif correlation >= 0.6:
            comments.append("OK correlation")
        else:
            comments.append("Low correlation")
        
        comment_str = ", ".join(comments)
        
        print(f"{strategy.upper():<15} {coverage:<10.1f} {correlation:<12.3f} {r_squared:<8.3f} {comment_str}")
    
    print("-" * 75)
    
    # Analysis and recommendations
    if len(successful_strategies) >= 2:
        print("\nMoE STRATEGY ANALYSIS:")
        
        # Compare MoE strategies to baseline
        baseline_results = all_results.get("concatenate", {})
        moe_strategies = [s for s in successful_strategies if s.startswith("moe_")]
        
        if baseline_results and moe_strategies:
            baseline_cov = baseline_results["coverage"]
            baseline_corr = baseline_results["correlation"]
            baseline_r2 = baseline_results["r_squared"]
            
            print(f"• Baseline (Concatenation): Coverage={baseline_cov:.1f}%, Corr={baseline_corr:.3f}, R²={baseline_r2:.3f}")
            
            for moe_strategy in moe_strategies:
                moe_results = all_results[moe_strategy]
                moe_cov = moe_results["coverage"]
                moe_corr = moe_results["correlation"]
                moe_r2 = moe_results["r_squared"]
                
                cov_diff = moe_cov - baseline_cov
                corr_diff = moe_corr - baseline_corr
                r2_diff = moe_r2 - baseline_r2
                
                print(f"• {moe_strategy.upper()}: "
                      f"Coverage={moe_cov:.1f}% ({cov_diff:+.1f}%), "
                      f"Corr={moe_corr:.3f} ({corr_diff:+.3f}), "
                      f"R²={moe_r2:.3f} ({r2_diff:+.3f})")
        
        # Find best MoE strategy
        moe_results = {s: all_results[s] for s in moe_strategies}
        if moe_results:
            best_moe = max(moe_results.keys(), 
                          key=lambda x: moe_results[x]["coverage"] + 10 * moe_results[x]["correlation"])
            best_results = moe_results[best_moe]
            
            print(f"\n• Best MoE strategy: {best_moe.upper()}")
            print(f"  - Coverage: {best_results['coverage']:.1f}%")
            print(f"  - Correlation: {best_results['correlation']:.3f}")
            print(f"  - R²: {best_results['r_squared']:.3f}")
    
    # Overall recommendation
    if successful_strategies:
        print(f"\nOVERALL RECOMMENDATION:")
        best_strategy = max(successful_strategies, 
                           key=lambda x: all_results[x]["coverage"] + 10 * all_results[x]["correlation"])
        best_results = all_results[best_strategy]
        print(f"• Best performing strategy: {best_strategy.upper()}")
        print(f"  - Coverage: {best_results['coverage']:.1f}%")
        print(f"  - Correlation: {best_results['correlation']:.3f}")
        print(f"  - R²: {best_results['r_squared']:.3f}")
        
        # Strategy-specific insights
        if best_strategy == "moe_average":
            print("  - Simple averaging of experts works well")
        elif best_strategy == "moe_gating":
            print("  - Learned gating successfully adapts expert weights")
        elif best_strategy == "moe_learned":
            print("  - Fixed learned weights provide good balance")
        elif best_strategy == "concatenate":
            print("  - Concatenation baseline is competitive with MoE strategies")


if __name__ == "__main__":
    print("Starting Mixture of Experts test with Votes and Surveys...")
    print()
    
    # Run the test
    results, data, true_ideal_points = test_mixture_of_experts_votes_surveys()
    
    print(f"\n{'='*75}")
    print("MoE TEST COMPLETED!")
    print(f"{'='*75}")
    
    # Additional analysis comparing MoE strategies
    if results:
        successful_strategies = list(results.keys())
        moe_strategies = [s for s in successful_strategies if s.startswith("moe_")]
        
        if len(moe_strategies) > 1:
            print(f"\nMoE STRATEGY COMPARISON:")
            best_moe = max(moe_strategies, key=lambda x: results[x]["r_squared"])
            worst_moe = min(moe_strategies, key=lambda x: results[x]["r_squared"])
            
            best_r2 = results[best_moe]["r_squared"]
            worst_r2 = results[worst_moe]["r_squared"]
            improvement = ((best_r2 - worst_r2) / worst_r2) * 100
            
            print(f"• Best MoE: {best_moe.upper()} (R² = {best_r2:.3f})")
            print(f"• Worst MoE: {worst_moe.upper()} (R² = {worst_r2:.3f})")
            print(f"• Improvement: {improvement:.1f}%")
        
        # Compare best MoE vs concatenation
        if "concatenate" in results and moe_strategies:
            concat_r2 = results["concatenate"]["r_squared"]
            best_moe_r2 = max(results[s]["r_squared"] for s in moe_strategies)
            best_moe_strategy = max(moe_strategies, key=lambda x: results[x]["r_squared"])
            
            if best_moe_r2 > concat_r2:
                improvement = ((best_moe_r2 - concat_r2) / concat_r2) * 100
                print(f"\n• MoE improvement over concatenation: {improvement:.1f}%")
                print(f"• Best MoE strategy: {best_moe_strategy.upper()}")
            else:
                print(f"\n• Concatenation baseline remains competitive")
    
    print("\nMoE test script finished.")
