# Product of Experts Test Script: Votes + Surveys
# ===============================================
# Test script for Product of Experts fusion strategy 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")

from models import create_multimodal_vae, coverage_rate, _linear_align


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_product_of_experts_votes_surveys():
    """
    Test Product of Experts fusion with votes and surveys.
    """
    print("Product 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 fusion strategies for comparison
    fusion_strategies = ["concatenate", "poe"]
    all_results = {}
    
    for fusion_strategy in fusion_strategies:
        print(f"\n{'-'*40}")
        print(f"Testing fusion strategy: {fusion_strategy.upper()}")
        print(f"{'-'*40}")
        
        # Create model
        print("Creating model...")
        model = create_multimodal_vae(
            modality_specs=modality_specs,
            fusion_strategy=fusion_strategy,
            latent_dim=1,
            encoder_hidden_dims=(64, 32),
            decoder_hidden_dims=(32,),
            activation="relu",
            dropout_rate=0.1
        )
        
        # Train model
        print("Training model...")
        model.fit(
            data=model_data,
            num_steps=6000,
            lr=1e-3,
            seed=42,
            progress_bar=False
        )
        
        # Sample from posterior
        print("Sampling 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)
        
        print(f"Theta samples shape: {theta_samples.shape}")
        
        # 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(f"Results for {fusion_strategy.upper()}:")
        print(f"  Correlation: {correlation:.3f}")
        print(f"  R²: {r_squared:.3f}")
        print(f"  Coverage: {cov:.1f}% ({n_in}/{len(true_ideal_points)})")
        print(f"  Slope: {slope:.3f}, Intercept: {intercept:.3f}")
    
    # Create visualization
    print(f"\n{'-'*40}")
    print("Creating visualizations...")
    print(f"{'-'*40}")
    
    create_comparison_plots(all_results, true_ideal_points, fusion_strategies)
    
    # Print summary comparison
    print_summary_comparison(all_results, fusion_strategies)
    
    return all_results, data, true_ideal_points


def create_comparison_plots(all_results: Dict[str, Dict], 
                          true_ideal_points: np.ndarray, 
                          fusion_strategies: list):
    """Create comparison plots for different fusion strategies."""
    
    fig, axes = plt.subplots(2, 3, figsize=(18, 12))
    
    # Color scheme
    colors = {"concatenate": "blue", "poe": "red"}
    strategy_names = {"concatenate": "Concatenation", "poe": "Product of Experts"}
    
    for i, strategy in enumerate(fusion_strategies):
        if strategy not in all_results:
            continue
            
        results = all_results[strategy]
        color = colors[strategy]
        name = strategy_names[strategy]
        
        # 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)
    
    # Plot 3: Coverage comparison
    if len(all_results) >= 2:
        strategies = list(all_results.keys())
        coverages = [all_results[s]["coverage"] for s in strategies]
        correlations = [all_results[s]["correlation"] for s in strategies]
        
        # Coverage bar plot
        bars = axes[0, 2].bar(strategies, coverages, 
                             color=[colors.get(s, 'gray') for s in strategies],
                             alpha=0.7)
        axes[0, 2].axhline(y=95, color='black', linestyle='--', alpha=0.7, label='Target (95%)')
        axes[0, 2].set_ylabel("Coverage (%)")
        axes[0, 2].set_title("Coverage Comparison")
        axes[0, 2].legend()
        
        # Add value labels
        for bar, cov in zip(bars, coverages):
            axes[0, 2].text(bar.get_x() + bar.get_width()/2, bar.get_height() + 1,
                           f'{cov:.1f}%', ha='center', va='bottom')
        
        # Correlation bar plot
        bars2 = axes[1, 2].bar(strategies, correlations,
                              color=[colors.get(s, 'gray') for s in strategies],
                              alpha=0.7)
        axes[1, 2].set_ylabel("Correlation")
        axes[1, 2].set_title("Correlation Comparison")
        axes[1, 2].set_ylim(0, 1)
        
        # Add value labels
        for bar, corr in zip(bars2, correlations):
            axes[1, 2].text(bar.get_x() + bar.get_width()/2, bar.get_height() + 0.02,
                           f'{corr:.3f}', ha='center', va='bottom')
    
    plt.tight_layout()
    plt.show()
    
    # Additional plot: Coverage visualization for PoE
    if "poe" in all_results:
        plot_coverage_visualization(all_results["poe"], true_ideal_points, "Product of Experts")


def plot_coverage_visualization(results: Dict, true_ideal_points: np.ndarray, title: str):
    """Create detailed coverage visualization."""
    
    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='blue', label="95% CI")
    axes[0].scatter(x_pos, mean_sorted, alpha=0.7, s=8, color='blue', 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_summary_comparison(all_results: Dict[str, Dict], fusion_strategies: list):
    """Print a summary comparison table."""
    
    print(f"\n{'='*60}")
    print("SUMMARY COMPARISON: VOTES + SURVEYS")
    print(f"{'='*60}")
    print(f"{'Strategy':<20} {'Coverage':<10} {'Correlation':<12} {'R²':<8} {'Comments'}")
    print("-" * 70)
    
    for strategy in fusion_strategies:
        if strategy not in all_results:
            print(f"{strategy.upper():<20} {'ERROR':<10} {'ERROR':<12} {'ERROR':<8} Failed to train")
            continue
            
        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():<20} {coverage:<10.1f} {correlation:<12.3f} {r_squared:<8.3f} {comment_str}")
    
    print("-" * 70)
    
    # Analysis and recommendations
    if len(all_results) >= 2:
        print("\nANALYSIS:")
        
        poe_results = all_results.get("poe", {})
        concat_results = all_results.get("concatenate", {})
        
        if poe_results and concat_results:
            poe_coverage = poe_results["coverage"]
            concat_coverage = concat_results["coverage"]
            poe_corr = poe_results["correlation"]
            concat_corr = concat_results["correlation"]
            
            print(f"• Coverage difference (PoE - Concat): {poe_coverage - concat_coverage:+.1f}%")
            print(f"• Correlation difference (PoE - Concat): {poe_corr - concat_corr:+.3f}")
            
            if poe_coverage > concat_coverage and poe_corr > concat_corr:
                print("• Product of Experts shows superior performance on both metrics")
            elif poe_coverage > concat_coverage:
                print("• Product of Experts shows better coverage but similar correlation")
            elif poe_corr > concat_corr:
                print("• Product of Experts shows better correlation but similar coverage")
            else:
                print("• Concatenation strategy performs competitively with PoE")
        
        print(f"\nRECOMMENDATION:")
        best_strategy = max(all_results.keys(), 
                           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}")


if __name__ == "__main__":
    print("Starting Product of Experts test with Votes and Surveys...")
    print()
    
    # Run the test
    results, data, true_ideal_points = test_product_of_experts_votes_surveys()
    
    print(f"\n{'='*60}")
    print("TEST COMPLETED SUCCESSFULLY!")
    print(f"{'='*60}")
    
    # Additional analysis
    if "poe" in results and "concatenate" in results:
        poe_perf = results["poe"]["r_squared"]
        concat_perf = results["concatenate"]["r_squared"]
        improvement = ((poe_perf - concat_perf) / concat_perf) * 100
        
        print(f"\nPRODUCT OF EXPERTS IMPROVEMENT:")
        print(f"• R² improvement over concatenation: {improvement:+.1f}%")
        
        if improvement > 5:
            print("• Product of Experts shows meaningful improvement")
        elif improvement > 0:
            print("• Product of Experts shows modest improvement")
        else:
            print("• Concatenation strategy is competitive")
    
    print("\nTest script finished.")
