"""
Test script demonstrating advanced variational families.
"""

import jax.numpy as jnp
import jax.random as random
from advanced_variational_families import AdvancedMultiModalVAE, demonstrate_advanced_posteriors

def create_test_data():
    """Create simple test data for demonstration."""
    key = random.PRNGKey(42)
    n_samples = 1000
    
    # Create synthetic multi-modal data
    data = {
        "embeddings": random.normal(key, (n_samples, 128)),  # Text embeddings
        "votes": random.bernoulli(random.split(key)[0], 0.6, (n_samples, 50)),  # Binary votes
    }
    
    modality_specs = {
        "embeddings": {
            "type": "gaussian",
            "dim": 128,
            "likelihood_params": {"scale": 1.0}
        },
        "votes": {
            "type": "bernoulli", 
            "dim": 50,
            "likelihood_params": {}
        }
    }
    
    return data, modality_specs


def test_single_family(family: str = "diagonal"):
    """Test a single variational family."""
    print(f"\nTesting {family} variational family...")
    
    data, modality_specs = create_test_data()
    
    model = AdvancedMultiModalVAE(
        modality_specs=modality_specs,
        latent_dim=5,
        variational_family=family,
        fusion_strategy="concatenate"
    )
    
    # Train with fewer steps for testing
    model.fit(data, num_steps=1000, progress_bar=False)
    
    print(f"✓ Training completed")
    
    # Sample from posterior
    samples = model.sample_posterior(data, num_samples=100)
    print(f"✓ Posterior sampling completed")
    print(f"  Latent samples shape: {samples['z'].shape}")
    
    return model


def compare_families():
    """Compare different variational families."""
    from advanced_variational_families import compare_variational_families
    
    data, modality_specs = create_test_data()
    
    # Test subset of families (some might be slow)
    families = ["diagonal", "full_cov", "low_rank"]
    
    results = compare_variational_families(
        data=data,
        modality_specs=modality_specs, 
        families=families
    )
    
    print("\n" + "="*60)
    print("COMPARISON RESULTS:")
    print("="*60)
    
    for family, result in results.items():
        if result['success']:
            print(f"{family:15}: Loss = {result['final_loss']:.4f}")
        else:
            print(f"{family:15}: FAILED - {result['error']}")


if __name__ == "__main__":
    print("Advanced Variational Families Demo")
    print("="*60)
    
    # Show theoretical differences
    demonstrate_advanced_posteriors()
    
    print("\n" + "="*60)
    print("PRACTICAL TESTS:")
    print("="*60)
    
    # Test individual families
    test_single_family("diagonal")
    test_single_family("full_cov") 
    
    # Compare multiple families
    print("\n" + "="*60)
    print("COMPARING FAMILIES:")
    print("="*60)
    compare_families()
