# Comprehensive Comparison Test Script
# ====================================
# Test script that compares votes only, survey only, PoE, and MoE average methods
# in a 2x2 panel showing politicians sorted by ideal points with confidence intervals.
# Displays correlation with ground truth and coverage rate for each method.

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, List
import warnings
warnings.filterwarnings("ignore")


def create_votes_data(num_politicians: int, num_bills: int, true_ideal_points: np.ndarray, seed: int = 42):
    """Create voting data using 2PL IRT model."""
    np.random.seed(seed)
    
    # Generate bill parameters
    bill_difficulty = np.random.normal(0, 0.8, num_bills)
    bill_discrimination = np.random.normal(0, 1.2, num_bills)
    
    # Generate votes: P(vote=1) = sigmoid(ideal_point * discrimination - 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(float)
    
    return jnp.array(votes, dtype=jnp.float32)


def create_survey_data(num_politicians: int, num_questions: int, true_ideal_points: np.ndarray, seed: int = 42):
    """Create survey data with ordinal responses."""
    np.random.seed(seed + 100)  # Different seed for surveys
    
    survey_responses = []
    categories_per_question = []
    
    for q in range(num_questions):
        # Each question has 3-7 response categories
        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
    
    return jnp.array(survey_responses, dtype=jnp.int32), categories_per_question


def create_comprehensive_data(num_politicians: int = 300, num_bills: int = 100, 
                             num_questions: int = 30, seed: int = 42):
    """Create comprehensive dataset with both votes and surveys."""
    np.random.seed(seed)
    
    # Generate true ideal points
    true_ideal_points = np.random.normal(0, 1, num_politicians)
    
    # Create votes and surveys
    votes = create_votes_data(num_politicians, num_bills, true_ideal_points, seed)
    surveys, categories_per_question = create_survey_data(num_politicians, num_questions, true_ideal_points, seed)
    
    data = {
        "votes": votes,
        "surveys": surveys,
        "categories_per_question": categories_per_question
    }
    
    return data, jnp.array(true_ideal_points, dtype=jnp.float32)


def test_method(method_name: str, data: Dict, true_ideal_points: np.ndarray, 
                num_steps: int = 5000, seed: int = 42) -> Dict[str, Any]:
    """Test a specific method and return results."""
    print(f"Testing {method_name}...")
    
    if method_name == "votes_only":
        model_data = {"votes": data["votes"]}
        modality_specs = {
            "votes": {
                "type": "bernoulli",
                "dim": data["votes"].shape[1],
                "likelihood_params": {}
            }
        }
        fusion_strategy = "concatenate"
        
    elif method_name == "survey_only":
        total_logits = sum(data["categories_per_question"])
        model_data = {"surveys": data["surveys"]}
        modality_specs = {
            "surveys": {
                "type": "categorical_multi",
                "dim": total_logits,
                "likelihood_params": {
                    "categories_per_question": data["categories_per_question"]
                }
            }
        }
        fusion_strategy = "concatenate"
        
    elif method_name == "poe":
        total_logits = sum(data["categories_per_question"])
        model_data = {"votes": data["votes"], "surveys": data["surveys"]}
        modality_specs = {
            "votes": {
                "type": "bernoulli",
                "dim": data["votes"].shape[1],
                "likelihood_params": {}
            },
            "surveys": {
                "type": "categorical_multi",
                "dim": total_logits,
                "likelihood_params": {
                    "categories_per_question": data["categories_per_question"]
                }
            }
        }
        fusion_strategy = "poe"
        
    elif method_name == "moe_average":
        total_logits = sum(data["categories_per_question"])
        model_data = {"votes": data["votes"], "surveys": data["surveys"]}
        modality_specs = {
            "votes": {
                "type": "bernoulli",
                "dim": data["votes"].shape[1],
                "likelihood_params": {}
            },
            "surveys": {
                "type": "categorical_multi",
                "dim": total_logits,
                "likelihood_params": {
                    "categories_per_question": data["categories_per_question"]
                }
            }
        }
        fusion_strategy = "moe_average"
    
    else:
        raise ValueError(f"Unknown method: {method_name}")
    
    # Create and train model
    model = create_multimodal_vae(
        modality_specs=modality_specs,
        fusion_strategy=fusion_strategy,
        latent_dim=1,
        encoder_hidden_dims=(32, 16),
        decoder_hidden_dims=(),
        activation="relu"
    )
    
    model.fit(model_data, num_steps=num_steps, lr=1e-3, seed=seed, progress_bar=False)
    
    # Sample from posterior
    posterior = model.sample_posterior(model_data, num_samples=1000, seed=seed)
    
    if "z" not in posterior:
        print(f"Warning: No 'z' samples found for {method_name}")
        return None
    
    # Extract and process latent samples
    theta_samples = np.asarray(posterior["z"])
    print(f"Original theta samples shape: {theta_samples.shape}")
    
    # Handle different dimensionalities
    if theta_samples.ndim == 3 and theta_samples.shape[-1] == 1:
        theta_samples = theta_samples.squeeze(-1)
    elif theta_samples.ndim == 3 and theta_samples.shape[-1] > 1:
        # For multidimensional latent space, use first dimension
        theta_samples = theta_samples[:, :, 0]
    elif theta_samples.ndim == 2:
        # Already in correct shape (samples, politicians)
        pass
    else:
        print(f"Unexpected theta_samples shape: {theta_samples.shape}")
        return None
    
    print(f"Final 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)
    
    return {
        "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
    }


def plot_comprehensive_comparison(all_results: Dict[str, Dict], true_ideal_points: np.ndarray):
    """Create a 2x2 panel plot comparing all methods."""
    
    # Define methods and their display names
    methods = ["votes_only", "survey_only", "poe", "moe_average"]
    method_names = {
        "votes_only": "Votes Only",
        "survey_only": "Survey Only", 
        "poe": "Votes + Survey (PoE)",
        "moe_average": "Votes + Survey (MoE Average)"
    }
    
    # Filter successful methods
    successful_methods = [m for m in methods if m in all_results and all_results[m] is not None]
    
    if len(successful_methods) == 0:
        print("No successful methods to plot!")
        return
    
    # Create 2x2 subplot grid
    fig, axes = plt.subplots(2, 2, figsize=(16, 12))
    axes = axes.flatten()
    
    # Colors for each method - using black and grey
    colors = {
        "votes_only": "black",
        "survey_only": "grey", 
        "poe": "black",
        "moe_average": "grey"
    }
    
    for i, method in enumerate(successful_methods[:4]):  # Show up to 4 methods
        results = all_results[method]
        ax = axes[i]
        
        # Sort politicians by true ideal points for visualization
        sort_idx = np.argsort(true_ideal_points)
        true_sorted = true_ideal_points[sort_idx]
        
        # Align all posterior samples and sort
        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]
        
        # Create x-axis positions
        x_pos = np.arange(len(true_sorted))
        
        # Plot confidence intervals
        ax.fill_between(x_pos,
                       mean_sorted - 1.96 * std_sorted,
                       mean_sorted + 1.96 * std_sorted,
                       alpha=0.3, color=colors[method], label="95% CI")
        
        # Plot posterior means
        ax.scatter(x_pos, mean_sorted, alpha=0.6, s=8, color=colors[method], 
                  label="Posterior Mean", zorder=3)
        
        # Plot true values
        ax.scatter(x_pos, true_sorted, alpha=0.7, s=8, color='black', 
                  label="True Values", zorder=4)
        
        # Add perfect correlation line for reference (invisible, just for alignment)
        ax.plot([], [], 'k--', alpha=0, label="")  # Invisible line for consistent legend
        
        # Formatting
        ax.set_xlabel("Politicians (sorted by true ideal point)", fontsize=14)
        ax.set_ylabel("Ideal Point", fontsize=14)
        ax.set_title(f"{method_names[method]}\n"
                    f"Correlation: {results['correlation']:.3f} | "
                    f"Coverage: {results['coverage']:.1f}%",
                    fontsize=14, fontweight='bold')
        ax.grid(True, alpha=0.3)
        ax.legend(loc='upper left', fontsize=12)
        ax.tick_params(axis='both', which='major', labelsize=12)
        ax.tick_params(axis='both', which='major', labelsize=12)
        
        # Add text box with additional metrics - REMOVED as requested
    
    # Hide unused subplots
    for i in range(len(successful_methods), 4):
        axes[i].set_visible(False)
    
    plt.tight_layout()
    plt.suptitle("Comprehensive Comparison: Political Ideal Point Estimation Methods", 
                 fontsize=18, fontweight='bold', y=1.02)
    plt.show()


def create_summary_metrics_plot(all_results: Dict[str, Dict]):
    """Create a summary plot of key metrics across methods."""
    
    methods = ["votes_only", "survey_only", "poe", "moe_average"]
    method_names = {
        "votes_only": "Votes Only",
        "survey_only": "Survey Only", 
        "poe": "PoE",
        "moe_average": "MoE Avg"
    }
    
    # Filter successful methods
    successful_methods = [m for m in methods if m in all_results and all_results[m] is not None]
    
    if len(successful_methods) == 0:
        return
    
    # Extract metrics
    correlations = [all_results[m]["correlation"] for m in successful_methods]
    coverages = [all_results[m]["coverage"] for m in successful_methods]
    method_labels = [method_names[m] for m in successful_methods]
    
    # Colors - using black and grey
    colors = ["black", "grey"] * (len(successful_methods) // 2 + 1)
    colors = colors[:len(successful_methods)]
    
    # Create figure with 2 subplots (removed R² plot)
    fig, axes = plt.subplots(1, 2, figsize=(12, 5))
    
    # Plot 1: Correlation
    bars1 = axes[0].bar(method_labels, correlations, color=colors, alpha=0.7)
    axes[0].set_ylabel("Correlation with Ground Truth", fontsize=14)
    axes[0].set_title("Correlation Comparison", fontsize=14)
    axes[0].set_ylim(0, 1)
    axes[0].tick_params(axis='both', which='major', labelsize=12)
    for bar, corr in zip(bars1, correlations):
        axes[0].text(bar.get_x() + bar.get_width()/2, bar.get_height() + 0.02,
                    f'{corr:.3f}', ha='center', va='bottom', fontweight='bold', fontsize=12)
    
    # Plot 2: Coverage
    bars2 = axes[1].bar(method_labels, coverages, color=colors, alpha=0.7)
    axes[1].axhline(y=95, color='black', linestyle='--', alpha=0.7, label='Target (95%)')
    axes[1].set_ylabel("Coverage Rate (%)", fontsize=14)
    axes[1].set_title("Coverage Comparison", fontsize=14)
    axes[1].legend(fontsize=12)
    axes[1].tick_params(axis='both', which='major', labelsize=12)
    for bar, cov in zip(bars2, coverages):
        axes[1].text(bar.get_x() + bar.get_width()/2, bar.get_height() + 1,
                    f'{cov:.1f}%', ha='center', va='bottom', fontweight='bold', fontsize=12)
    
    # Rotate x-axis labels if needed
    for ax in axes:
        ax.tick_params(axis='x', rotation=45, labelsize=12)
        ax.grid(True, alpha=0.3)
    
    plt.tight_layout()
    plt.show()


def print_comprehensive_summary(all_results: Dict[str, Dict], true_ideal_points: np.ndarray):
    """Print comprehensive summary table."""
    
    methods = ["votes_only", "survey_only", "poe", "moe_average"]
    method_names = {
        "votes_only": "Votes Only",
        "survey_only": "Survey Only", 
        "poe": "Votes + Survey (PoE)",
        "moe_average": "Votes + Survey (MoE Average)"
    }
    
    print(f"\n{'='*90}")
    print("COMPREHENSIVE COMPARISON: POLITICAL IDEAL POINT ESTIMATION")
    print(f"{'='*90}")
    print(f"{'Method':<25} {'Correlation':<12} {'Coverage':<10} {'CI Width':<10} {'Performance'}")
    print("-" * 75)
    
    successful_methods = []
    
    for method in methods:
        if method in all_results and all_results[method] is not None:
            successful_methods.append(method)
            results = all_results[method]
            
            # Calculate average CI width
            theta_samples = results["theta_samples"]
            slope, intercept = results["slope"], results["intercept"]
            aligned_samples = theta_samples * slope + intercept
            ci_width = np.mean(2 * 1.96 * aligned_samples.std(axis=0))
            
            correlation = results["correlation"]
            coverage = results["coverage"]
            
            # Performance assessment
            performance = ""
            if correlation >= 0.8 and coverage >= 93:
                performance = "Excellent"
            elif correlation >= 0.7 and coverage >= 90:
                performance = "Good"
            elif correlation >= 0.6 and coverage >= 87:
                performance = "Fair"
            else:
                performance = "Poor"
            
            print(f"{method_names[method]:<25} {correlation:<12.3f} "
                  f"{coverage:<10.1f} {ci_width:<10.2f} {performance}")
        else:
            print(f"{method_names[method]:<25} {'FAILED':<12} {'FAILED':<10} {'FAILED':<10} Failed")
    
    print("-" * 75)
    
    # Additional analysis
    if len(successful_methods) >= 2:
        print("\nKEY FINDINGS:")
        
        # Find best method overall
        best_method = max(successful_methods, 
                         key=lambda x: all_results[x]["correlation"] * 0.7 + all_results[x]["coverage"] * 0.003)
        best_results = all_results[best_method]
        
        print(f"• Best overall method: {method_names[best_method]}")
        print(f"  - Correlation: {best_results['correlation']:.3f}")
        print(f"  - Coverage: {best_results['coverage']:.1f}%")
        
        # Compare multimodal vs unimodal
        unimodal_methods = [m for m in successful_methods if m in ["votes_only", "survey_only"]]
        multimodal_methods = [m for m in successful_methods if m in ["poe", "moe_average"]]
        
        if unimodal_methods and multimodal_methods:
            best_unimodal = max(unimodal_methods, key=lambda x: all_results[x]["correlation"])
            best_multimodal = max(multimodal_methods, key=lambda x: all_results[x]["correlation"])
            
            uni_corr = all_results[best_unimodal]["correlation"]
            multi_corr = all_results[best_multimodal]["correlation"]
            improvement = ((multi_corr - uni_corr) / uni_corr) * 100
            
            print(f"\n• Multimodal vs Unimodal:")
            print(f"  - Best unimodal: {method_names[best_unimodal]} (r = {uni_corr:.3f})")
            print(f"  - Best multimodal: {method_names[best_multimodal]} (r = {multi_corr:.3f})")
            print(f"  - Improvement: {improvement:+.1f}%")
        
        # Compare fusion strategies
        if "poe" in successful_methods and "moe_average" in successful_methods:
            poe_corr = all_results["poe"]["correlation"]
            moe_corr = all_results["moe_average"]["correlation"]
            poe_cov = all_results["poe"]["coverage"]
            moe_cov = all_results["moe_average"]["coverage"]
            
            print(f"\n• Fusion Strategy Comparison:")
            print(f"  - PoE: r = {poe_corr:.3f}, coverage = {poe_cov:.1f}%")
            print(f"  - MoE Average: r = {moe_corr:.3f}, coverage = {moe_cov:.1f}%")
            
            if poe_corr > moe_corr and poe_cov > moe_cov:
                print("  - PoE dominates on both metrics")
            elif moe_corr > poe_corr and moe_cov > poe_cov:
                print("  - MoE Average dominates on both metrics")
            else:
                print("  - Trade-off between correlation and coverage")
    
    print(f"Test completed successfully!")


def run_comprehensive_test():
    """Run comprehensive comparison test."""
    print("Starting Comprehensive Comparison Test")
    print("=" * 50)
    
    # Generate data
    print("Generating comprehensive dataset...")
    data, true_ideal_points = create_comprehensive_data(
        num_politicians=1000,
        num_bills=500,
        num_questions=100,
        seed=42
    )
    
    print(f"Dataset created:")
    print(f"  Politicians: {len(true_ideal_points)}")
    print(f"  Votes: {data['votes'].shape}")
    print(f"  Surveys: {data['surveys'].shape}")
    print(f"  Survey categories: {len(data['categories_per_question'])} questions")
    
    # Test all methods
    methods = ["votes_only", "survey_only", "poe", "moe_average"]
    all_results = {}
    
    for method in methods:
        print(f"\n{'-'*30}")
        try:
            result = test_method(method, data, true_ideal_points, num_steps=5000, seed=42)
            all_results[method] = result
            if result:
                print(f"{method}: r = {result['correlation']:.3f}, coverage = {result['coverage']:.1f}%")
            else:
                print(f"{method}: FAILED")
        except Exception as e:
            print(f"{method}: ERROR - {str(e)}")
            all_results[method] = None
    
    # Create visualizations
    print(f"\n{'-'*30}")
    print("Creating visualizations...")
    
    # Main comparison plot
    plot_comprehensive_comparison(all_results, true_ideal_points)
    
    # Summary metrics plot
    create_summary_metrics_plot(all_results)
    
    # Print summary
    print_comprehensive_summary(all_results, true_ideal_points)
    
    return all_results, data, true_ideal_points


if __name__ == "__main__":
    print("Political Ideal Point Estimation: Comprehensive Comparison")
    print("=" * 60)
    print("This script compares four methods:")
    print("1. Votes Only")
    print("2. Survey Only") 
    print("3. Votes + Survey (Product of Experts)")
    print("4. Votes + Survey (Mixture of Experts Average)")
    print()
    
    # Run the comprehensive test
    results, data, true_ideal_points = run_comprehensive_test()
    
    print(f"\n{'='*60}")
    print("COMPREHENSIVE TEST COMPLETED!")
    print(f"{'='*60}")
    
    # Final recommendations
    successful_methods = [m for m in results.keys() if results[m] is not None]
    
    if successful_methods:
        print(f"\nSUCCESSFUL METHODS: {len(successful_methods)}/4")
        
        # Best method by different criteria
        if len(successful_methods) > 1:
            best_corr_method = max(successful_methods, key=lambda x: results[x]["correlation"])
            best_cov_method = max(successful_methods, key=lambda x: results[x]["coverage"])
            best_overall = max(successful_methods, 
                              key=lambda x: results[x]["correlation"] * 0.7 + results[x]["coverage"] * 0.003)
            
            print(f"• Highest correlation: {best_corr_method} ({results[best_corr_method]['correlation']:.3f})")
            print(f"• Best coverage: {best_cov_method} ({results[best_cov_method]['coverage']:.1f}%)")
            print(f"• Best overall: {best_overall}")
        
        print(f"\nRecommendation: Use the method that best balances correlation and coverage for your specific application.")
    else:
        print("No methods completed successfully. Check data generation and model parameters.")
    
    print("\nComprehensive test script finished.")
