"""
Test script for the enhanced MultiModalVAE with advanced variational families.
"""

import jax.numpy as jnp
import jax.random as random
from models import MultiModalVAE, create_multimodal_vae

def test_variational_families():
    """Test different variational families in the updated MultiModalVAE."""
    
    # Create test data
    key = random.PRNGKey(42)
    n_samples = 500
    
    data = {
        "embeddings": random.normal(key, (n_samples, 64)),
        "votes": random.bernoulli(random.split(key)[0], 0.6, (n_samples, 20)),
    }
    
    modality_specs = {
        "embeddings": {
            "type": "gaussian",
            "dim": 64,
            "likelihood_params": {"scale": 1.0}
        },
        "votes": {
            "type": "bernoulli", 
            "dim": 20,
            "likelihood_params": {}
        }
    }
    
    # Test different variational families
    families_to_test = ["diagonal", "full_cov", "low_rank"]  # Start with these
    
    results = {}
    
    for family in families_to_test:
        print(f"\nTesting variational family: {family}")
        print("-" * 40)
        
        try:
            # Create model with specific variational family
            model = MultiModalVAE(
                modality_specs=modality_specs,
                latent_dim=8,
                variational_family=family,
                fusion_strategy="concatenate"
            )
            
            print(f"✓ Model created successfully")
            
            # Train with fewer steps for testing
            model.fit(data, num_steps=500, progress_bar=False)
            print(f"✓ Training completed")
            
            # Sample from posterior
            samples = model.sample_posterior(data, num_samples=50)
            print(f"✓ Posterior sampling completed")
            print(f"  Latent samples shape: {samples['z'].shape}")
            
            # Store results
            results[family] = {
                'success': True,
                'final_loss': model.losses[-1] if hasattr(model, 'losses') else None,
                'model': model
            }
            
        except Exception as e:
            print(f"✗ Failed: {e}")
            results[family] = {
                'success': False,
                'error': str(e),
                'model': None
            }
    
    # Print summary
    print("\n" + "="*60)
    print("SUMMARY")
    print("="*60)
    
    for family, result in results.items():
        if result['success']:
            loss_str = f"Loss: {result['final_loss']:.4f}" if result['final_loss'] else "Loss: N/A"
            print(f"{family:15}: ✓ SUCCESS - {loss_str}")
        else:
            print(f"{family:15}: ✗ FAILED - {result['error']}")
    
    return results


def test_convenience_function():
    """Test the updated convenience function."""
    print("\nTesting convenience function...")
    
    modality_specs = {
        "text": {
            "type": "gaussian",
            "dim": 32,
            "likelihood_params": {"scale": 1.0}
        }
    }
    
    # Test creating models with different variational families
    for family in ["diagonal", "full_cov"]:
        model = create_multimodal_vae(
            modality_specs=modality_specs,
            variational_family=family,
            latent_dim=5
        )
        print(f"✓ Created model with {family} variational family")


def demonstrate_usage():
    """Show how to use the new variational families."""
    print("\nUSAGE EXAMPLES:")
    print("="*60)
    
    print("""
# 1. Diagonal Gaussian (default - your current approach)
model = MultiModalVAE(
    modality_specs=specs,
    variational_family="diagonal"  # Default
)

# 2. Full covariance Gaussian (captures correlations)
model = MultiModalVAE(
    modality_specs=specs,
    variational_family="full_cov"
)

# 3. Low-rank approximation (good compromise)
model = MultiModalVAE(
    modality_specs=specs,
    variational_family="low_rank"
)

# 4. Inverse Autoregressive Flow (very flexible)
model = MultiModalVAE(
    modality_specs=specs,
    variational_family="iaf"
)

# 5. Block Neural Autoregressive Flow (most flexible)
model = MultiModalVAE(
    modality_specs=specs,
    variational_family="bnaf"
)
""")


if __name__ == "__main__":
    print("Testing Enhanced MultiModalVAE with Advanced Variational Families")
    print("="*70)
    
    # Show usage examples
    demonstrate_usage()
    
    # Test different families
    results = test_variational_families()
    
    # Test convenience function
    test_convenience_function()
    
    print(f"\n✓ All tests completed!")
    print(f"✓ You can now use variational_family='full_cov', 'low_rank', 'iaf', or 'bnaf'")
    print(f"✓ in your MultiModalVAE constructor!")
