# Simple Survey-Only Test Script
# ==============================
# A focused test script for survey data following the pattern of test_votes_only.py

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

from models import create_multimodal_vae, coverage_rate, _linear_align


def create_survey_data_simple(num_politicians: int = 1000, num_questions: int = 100, seed: int = 42):
    """Create survey data with 50 questions, each having 2-6 possible responses."""
    np.random.seed(seed)
    
    # Generate true ideal points
    true_ideal_points = np.random.normal(0, 1, num_politicians)
    
    # Store question metadata for later use
    question_info = []
    
    # Generate survey responses (categorical with varying number of choices)
    surveys = []
    
    for q in range(num_questions):
        # Each question has 2-6 response categories
        K = np.random.randint(2, 7)  # 2 to 6 inclusive
        question_info.append(K)
        
        # Generate question parameters - how ideal points map to responses
        difficulty = np.random.normal(0, 0.8)  # Question difficulty
        discrimination = np.random.normal(0, 1.2)  # How well question discriminates
        
        # Create response probabilities based on ideal points
        # Use ordinal model where higher ideal points prefer higher categories
        logits = (true_ideal_points * discrimination - difficulty)[:, None]
        
        # Create category thresholds (K-1 thresholds for K categories)
        thresholds = np.sort(np.random.normal(0, 0.5, K-1))
        
        # Compute probabilities for each category using cumulative logistic
        probs = np.zeros((num_politicians, K))
        for i in range(num_politicians):
            # Compute cumulative probabilities
            cum_probs = 1 / (1 + np.exp(thresholds - logits[i, 0]))
            
            # Convert to category probabilities
            for k in range(K):
                if k == 0:
                    probs[i, k] = cum_probs[0]
                elif k == K-1:
                    probs[i, k] = 1 - cum_probs[k-1]
                else:
                    probs[i, k] = cum_probs[k] - cum_probs[k-1]
            
            # Ensure probabilities are valid
            probs[i] = np.maximum(probs[i], 1e-8)
            probs[i] /= probs[i].sum()
        
        # Sample responses
        responses = [np.random.choice(K, p=probs[i]) for i in range(num_politicians)]
        surveys.append(responses)
    
    surveys = np.array(surveys).T  # Shape: (num_politicians, num_questions)
    
    # Return the raw categorical responses and metadata
    data = {
        "raw_responses": surveys,
        "question_info": question_info
    }
    
    return data, jnp.array(true_ideal_points, dtype=jnp.float32)


def test_surveys_only():
    """Test surveys-only model using proper categorical_multi likelihood."""
    print("Testing Surveys-Only Model with Categorical Multi Likelihood")
    print("=" * 60)
    
    # Create data
    data, true_ideal_points = create_survey_data_simple(num_politicians=300, num_questions=50)
    
    # Use raw responses for categorical model
    raw_responses = data["raw_responses"]  # Shape: (num_politicians, num_questions)
    question_info = data["question_info"]  # List of number of categories per question
    
    print(f"Raw survey data shape: {raw_responses.shape}")
    print(f"True ideal points shape: {true_ideal_points.shape}")
    print(f"Questions with categories: {question_info[:10]}...")  # Show first 10
    print(f"Response range: {raw_responses.min()} to {raw_responses.max()}")
    
    # Calculate total number of logits needed for decoder output
    total_logits = sum(question_info)
    
    # Use the new categorical_multi modality
    survey_data = {
        "surveys": jnp.array(raw_responses, dtype=jnp.int32)
    }
    
    modality_specs = {
        "surveys": {
            "type": "categorical_multi",
            "dim": total_logits,  # Total number of logits across all questions
            "likelihood_params": {
                "categories_per_question": question_info
            }
        }
    }
    
    print(f"\nUsing categorical_multi modality")
    print(f"Survey data shape: {survey_data['surveys'].shape}")
    print(f"Total decoder logits needed: {total_logits}")
    print(f"Categories per question: {question_info[:5]}... (showing first 5)")
    
    # Create and train model
    print("\nCreating model...")
    model = create_multimodal_vae(
        modality_specs=modality_specs,
        fusion_strategy="concatenate",
        latent_dim=2,  # Use 2D latent space for better representation
        encoder_hidden_dims=(32,16),
        decoder_hidden_dims=(),
        activation="relu"
    )
    
    print("Training model...")
    model.fit(survey_data, num_steps=5000, lr=1e-3, seed=42, progress_bar=False)
    
    # Sample from posterior
    print("Sampling from posterior...")
    posterior = model.sample_posterior(survey_data, num_samples=1000, seed=42)
    
    print(f"Posterior keys: {list(posterior.keys())}")
    
    # Check for latent samples
    if "z" in posterior:
        theta_samples = np.asarray(posterior["z"])
        print(f"Found 'z' samples with shape: {theta_samples.shape}")
    else:
        print("No 'z' samples found!")
        return
    
    # For 2D latent space, use first dimension for comparison
    if theta_samples.shape[-1] == 2:
        theta_samples_1d = theta_samples[:, :, 0]  # Use first dimension
        print(f"Using first dimension of 2D latent space: {theta_samples_1d.shape}")
    else:
        theta_samples_1d = theta_samples.squeeze(-1) if theta_samples.ndim == 3 else theta_samples
    
    # Compute posterior mean
    theta_mean = theta_samples_1d.mean(axis=0)
    
    # Align to truth
    theta_aligned, slope, intercept = _linear_align(theta_mean, true_ideal_points)
    
    # Compute correlation
    correlation = np.corrcoef(true_ideal_points, theta_aligned)[0, 1]
    
    # Compute R-squared
    fit = linregress(true_ideal_points, theta_aligned)
    r_squared = fit.rvalue**2
    
    # Calculate coverage rate
    cov, n_in = coverage_rate(theta_samples_1d, true_ideal_points)
    
    print(f"\nResults:")
    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}")
    print(f"Intercept: {intercept:.3f}")
    
    # Plot results
    plt.figure(figsize=(15, 8))
    
    plt.subplot(2, 2, 1)
    plt.scatter(true_ideal_points, theta_aligned, alpha=0.6)
    plt.plot([true_ideal_points.min(), true_ideal_points.max()], 
             [true_ideal_points.min(), true_ideal_points.max()], 'r--', alpha=0.5)
    plt.xlabel("True Ideal Points")
    plt.ylabel("Estimated Ideal Points")
    plt.title(f"Categorical Survey Model (R² = {r_squared:.3f})")
    plt.grid(True, alpha=0.3)
    
    plt.subplot(2, 2, 2)
    plt.hist(theta_aligned - true_ideal_points, bins=20, alpha=0.7)
    plt.xlabel("Estimation Error")
    plt.ylabel("Frequency")
    plt.title("Error Distribution")
    plt.grid(True, alpha=0.3)
    
    # Coverage visualization
    plt.subplot(2, 2, 3)
    sort_idx = np.argsort(true_ideal_points)
    x_sorted = true_ideal_points[sort_idx]
    theta_aligned_samples = theta_samples_1d * slope + intercept  # (S, D)
    mean_sorted = theta_aligned_samples.mean(axis=0)[sort_idx]
    std_sorted = theta_aligned_samples.std(axis=0)[sort_idx]
    
    plt.fill_between(range(len(x_sorted)),
                     mean_sorted - 1.96 * std_sorted,
                     mean_sorted + 1.96 * std_sorted,
                     alpha=0.3, label="95% CI")
    plt.scatter(range(len(x_sorted)), mean_sorted, alpha=0.7, s=20, label="Posterior Mean")
    plt.scatter(range(len(x_sorted)), x_sorted, alpha=0.7, s=20, color="red", label="True Values")
    plt.xlabel("Politician (sorted by true ideal point)")
    plt.ylabel("Ideal Point")
    plt.title(f"Coverage Visualization ({cov:.1f}%)")
    plt.legend()
    plt.grid(True, alpha=0.3)
    
    # Residuals plot
    plt.subplot(2, 2, 4)
    residuals = theta_aligned - true_ideal_points
    plt.scatter(true_ideal_points, residuals, alpha=0.6)
    plt.axhline(y=0, color='r', linestyle='--', alpha=0.5)
    plt.xlabel("True Ideal Points")
    plt.ylabel("Residuals")
    plt.title("Residuals vs True Values")
    plt.grid(True, alpha=0.3)
    
    plt.tight_layout()
    plt.show()
    
    print("\nTest completed successfully!")
    return correlation, r_squared


if __name__ == "__main__":
    test_surveys_only()
