# Simple Poisson Test Script
# ==========================
# A focused test script for Poisson likelihood 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_poisson_data(num_politicians: int = 500, num_activities: int = 30, seed: int = 42):
    """Create synthetic count data (e.g., number of bills sponsored, speeches given, etc.)."""
    np.random.seed(seed)
    
    # Generate true ideal points (1D) - reduced variance for stability
    true_ideal_points = np.random.normal(0, 0.8, num_politicians)
    
    # Generate count data using Poisson distribution
    count_data = np.zeros((num_politicians, num_activities), dtype=int)
    
    print(f"Creating Poisson count data...")
    
    for a in range(num_activities):
        # More conservative activity parameters to prevent very large counts
        baseline_rate = np.random.uniform(0.8, 3.0)  # Reduced range
        sensitivity = np.random.uniform(0.1, 0.5)    # Much reduced sensitivity
        
        for p in range(num_politicians):
            # Log-rate depends on ideal point but more conservatively
            log_rate = np.log(baseline_rate) + sensitivity * true_ideal_points[p]
            rate = np.exp(log_rate)
            
            # Sample from Poisson distribution
            count_data[p, a] = np.random.poisson(rate)
    
    data = {
        "count_data": jnp.array(count_data, dtype=jnp.int32)
    }
    
    print(f"Generated Poisson count data with:")
    print(f"  Politicians: {num_politicians}")
    print(f"  Activities: {num_activities}")
    print(f"  Count range: [{count_data.min()}, {count_data.max()}]")
    print(f"  Mean counts: {count_data.mean():.2f}")
    print(f"  True ideal points range: [{true_ideal_points.min():.3f}, {true_ideal_points.max():.3f}]")
    
    return data, true_ideal_points


def test_poisson():
    """Test Poisson count model."""
    print("Testing Poisson Likelihood Model")
    print("=" * 35)
    
    # Create data
    data, true_ideal_points = create_poisson_data(
        num_politicians=800, 
        num_activities=25
    )
    
    print(f"Count data shape: {data['count_data'].shape}")
    print(f"True ideal points shape: {true_ideal_points.shape}")
    
    # Create model spec
    modality_specs = {
        "count_data": {
            "type": "poisson",
            "dim": data["count_data"].shape[1],  # Number of activities
            "likelihood_params": {}
        }
    }
    
    # Create and train model
    print("\nCreating model...")
    model = create_multimodal_vae(
        modality_specs=modality_specs,
        fusion_strategy="concatenate",
        latent_dim=1,
        encoder_hidden_dims=(32, 16),
        decoder_hidden_dims=(),
        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"Poisson 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 * 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_poisson()
