# Comprehensive Test Script for Multi-Modal Models
# ===================================================
# Tests different fusion strategies on simulated political data:
# - Votes only (Bernoulli likelihood)
# - Text embeddings only (Gaussian likelihood) 
# - Survey answers only (Categorical likelihood)
# - All three combined with Product of Experts (PoE)
# - All three combined with Mixture of Experts (MoE)

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

from models import MultiModalVAE, create_multimodal_vae


def create_synthetic_political_data(
    num_politicians: int = 500,
    num_bills: int = 100, 
    num_survey_questions: int = 20,
    text_embedding_dim: int = 50,
    latent_dim: int = 1,
    seed: int = 42
) -> Tuple[Dict[str, jnp.ndarray], jnp.ndarray]:
    """
    Create synthetic political data with three modalities:
    - Voting records (binary)
    - Text embeddings (continuous)
    - Survey responses (categorical)
    
    Returns:
        data: Dictionary with modality data
        true_ideal_points: Ground truth ideal points
    """
    np.random.seed(seed)
    rng = random.PRNGKey(seed)
    
    # Generate true latent ideal points
    true_ideal_points = np.random.normal(0, 1, (num_politicians, latent_dim))
    
    # =====================================================
    # 1. Generate Voting Records (Bernoulli likelihood)
    # =====================================================
    # Generate bill parameters
    bill_difficulty = np.random.normal(0, 0.5, num_bills)  # delta (difficulty)
    bill_discrimination = np.random.normal(0, 1, num_bills)  # beta (discrimination)
    
    # Generate votes based on 2PL IRT model: P(vote=1) = sigmoid(discrimination * ideal_point - difficulty)
    if latent_dim == 1:
        voting_scores = true_ideal_points[:, 0]
    else:
        voting_scores = np.sum(true_ideal_points, axis=1)  # Project to 1D
    
    vote_logits = voting_scores[:, 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 Text Embeddings (Gaussian likelihood)
    # =====================================================
    # Text embeddings are linear transformation of ideal points + noise
    text_weights = np.random.normal(0, 1, (latent_dim, text_embedding_dim))
    text_mean = true_ideal_points @ text_weights
    text_embeddings = text_mean + np.random.normal(0, 0.3, text_mean.shape)
    
    # =====================================================
    # 3. Generate Survey Responses (Categorical likelihood)  
    # =====================================================
    # Simplify: treat as single categorical with multiple categories
    max_categories = 5  # Maximum categories per question
    
    # Generate a single combined survey score based on ideal points
    if latent_dim == 1:
        survey_scores = true_ideal_points[:, 0]
    else:
        survey_scores = np.sum(true_ideal_points, axis=1)
    
    # Survey responses: Convert to proper format for categorical
    survey_boundaries = np.linspace(survey_scores.min(), survey_scores.max(), num_survey_questions + 1)
    survey_responses = np.digitize(survey_scores, survey_boundaries[1:-1])
    survey_responses = np.clip(survey_responses, 0, num_survey_questions - 1)
    
    # For categorical likelihood, we need class indices (not one-hot)
    # Ensure it's 1D array of integers
    survey_responses = survey_responses.astype(np.int32).ravel()
    
    data = {
        "votes": jnp.array(votes, dtype=jnp.float32),
        "text_embeddings": jnp.array(text_embeddings, dtype=jnp.float32), 
        "survey_responses": jnp.array(survey_responses, dtype=jnp.int32)
    }
    
    return data, jnp.array(true_ideal_points, dtype=jnp.float32)


# ============================================================================
# Model Specification Functions  
# ============================================================================

def get_votes_only_spec(num_bills: int) -> Dict[str, Dict[str, Any]]:
    """Specification for votes-only model."""
    return {
        "votes": {
            "type": "bernoulli",
            "dim": num_bills,
            "likelihood_params": {}
        }
    }

def get_text_only_spec(text_embedding_dim: int) -> Dict[str, Dict[str, Any]]:
    """Specification for text-only model.""" 
    return {
        "text_embeddings": {
            "type": "gaussian", 
            "dim": text_embedding_dim,
            "likelihood_params": {"scale": 1.0}
        }
    }

def get_survey_only_spec(num_survey_categories: int) -> Dict[str, Dict[str, Any]]:
    """Specification for survey-only model."""
    return {
        "survey_responses": {
            "type": "categorical",
            "dim": num_survey_categories,  # Number of possible response categories
            "likelihood_params": {}
        }
    }

def get_multimodal_spec(num_bills: int, text_embedding_dim: int, num_survey_categories: int) -> Dict[str, Dict[str, Any]]:
    """Specification for full multi-modal model."""
    return {
        "votes": {
            "type": "bernoulli",
            "dim": num_bills,
            "likelihood_params": {}
        },
        "text_embeddings": {
            "type": "gaussian",
            "dim": text_embedding_dim, 
            "likelihood_params": {"scale": 1.0}
        },
        "survey_responses": {
            "type": "categorical",
            "dim": num_survey_categories,
            "likelihood_params": {}
        }
    }


# ============================================================================
# Alignment and Coverage Functions (from existing code)
# ============================================================================

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

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

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

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

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

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

    aligned = theta_samples * slope + intercept

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

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

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

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


# ============================================================================
# Evaluation and Plotting Functions
# ============================================================================

def evaluate_model(model: MultiModalVAE, 
                  data: Dict[str, jnp.ndarray], 
                  true_ideal_points: jnp.ndarray,
                  model_name: str = "Model",
                  num_samples: int = 500) -> Dict[str, Any]:
    """
    Evaluate a fitted model and return metrics.
    
    Returns:
        results: Dictionary with coverage, correlation, and other metrics
    """
    print(f"\n{'='*50}")
    print(f"Evaluating {model_name}")
    print(f"{'='*50}")
    
    # Sample from posterior
    print("Sampling from posterior...")
    try:
        posterior = model.sample_posterior(data, num_samples=num_samples, seed=42)
        
        # Extract theta samples (latent ideal points)
        if "z" in posterior:
            theta_samples = np.asarray(posterior["z"])
        elif "latent" in posterior:
            theta_samples = np.asarray(posterior["latent"])
        else:
            print("Warning: No 'z' samples found in posterior")
            print(f"Available keys: {list(posterior.keys())}")
            # Try to find any latent variable
            for key in posterior.keys():
                if key not in ["obs_votes", "obs_text_embeddings", "obs_survey_responses"] and not key.startswith("obs_"):
                    theta_samples = np.asarray(posterior[key])
                    print(f"Using {key} as latent samples")
                    break
            else:
                return {"error": "No latent samples found"}
            
        # Ensure correct shape
        if theta_samples.ndim == 3 and theta_samples.shape[-1] == 1:
            theta_samples = theta_samples.squeeze(-1)  # Remove last dimension if it's 1
        
        print(f"Theta samples shape: {theta_samples.shape}")
        print(f"True ideal points shape: {true_ideal_points.shape}")
        
        # Flatten true ideal points if needed
        if true_ideal_points.ndim == 2 and true_ideal_points.shape[1] == 1:
            true_ideal_points_flat = true_ideal_points.ravel()
        else:
            true_ideal_points_flat = true_ideal_points.ravel()
        
        # Compute coverage
        cov, n_in = coverage_rate(theta_samples, true_ideal_points_flat)
        
        # Compute posterior mean and align to truth
        theta_mean = theta_samples.mean(axis=0)
        theta_mean_aligned, slope, intercept = _linear_align(theta_mean, true_ideal_points_flat)
        
        # Compute correlation
        correlation = np.corrcoef(true_ideal_points_flat, theta_mean_aligned)[0, 1]
        
        # Compute R-squared  
        fit = linregress(true_ideal_points_flat, theta_mean_aligned)
        r_squared = fit.rvalue**2
        
        results = {
            "coverage": cov,
            "n_in_ci": n_in,
            "correlation": correlation,
            "r_squared": r_squared,
            "slope": slope,
            "intercept": intercept,
            "theta_samples": theta_samples,
            "theta_mean_aligned": theta_mean_aligned,
            "true_ideal_points": true_ideal_points_flat,
            "error": None
        }
        
        print(f"Coverage: {cov:.1f}% ({n_in}/{len(true_ideal_points_flat)})")
        print(f"Correlation: {correlation:.3f}")
        print(f"R²: {r_squared:.3f}")
        
        return results
        
    except Exception as e:
        print(f"Error during evaluation: {e}")
        return {"error": str(e)}


def plot_model_results(results: Dict[str, Any], model_name: str, save_path: str = None):
    """Plot evaluation results for a single model."""
    if results.get("error"):
        print(f"Cannot plot {model_name}: {results['error']}")
        return
        
    fig, axes = plt.subplots(1, 2, figsize=(15, 6))
    
    # Plot 1: Posterior vs Truth scatter
    true_points = results["true_ideal_points"]
    aligned_mean = results["theta_mean_aligned"]
    correlation = results["correlation"]
    r_squared = results["r_squared"]
    
    axes[0].scatter(true_points, aligned_mean, alpha=0.6, s=20)
    
    # Add regression line
    x_vals = np.linspace(true_points.min(), true_points.max(), 100)
    slope = results["slope"]
    intercept = results["intercept"]
    axes[0].plot(x_vals, slope * x_vals + intercept, 'r-', label=f'Fitted (R²={r_squared:.3f})')
    axes[0].plot(x_vals, x_vals, 'k--', alpha=0.5, label='Perfect')
    
    axes[0].set_xlabel("True Ideal Points")
    axes[0].set_ylabel("Posterior Mean (Aligned)")
    axes[0].set_title(f"{model_name}: Posterior vs Truth")
    axes[0].legend()
    axes[0].grid(True, alpha=0.3)
    
    # Plot 2: Coverage visualization
    theta_samples = results["theta_samples"]
    sort_idx = np.argsort(true_points)
    x_sorted = true_points[sort_idx]
    
    # Align all samples
    aligned_samples = theta_samples * slope + intercept
    mean_sorted = aligned_samples.mean(axis=0)[sort_idx]
    std_sorted = aligned_samples.std(axis=0)[sort_idx]
    
    coverage = results["coverage"]
    
    axes[1].fill_between(range(len(x_sorted)),
                        mean_sorted - 1.96 * std_sorted,
                        mean_sorted + 1.96 * std_sorted,
                        alpha=0.3, label="95% CI")
    axes[1].scatter(range(len(x_sorted)), mean_sorted, alpha=0.7, s=15, label="Posterior Mean")
    axes[1].scatter(range(len(x_sorted)), x_sorted, alpha=0.7, s=15, color="red", label="True Values")
    
    axes[1].set_xlabel("Politician (sorted by true ideal point)")
    axes[1].set_ylabel("Ideal Point")
    axes[1].set_title(f"{model_name}: Coverage Visualization ({coverage:.1f}%)")
    axes[1].legend()
    axes[1].grid(True, alpha=0.3)
    
    plt.tight_layout()
    
    if save_path:
        plt.savefig(save_path, dpi=300, bbox_inches='tight')
    
    plt.show()


def plot_model_comparison(all_results: Dict[str, Dict[str, Any]], save_path: str = None):
    """Plot comparison across all models."""
    # Filter out models with errors
    valid_results = {name: res for name, res in all_results.items() if not res.get("error")}
    
    if not valid_results:
        print("No valid results to plot")
        return
        
    # Extract metrics
    model_names = list(valid_results.keys())
    coverages = [valid_results[name]["coverage"] for name in model_names]
    correlations = [valid_results[name]["correlation"] for name in model_names]
    r_squareds = [valid_results[name]["r_squared"] for name in model_names]
    
    fig, axes = plt.subplots(1, 3, figsize=(18, 5))
    
    # Coverage comparison
    bars1 = axes[0].bar(model_names, coverages, color='skyblue', alpha=0.7)
    axes[0].axhline(y=95, color='red', linestyle='--', alpha=0.7, label='Target (95%)')
    axes[0].set_ylabel("Coverage (%)")
    axes[0].set_title("Coverage Rate Comparison")
    axes[0].legend()
    axes[0].tick_params(axis='x', rotation=45)
    
    # Add value labels on bars
    for bar, coverage in zip(bars1, coverages):
        axes[0].text(bar.get_x() + bar.get_width()/2, bar.get_height() + 1, 
                    f'{coverage:.1f}%', ha='center', va='bottom', fontsize=10)
    
    # Correlation comparison  
    bars2 = axes[1].bar(model_names, correlations, color='lightgreen', alpha=0.7)
    axes[1].set_ylabel("Correlation")
    axes[1].set_title("Correlation with Truth")
    axes[1].tick_params(axis='x', rotation=45)
    axes[1].set_ylim(0, 1)
    
    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', fontsize=10)
    
    # R-squared comparison
    bars3 = axes[2].bar(model_names, r_squareds, color='salmon', alpha=0.7) 
    axes[2].set_ylabel("R²")
    axes[2].set_title("R² with Truth")
    axes[2].tick_params(axis='x', rotation=45)
    axes[2].set_ylim(0, 1)
    
    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', fontsize=10)
    
    plt.tight_layout()
    
    if save_path:
        plt.savefig(save_path, dpi=300, bbox_inches='tight')
    
    plt.show()


# ============================================================================
# Main Test Function
# ============================================================================

def run_comprehensive_test(
    num_politicians: int = 1000,
    num_bills: int = 1000,
    num_survey_questions: int = 50,
    text_embedding_dim: int = 300,
    latent_dim: int = 1,
    num_steps: int = 5000,
    num_samples: int = 100,
    seed: int = 42
):
    """
    Run comprehensive test of all model configurations.
    """
    print("="*80)
    print("COMPREHENSIVE MULTI-MODAL MODEL TEST")
    print("="*80)
    print(f"Parameters:")
    print(f"  Politicians: {num_politicians}")
    print(f"  Bills: {num_bills}")  
    print(f"  Survey Questions: {num_survey_questions}")
    print(f"  Text Embedding Dim: {text_embedding_dim}")
    print(f"  Latent Dim: {latent_dim}")
    print(f"  Training Steps: {num_steps}")
    print(f"  Posterior Samples: {num_samples}")
    print()
    
    # Generate synthetic data
    print("Generating synthetic political data...")
    data, true_ideal_points = create_synthetic_political_data(
        num_politicians=num_politicians,
        num_bills=num_bills,
        num_survey_questions=num_survey_questions,
        text_embedding_dim=text_embedding_dim,
        latent_dim=latent_dim,
        seed=seed
    )
    
    print(f"Data shapes:")
    for modality, values in data.items():
        print(f"  {modality}: {values.shape}")
    print(f"True ideal points: {true_ideal_points.shape}")
    print()
    
    # Define model configurations
    models_to_test = [
        ("Votes Only", get_votes_only_spec(num_bills), {"votes": data["votes"]}, "concatenate"),
        ("Text Only", get_text_only_spec(text_embedding_dim), {"text_embeddings": data["text_embeddings"]}, "concatenate"),
        ("Survey Only", get_survey_only_spec(num_survey_questions), {"survey_responses": data["survey_responses"]}, "concatenate"),
        ("All + PoE", get_multimodal_spec(num_bills, text_embedding_dim, num_survey_questions), data, "poe"),
        ("All + MoE", get_multimodal_spec(num_bills, text_embedding_dim, num_survey_questions), data, "moe_average"),
    ]
    
    # Train and evaluate each model
    all_results = {}
    
    for model_name, spec, model_data, fusion_strategy in models_to_test:
        print(f"\n{'='*60}")
        print(f"TRAINING: {model_name}")
        print(f"Fusion Strategy: {fusion_strategy}")
        print(f"{'='*60}")
        
        try:
            # Create model
            model = create_multimodal_vae(
                modality_specs=spec,
                fusion_strategy=fusion_strategy,
                latent_dim=latent_dim,
                encoder_hidden_dims=(64, 32),
                decoder_hidden_dims=(32,),
                activation="relu"
            )
            
            # Train model
            print(f"Training {model_name}...")
            model.fit(
                data=model_data,
                num_steps=num_steps,
                lr=1e-3,
                seed=seed,
                progress_bar=False
            )
            
            # Evaluate model
            results = evaluate_model(
                model=model,
                data=model_data,
                true_ideal_points=true_ideal_points,
                model_name=model_name,
                num_samples=num_samples
            )
            
            all_results[model_name] = results
            
            # Plot individual results
            plot_model_results(results, model_name)
            
        except Exception as e:
            print(f"Error training {model_name}: {e}")
            all_results[model_name] = {"error": str(e)}
    
    # Plot comparison across all models
    print(f"\n{'='*60}")
    print("OVERALL COMPARISON")
    print(f"{'='*60}")
    
    plot_model_comparison(all_results)
    
    # Print summary table
    print("\nSUMMARY TABLE:")
    print("-" * 80)
    print(f"{'Model':<15} {'Coverage':<12} {'Correlation':<12} {'R²':<10} {'Status'}")
    print("-" * 80)
    
    for model_name, results in all_results.items():
        if results.get("error"):
            print(f"{model_name:<15} {'ERROR':<12} {'ERROR':<12} {'ERROR':<10} {results['error'][:30]}")
        else:
            coverage = results["coverage"]
            correlation = results["correlation"] 
            r_squared = results["r_squared"]
            print(f"{model_name:<15} {coverage:<12.1f} {correlation:<12.3f} {r_squared:<10.3f} OK")
    
    print("-" * 80)
    
    return all_results, data, true_ideal_points


# ============================================================================
# Entry Point
# ============================================================================

if __name__ == "__main__":
    # Run the comprehensive test
    results, data, true_ideal_points = run_comprehensive_test(
        num_politicians=1000,
        num_bills=1000,
        num_survey_questions=50, 
        text_embedding_dim=300,
        latent_dim=1,
        num_steps=5000,
        num_samples=300,
        seed=42
    )
    
    print("\nTest completed successfully!")

