#!/usr/bin/env python3
"""
Test script for joint VAE-predictor functionality.

This script demonstrates and tests the new joint training capability
where VAE parameters and outcome predictors are learned simultaneously.
"""

import sys
import os
sys.path.append(os.path.join(os.path.dirname(__file__), 'src2'))

import numpy as np
import jax
import jax.numpy as jnp
from jax import random
import matplotlib.pyplot as plt
from typing import Dict, Any
import time

from models import (
    create_joint_vae_predictor,
    MultiModalVAE,
    create_multimodal_vae,
    PredictorMLP
)


def generate_synthetic_data(n_samples: int = 1000, 
                          latent_dim: int = 2,
                          text_dim: int = 50,
                          covariate_dim: int = 3,
                          seed: int = 42) -> Dict[str, jnp.ndarray]:
    """Generate synthetic data for testing joint VAE-predictor."""
    rng = random.PRNGKey(seed)
    
    # Generate true latent ideal points
    rng, subkey = random.split(rng)
    true_ideal_points = random.normal(subkey, (n_samples, latent_dim)) * 2.0
    
    # Generate covariates (e.g., demographics)
    rng, subkey = random.split(rng)
    covariates = random.normal(subkey, (n_samples, covariate_dim))
    
    # Generate text embeddings from ideal points (with noise)
    rng, subkey = random.split(rng)
    text_loading = random.normal(subkey, (latent_dim, text_dim)) * 0.5
    text_noise = random.normal(subkey, (n_samples, text_dim)) * 0.1
    text_embeddings = true_ideal_points @ text_loading + text_noise
    
    # Generate outcomes based on ideal points and covariates
    # Outcome 1: Binary vote choice (based on first ideal point dimension)
    vote_prob = jax.nn.sigmoid(2.0 * true_ideal_points[:, 0] + 0.5 * covariates[:, 0])
    rng, subkey = random.split(rng)
    vote_choice = random.bernoulli(subkey, vote_prob).astype(jnp.float32)
    
    # Outcome 2: Multi-class party preference (3 classes)
    # Based on combination of ideal points and covariates
    party_logits = jnp.stack([
        -1.0 * true_ideal_points[:, 0] + 0.3 * covariates[:, 1],  # Party A
        0.5 * true_ideal_points[:, 0] + true_ideal_points[:, 1],   # Party B  
        0.8 * true_ideal_points[:, 1] - 0.5 * covariates[:, 2]    # Party C
    ], axis=1)
    party_probs = jax.nn.softmax(party_logits, axis=1)
    rng, subkey = random.split(rng)
    party_preference = random.categorical(subkey, party_logits)
    
    # Outcome 3: Continuous approval rating
    approval_mean = (
        1.5 * true_ideal_points[:, 0] - 
        0.8 * true_ideal_points[:, 1] + 
        0.6 * covariates[:, 0] + 
        2.0
    )
    rng, subkey = random.split(rng)
    approval_rating = approval_mean + random.normal(subkey, (n_samples,)) * 0.5
    
    # Convert party preference to one-hot for cross-entropy loss
    party_onehot = jax.nn.one_hot(party_preference, 3)
    
    return {
        'text_embeddings': text_embeddings,
        'covariates': covariates,
        'true_ideal_points': true_ideal_points,
        'vote_choice': vote_choice,
        'party_preference': party_preference,
        'party_onehot': party_onehot,
        'approval_rating': approval_rating,
        'vote_prob': vote_prob,
        'party_probs': party_probs,
        'approval_mean': approval_mean
    }


def test_basic_joint_training():
    """Test basic joint VAE-predictor training."""
    print("=" * 60)
    print("TEST 1: Basic Joint Training")
    print("=" * 60)
    
    # Generate synthetic data
    data = generate_synthetic_data(n_samples=500, latent_dim=2, seed=42)
    
    # Define modality specifications
    modality_specs = {
        "text_embeddings": {
            "type": "gaussian",
            "dim": data['text_embeddings'].shape[1]
        }
    }
    
    # Define predictor specifications
    predictor_specs = {
        "vote_choice": {
            "output_dim": 1,
            "loss_fn": "binary_cross_entropy",
            "hidden_dims": (64, 32),
            "output_activation": "sigmoid"
        },
        "approval_rating": {
            "output_dim": 1,
            "loss_fn": "mse",
            "hidden_dims": (64, 32),
            "obs_scale": 0.5
        }
    }
    
    # Create joint model
    model = create_joint_vae_predictor(
        modality_specs=modality_specs,
        predictor_specs=predictor_specs,
        latent_dim=2,
        covariate_dim=3,
        encoder_hidden_dims=(128, 64),
        decoder_hidden_dims=(64,)
    )
    
    print(f"Created model with {len(model.predictor_mlps)} predictors")
    print(f"Predictor names: {list(model.predictor_mlps.keys())}")
    
    # Prepare data for training
    train_data = {"text_embeddings": data['text_embeddings']}
    train_outcomes = {
        "vote_choice": data['vote_choice'],
        "approval_rating": data['approval_rating']
    }
    
    # Joint training
    print("\nStarting joint training...")
    start_time = time.time()
    
    model.fit(
        data=train_data,
        covariates=data['covariates'],
        outcomes=train_outcomes,
        num_steps=2000,
        lr=1e-3,
        seed=123,
        progress_bar=True
    )
    
    train_time = time.time() - start_time
    print(f"Training completed in {train_time:.2f} seconds")
    
    # Test posterior sampling
    print("\nSampling from posterior...")
    posterior_samples = model.sample_posterior(
        data=train_data,
        covariates=data['covariates'],
        outcomes=train_outcomes,
        num_samples=100,
        seed=456
    )
    
    # Extract results
    inferred_ideal_points = posterior_samples["z"].mean(axis=0)  # Average over samples
    vote_predictions = posterior_samples["predictor_vote_choice_output"].mean(axis=0)
    approval_predictions = posterior_samples["predictor_approval_rating_output"].mean(axis=0)
    
    # Evaluate predictions
    vote_accuracy = np.mean((vote_predictions.squeeze() > 0.5) == data['vote_choice'])
    approval_mse = np.mean((approval_predictions.squeeze() - data['approval_rating']) ** 2)
    
    print(f"\nResults:")
    print(f"Vote choice accuracy: {vote_accuracy:.3f}")
    print(f"Approval rating MSE: {approval_mse:.3f}")
    
    # Correlation with true ideal points
    true_points = data['true_ideal_points']
    for dim in range(2):
        corr = np.corrcoef(true_points[:, dim], inferred_ideal_points[:, dim])[0, 1]
        print(f"Ideal point dimension {dim} correlation: {corr:.3f}")
    
    return model, data, posterior_samples


def test_structured_prior_joint_training():
    """Test joint training with structured prior."""
    print("\n" + "=" * 60)
    print("TEST 2: Joint Training with Structured Prior")
    print("=" * 60)
    
    # Generate synthetic data
    data = generate_synthetic_data(n_samples=400, latent_dim=2, seed=84)
    
    # Define modality and predictor specifications
    modality_specs = {
        "text_embeddings": {
            "type": "gaussian", 
            "dim": data['text_embeddings'].shape[1]
        }
    }
    
    predictor_specs = {
        "party_preference": {
            "output_dim": 3,
            "loss_fn": "cross_entropy",
            "hidden_dims": (64,),
            "output_activation": "softmax"
        }
    }
    
    # Create model with structured prior
    model = create_joint_vae_predictor(
        modality_specs=modality_specs,
        predictor_specs=predictor_specs,
        latent_dim=2,
        covariate_dim=3,
        use_structured_prior=True,
        encoder_hidden_dims=(64, 32),
        decoder_hidden_dims=(32,)
    )
    
    print("Created model with structured prior")
    
    # Training data
    train_data = {"text_embeddings": data['text_embeddings']}
    train_outcomes = {"party_preference": data['party_preference']}
    
    # Joint training
    print("Starting joint training with structured prior...")
    model.fit(
        data=train_data,
        covariates=data['covariates'],
        outcomes=train_outcomes,
        num_steps=1500,
        lr=1e-3,
        seed=789
    )
    
    # Sample posterior
    posterior_samples = model.sample_posterior(
        data=train_data,
        covariates=data['covariates'],
        outcomes=train_outcomes,
        num_samples=100
    )
    
    # Extract predictions
    party_predictions = posterior_samples["predictor_party_preference_output"].mean(axis=0)
    predicted_classes = jnp.argmax(party_predictions, axis=1)
    
    # Evaluate
    party_accuracy = np.mean(predicted_classes == data['party_preference'])
    print(f"\nParty preference accuracy: {party_accuracy:.3f}")
    
    # Check structured prior parameters
    if 'gamma' in posterior_samples:
        gamma_mean = posterior_samples['gamma'].mean(axis=0)
        print(f"Gamma shape (covariates -> latent): {gamma_mean.shape}")
        print(f"Gamma mean effects:\n{gamma_mean}")
    
    return model, data


def test_comparison_with_separate_training():
    """Compare joint training vs separate training."""
    print("\n" + "=" * 60)
    print("TEST 3: Joint vs Separate Training Comparison")
    print("=" * 60)
    
    # Generate synthetic data
    data = generate_synthetic_data(n_samples=600, latent_dim=2, seed=100)
    
    modality_specs = {
        "text_embeddings": {
            "type": "gaussian",
            "dim": data['text_embeddings'].shape[1]
        }
    }
    
    train_data = {"text_embeddings": data['text_embeddings']}
    
    # Method 1: Joint training
    print("Method 1: Joint Training")
    predictor_specs = {
        "vote_choice": {
            "output_dim": 1,
            "loss_fn": "binary_cross_entropy",
            "output_activation": "sigmoid"
        }
    }
    
    joint_model = create_joint_vae_predictor(
        modality_specs=modality_specs,
        predictor_specs=predictor_specs,
        latent_dim=2,
        covariate_dim=3
    )
    
    start_time = time.time()
    joint_model.fit(
        data=train_data,
        covariates=data['covariates'],
        outcomes={"vote_choice": data['vote_choice']},
        num_steps=1500,
        lr=1e-3,
        progress_bar=False
    )
    joint_time = time.time() - start_time
    
    # Get joint predictions
    joint_samples = joint_model.sample_posterior(
        data=train_data,
        covariates=data['covariates'],
        outcomes={"vote_choice": data['vote_choice']},
        num_samples=100
    )
    joint_predictions = joint_samples["predictor_vote_choice_output"].mean(axis=0).squeeze()
    joint_accuracy = np.mean((joint_predictions > 0.5) == data['vote_choice'])
    
    # Method 2: Separate training (VAE first, then predictor)
    print("Method 2: Separate Training")
    
    # Train VAE only
    separate_model = create_multimodal_vae(
        modality_specs=modality_specs,
        latent_dim=2,
        covariate_dim=3
    )
    
    start_time = time.time()
    separate_model.fit(
        data=train_data,
        covariates=data['covariates'],
        num_steps=1500,
        lr=1e-3,
        progress_bar=False
    )
    
    # Extract ideal points
    vae_samples = separate_model.sample_posterior(
        data=train_data,
        covariates=data['covariates'],
        num_samples=100
    )
    ideal_points = vae_samples["z"].mean(axis=0)  # Use posterior mean
    
    # Train separate predictor
    predictor = PredictorMLP(
        output_dim=1,
        input_dim=2 + 3,  # latent_dim + covariate_dim
        hidden_dims=(128, 64),
        output_activation="sigmoid"
    )
    
    # Simple training loop for predictor
    from jax import grad
    predictor_input = jnp.concatenate([ideal_points, data['covariates']], axis=1)
    
    rng = random.PRNGKey(42)
    params = predictor.init(rng, predictor_input[:1])["params"]
    
    def loss_fn(params, X, y):
        preds = predictor.apply({"params": params}, X)
        return jnp.mean(-(y * jnp.log(preds + 1e-15) + (1-y) * jnp.log(1-preds + 1e-15)))
    
    # Simple gradient descent
    lr = 1e-3
    for i in range(1000):
        loss_val = loss_fn(params, predictor_input, data['vote_choice'])
        grads = grad(loss_fn)(params, predictor_input, data['vote_choice'])
        params = jax.tree_util.tree_map(lambda p, g: p - lr * g, params, grads)
    
    separate_time = time.time() - start_time
    
    # Get separate predictions
    separate_predictions = predictor.apply({"params": params}, predictor_input).squeeze()
    separate_accuracy = np.mean((separate_predictions > 0.5) == data['vote_choice'])
    
    # Compare results
    print(f"\nComparison Results:")
    print(f"Joint Training   - Accuracy: {joint_accuracy:.3f}, Time: {joint_time:.2f}s")
    print(f"Separate Training - Accuracy: {separate_accuracy:.3f}, Time: {separate_time:.2f}s")
    print(f"Joint training accuracy improvement: {joint_accuracy - separate_accuracy:.3f}")
    
    return joint_model, separate_model, data


def test_multiple_outcomes():
    """Test joint training with multiple diverse outcomes."""
    print("\n" + "=" * 60)
    print("TEST 4: Multiple Diverse Outcomes")
    print("=" * 60)
    
    # Generate synthetic data
    data = generate_synthetic_data(n_samples=800, latent_dim=3, seed=200)
    
    modality_specs = {
        "text_embeddings": {
            "type": "gaussian",
            "dim": data['text_embeddings'].shape[1]
        }
    }
    
    # Multiple outcome types
    predictor_specs = {
        "vote_binary": {
            "output_dim": 1,
            "loss_fn": "binary_cross_entropy",
            "output_activation": "sigmoid",
            "hidden_dims": (64, 32)
        },
        "party_multiclass": {
            "output_dim": 3,
            "loss_fn": "cross_entropy", 
            "output_activation": "softmax",
            "hidden_dims": (64, 32)
        },
        "approval_continuous": {
            "output_dim": 1,
            "loss_fn": "mse",
            "hidden_dims": (64, 32),
            "obs_scale": 0.5
        }
    }
    
    model = create_joint_vae_predictor(
        modality_specs=modality_specs,
        predictor_specs=predictor_specs,
        latent_dim=3,
        covariate_dim=3,
        encoder_hidden_dims=(128, 64),
        decoder_hidden_dims=(64,)
    )
    
    train_data = {"text_embeddings": data['text_embeddings']}
    train_outcomes = {
        "vote_binary": data['vote_choice'],
        "party_multiclass": data['party_preference'],
        "approval_continuous": data['approval_rating']
    }
    
    print("Training model with multiple outcome types...")
    model.fit(
        data=train_data,
        covariates=data['covariates'],
        outcomes=train_outcomes,
        num_steps=2000,
        lr=1e-3,
        progress_bar=True
    )
    
    # Evaluate all outcomes
    posterior_samples = model.sample_posterior(
        data=train_data,
        covariates=data['covariates'],
        outcomes=train_outcomes,
        num_samples=100
    )
    
    # Binary vote prediction
    vote_preds = posterior_samples["predictor_vote_binary_output"].mean(axis=0).squeeze()
    vote_acc = np.mean((vote_preds > 0.5) == data['vote_choice'])
    
    # Multiclass party prediction
    party_preds = posterior_samples["predictor_party_multiclass_output"].mean(axis=0)
    party_classes = jnp.argmax(party_preds, axis=1)
    party_acc = np.mean(party_classes == data['party_preference'])
    
    # Continuous approval prediction
    approval_preds = posterior_samples["predictor_approval_continuous_output"].mean(axis=0).squeeze()
    approval_mse = np.mean((approval_preds - data['approval_rating']) ** 2)
    
    print(f"\nMultiple Outcome Results:")
    print(f"Binary vote accuracy: {vote_acc:.3f}")
    print(f"Party preference accuracy: {party_acc:.3f}")
    print(f"Approval rating MSE: {approval_mse:.3f}")
    
    return model, data


def visualize_results(model, data, posterior_samples):
    """Create visualizations of the results."""
    try:
        import matplotlib.pyplot as plt
        
        fig, axes = plt.subplots(2, 2, figsize=(12, 10))
        
        # Plot 1: True vs Inferred Ideal Points (dimension 1)
        true_points = data['true_ideal_points']
        inferred_points = posterior_samples["z"].mean(axis=0)
        
        axes[0, 0].scatter(true_points[:, 0], inferred_points[:, 0], alpha=0.6)
        axes[0, 0].plot([-3, 3], [-3, 3], 'r--', alpha=0.5)
        axes[0, 0].set_xlabel('True Ideal Points (Dim 1)')
        axes[0, 0].set_ylabel('Inferred Ideal Points (Dim 1)')
        axes[0, 0].set_title('True vs Inferred Ideal Points')
        
        # Plot 2: Vote Choice Predictions
        if "predictor_vote_choice_output" in posterior_samples:
            vote_preds = posterior_samples["predictor_vote_choice_output"].mean(axis=0).squeeze()
            vote_true = data['vote_choice']
            
            axes[0, 1].scatter(true_points[:, 0], vote_preds, 
                             c=vote_true, cmap='RdYlBu', alpha=0.6)
            axes[0, 1].set_xlabel('True Ideal Points (Dim 1)')
            axes[0, 1].set_ylabel('Predicted Vote Probability')
            axes[0, 1].set_title('Vote Choice Predictions')
            
        # Plot 3: Approval Rating Predictions
        if "predictor_approval_rating_output" in posterior_samples:
            approval_preds = posterior_samples["predictor_approval_rating_output"].mean(axis=0).squeeze()
            approval_true = data['approval_rating']
            
            axes[1, 0].scatter(approval_true, approval_preds, alpha=0.6)
            axes[1, 0].plot([approval_true.min(), approval_true.max()], 
                           [approval_true.min(), approval_true.max()], 'r--', alpha=0.5)
            axes[1, 0].set_xlabel('True Approval Rating')
            axes[1, 0].set_ylabel('Predicted Approval Rating')
            axes[1, 0].set_title('Approval Rating Predictions')
        
        # Plot 4: Ideal Points in 2D space colored by outcomes
        if inferred_points.shape[1] >= 2:
            vote_colors = data['vote_choice'] if 'vote_choice' in data else np.zeros(len(inferred_points))
            scatter = axes[1, 1].scatter(inferred_points[:, 0], inferred_points[:, 1], 
                                       c=vote_colors, cmap='RdYlBu', alpha=0.6)
            axes[1, 1].set_xlabel('Ideal Point Dimension 1')
            axes[1, 1].set_ylabel('Ideal Point Dimension 2')
            axes[1, 1].set_title('Inferred Ideal Points (colored by vote)')
            plt.colorbar(scatter, ax=axes[1, 1])
        
        plt.tight_layout()
        plt.savefig('joint_vae_predictor_results.png', dpi=150, bbox_inches='tight')
        plt.show()
        
        print("\nVisualization saved as 'joint_vae_predictor_results.png'")
        
    except ImportError:
        print("Matplotlib not available, skipping visualization")
    except Exception as e:
        print(f"Error creating visualization: {e}")


def main():
    """Run all tests."""
    print("🚀 Testing Joint VAE-Predictor Functionality")
    print("=" * 60)
    
    # Run tests
    try:
        # Test 1: Basic joint training
        model1, data1, samples1 = test_basic_joint_training()
        
        # Test 2: Structured prior
        model2, data2 = test_structured_prior_joint_training()
        
        # Test 3: Comparison with separate training
        joint_model, separate_model, data3 = test_comparison_with_separate_training()
        
        # Test 4: Multiple outcomes
        model4, data4 = test_multiple_outcomes()
        
        # Create visualization for the first test
        visualize_results(model1, data1, samples1)
        
        print("\n" + "=" * 60)
        print("✅ ALL TESTS COMPLETED SUCCESSFULLY!")
        print("=" * 60)
        print("\nKey findings:")
        print("1. Joint training successfully learns VAE and predictors together")
        print("2. Structured priors work with joint prediction tasks") 
        print("3. Joint training can outperform separate training")
        print("4. Multiple diverse outcome types can be handled simultaneously")
        print("5. The framework is flexible and extensible")
        
    except Exception as e:
        print(f"\n❌ TEST FAILED: {e}")
        import traceback
        traceback.print_exc()
        return False
        
    return True


if __name__ == "__main__":
    success = main()
    sys.exit(0 if success else 1)