# Simple Text Embeddings Test Script
# ===================================
# A focused test script for Gaussian likelihood on text embeddings 
# following the pattern of test_votes_only.py and test_surveys_simple.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_text_embeddings_data(num_politicians: int = 500, embedding_dim: int = 100, seed: int = 42):
    """Create synthetic text embeddings data."""
    np.random.seed(seed)
    
    # Generate true ideal points (1D)
    true_ideal_points = np.random.normal(0, 1, num_politicians)
    
    # Generate text embeddings using a linear factor model
    # Each dimension of the embedding relates to the ideal point plus noise
    factor_loadings = np.random.normal(0, 0.5, embedding_dim)  # How each dimension relates to ideal point
    noise_scale = 0.3
    
    # Generate embeddings: each dimension = loading * ideal_point + noise
    embeddings = np.zeros((num_politicians, embedding_dim))
    for i in range(num_politicians):
        embeddings[i] = factor_loadings * true_ideal_points[i] + np.random.normal(0, noise_scale, embedding_dim)
    
    data = {
        "text_embeddings": jnp.array(embeddings, dtype=jnp.float32)
    }
    
    print(f"Generated text embeddings with:")
    print(f"  Politicians: {num_politicians}")
    print(f"  Embedding dimension: {embedding_dim}")
    print(f"  Embedding range: [{embeddings.min():.3f}, {embeddings.max():.3f}]")
    print(f"  True ideal points range: [{true_ideal_points.min():.3f}, {true_ideal_points.max():.3f}]")
    
    return data, true_ideal_points


def test_text_embeddings():
    """Test text embeddings model with Gaussian likelihood."""
    print("Testing Text Embeddings Model with Gaussian Likelihood")
    print("=" * 55)
    
    # Create data
    data, true_ideal_points = create_text_embeddings_data(num_politicians=1000, embedding_dim=128)
    print(f"Text embeddings shape: {data['text_embeddings'].shape}")
    print(f"True ideal points shape: {true_ideal_points.shape}")
    
    # Create model spec
    modality_specs = {
        "text_embeddings": {
            "type": "gaussian",
            "dim": data["text_embeddings"].shape[1],
            "likelihood_params": {"scale": 1.0}
        }
    }
    
    # Create and train model
    print("\nCreating model...")
    model = create_multimodal_vae(
        modality_specs=modality_specs,
        fusion_strategy="concatenate",
        latent_dim=1,
        encoder_hidden_dims=(64, 32),
        decoder_hidden_dims=(32,),
        activation="relu"
    )
    
    print("Training model...")
    model.fit(data, num_steps=5000, lr=1e-3, seed=42, progress_bar=False)
    
    # Sample from posterior
    print("Sampling from posterior...")
    posterior = model.sample_posterior(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
    
    # Squeeze if needed
    if theta_samples.ndim == 3 and theta_samples.shape[-1] == 1:
        theta_samples = theta_samples.squeeze(-1)
    
    # Compute posterior mean
    theta_mean = theta_samples.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, 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"Text Embeddings (Gaussian) (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 * 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_text_embeddings()
