#!/usr/bin/env python3
"""
Test script for logistic normal priors (standard and structured).

This script demonstrates the usage of both standard and structured
logistic normal priors in the MultiModalVAE framework.
"""

import numpy as np
import jax.numpy as jnp
from jax import random

from models import (
    create_logistic_normal_vae,
    create_structured_logistic_normal_vae,
    MultiModalVAE
)


def test_standard_logistic_normal():
    """Test standard logistic normal prior (no covariates)."""
    print("\n=== Testing Standard Logistic Normal Prior ===")
    
    # Create synthetic data
    np.random.seed(42)
    N = 100
    embedding_dim = 50
    latent_dim = 5
    
    # Generate some synthetic embedding data
    embeddings = np.random.randn(N, embedding_dim).astype(np.float32)
    
    # Create model with standard logistic normal prior
    modality_specs = {
        "embeddings": {"type": "gaussian", "dim": embedding_dim}
    }
    
    model = create_logistic_normal_vae(
        modality_specs=modality_specs,
        latent_dim=latent_dim,
        encoder_hidden_dims=(64,),
        decoder_hidden_dims=(64,)
    )
    
    print(f"✓ Created model with standard logistic normal prior (latent_dim={latent_dim})")
    
    # Prepare data
    data = {"embeddings": jnp.array(embeddings)}
    
    # Fit model (short training for testing)
    print("Training model...")
    model.fit(
        data=data,
        num_steps=100,
        lr=1e-3,
        batch_size=32,  # Test with mini-batch training
        seed=42,
        progress_bar=False
    )
    
    print("✓ Training completed successfully")
    
    # Sample from posterior
    print("Sampling from posterior...")
    posterior_samples = model.sample_posterior(
        data=data,
        num_samples=50,
        seed=123
    )
    
    # Check that z is on the simplex (all elements positive, sum to 1)
    if "z_simplex" in posterior_samples:
        z_samples = posterior_samples["z_simplex"]  # Shape: (num_samples, N, latent_dim)
        simplex_var_name = "z_simplex"
    else:
        # Fallback to z (for Gaussian priors)
        z_samples = posterior_samples["z"]  # Shape: (num_samples, N, latent_dim)
        simplex_var_name = "z"
    
    # Check positivity
    assert jnp.all(z_samples >= 0), "Logistic normal samples should be non-negative"
    
    # Check that each sample sums to 1 (simplex constraint)
    z_sums = jnp.sum(z_samples, axis=-1)  # Shape: (num_samples, N)
    assert jnp.allclose(z_sums, 1.0, atol=1e-6), "Logistic normal samples should sum to 1"
    
    print(f"✓ Posterior samples shape: {z_samples.shape}")
    print(f"✓ All samples are on simplex (positive and sum to 1) - using {simplex_var_name}")
    print(f"✓ Sample range: [{jnp.min(z_samples):.4f}, {jnp.max(z_samples):.4f}]")
    print(f"✓ Mean simplex entry: {jnp.mean(z_samples):.4f}")
    
    # Generate new samples
    print("Generating new samples...")
    generated = model.generate_samples(num_samples=20, seed=456)
    
    print(f"✓ Generated {generated['embeddings'].shape[0]} new samples")
    print("✓ Standard logistic normal prior test passed!")
    
    return model, data


def test_structured_logistic_normal():
    """Test structured logistic normal prior with covariates."""
    print("\n=== Testing Structured Logistic Normal Prior ===")
    
    # Create synthetic data with covariates
    np.random.seed(42)
    N = 120
    embedding_dim = 40
    latent_dim = 4
    covariate_dim = 3
    
    # Generate synthetic embedding data
    embeddings = np.random.randn(N, embedding_dim).astype(np.float32)
    
    # Generate synthetic covariates (e.g., party affiliation, region, etc.)
    covariates = np.random.randn(N, covariate_dim).astype(np.float32)
    
    # Create model with structured logistic normal prior
    modality_specs = {
        "embeddings": {"type": "gaussian", "dim": embedding_dim}
    }
    
    model = create_structured_logistic_normal_vae(
        modality_specs=modality_specs,
        covariate_dim=covariate_dim,
        latent_dim=latent_dim,
        encoder_hidden_dims=(64,),
        decoder_hidden_dims=(64,)
    )
    
    print(f"✓ Created model with structured logistic normal prior")
    print(f"  - latent_dim: {latent_dim}")
    print(f"  - covariate_dim: {covariate_dim}")
    
    # Prepare data
    data = {"embeddings": jnp.array(embeddings)}
    covariates_jax = jnp.array(covariates)
    
    # Fit model (short training for testing)
    print("Training model with covariates...")
    model.fit(
        data=data,
        covariates=covariates_jax,
        num_steps=150,
        lr=1e-3,
        batch_size=None,  # Use full-batch for compatibility with AutoGuides
        seed=42,
        progress_bar=False
    )
    
    print("✓ Training with structured prior completed successfully")
    
    # Sample from posterior
    print("Sampling from posterior...")
    posterior_samples = model.sample_posterior(
        data=data,
        covariates=covariates_jax,
        num_samples=50,
        seed=123
    )
    
    # Check simplex constraint
    if "z_simplex" in posterior_samples:
        z_samples = posterior_samples["z_simplex"]  # Shape: (num_samples, N, latent_dim)
        simplex_var_name = "z_simplex"
    else:
        # Fallback to z (for Gaussian priors)
        z_samples = posterior_samples["z"]  # Shape: (num_samples, N, latent_dim)
        simplex_var_name = "z"
    
    # Check positivity
    assert jnp.all(z_samples >= 0), "Structured logistic normal samples should be non-negative"
    
    # Check that each sample sums to 1 (simplex constraint)
    z_sums = jnp.sum(z_samples, axis=-1)  # Shape: (num_samples, N)
    assert jnp.allclose(z_sums, 1.0, atol=1e-6), "Structured logistic normal samples should sum to 1"
    
    print(f"✓ Posterior samples shape: {z_samples.shape}")
    print(f"✓ All samples are on simplex (positive and sum to 1) - using {simplex_var_name}")
    
    # Check that we learned structured prior parameters
    if "gamma" in posterior_samples:
        gamma_samples = posterior_samples["gamma"]
        print(f"✓ Learned regression coefficients gamma shape: {gamma_samples.shape}")
        print(f"✓ Gamma range: [{jnp.min(gamma_samples):.4f}, {jnp.max(gamma_samples):.4f}]")
    
    if "prior_mean" in posterior_samples:
        prior_mean_samples = posterior_samples["prior_mean"]
        print(f"✓ Prior mean varies with covariates, shape: {prior_mean_samples.shape}")
    
    # Generate new samples with different covariates
    print("Generating new samples with custom covariates...")
    new_covariates = np.random.randn(10, covariate_dim).astype(np.float32)
    generated = model.generate_samples(
        num_samples=10, 
        covariates=jnp.array(new_covariates),
        seed=456
    )
    
    print(f"✓ Generated {generated['embeddings'].shape[0]} new samples with custom covariates")
    print("✓ Structured logistic normal prior test passed!")
    
    return model, data, covariates_jax


def test_simplex_properties():
    """Test mathematical properties of the simplex constraint."""
    print("\n=== Testing Simplex Properties ===")
    
    # Quick test with minimal model
    modality_specs = {"test": {"type": "gaussian", "dim": 10}}
    
    model = MultiModalVAE(
        modality_specs=modality_specs,
        latent_dim=5,
        prior_type="logistic_normal",
        encoder_hidden_dims=(32,),
        decoder_hidden_dims=(32,)
    )
    
    # Generate synthetic data
    N = 50
    test_data = {"test": jnp.array(np.random.randn(N, 10).astype(np.float32))}
    
    # Quick training
    model.fit(test_data, num_steps=50, lr=1e-2, progress_bar=False, seed=42)
    
    # Sample and check properties
    samples = model.sample_posterior(test_data, num_samples=100, seed=42)
    if "z_simplex" in samples:
        z = samples["z_simplex"]
        simplex_var_name = "z_simplex"
    else:
        z = samples["z"]
        simplex_var_name = "z"
    
    # Test 1: All entries non-negative
    assert jnp.all(z >= 0), "All simplex entries must be non-negative"
    
    # Test 2: Each row sums to 1
    row_sums = jnp.sum(z, axis=-1)
    assert jnp.allclose(row_sums, 1.0, atol=1e-6), "Each simplex row must sum to 1"
    
    # Test 3: No entry should be exactly 0 or 1 (except in degenerate cases)
    # This is a property of the softmax transform
    assert jnp.all(z > 1e-10), "Softmax should not produce exact zeros"
    assert jnp.all(z < 1.0 - 1e-10), "Softmax should not produce exact ones in multi-dimensional case"
    
    print(f"✓ Simplex constraint satisfied: all entries non-negative and sum to 1 - using {simplex_var_name}")
    print("✓ Softmax properties verified: no exact zeros or ones")
    print(f"✓ Entry range: [{jnp.min(z):.6f}, {jnp.max(z):.6f}]")


def main():
    """Run all tests for logistic normal priors."""
    print("Testing Logistic Normal Priors Implementation")
    print("=" * 50)
    
    try:
        # Test 1: Standard logistic normal prior
        model1, data1 = test_standard_logistic_normal()
        
        # Test 2: Structured logistic normal prior
        model2, data2, covariates2 = test_structured_logistic_normal()
        
        # Test 3: Mathematical properties
        test_simplex_properties()
        
        print("\n" + "=" * 50)
        print("🎉 ALL TESTS PASSED! 🎉")
        print("\nSummary:")
        print("✓ Standard logistic normal prior works correctly")
        print("✓ Structured logistic normal prior works correctly") 
        print("✓ Simplex constraints are properly enforced")
        print("✓ Integration with existing VAE framework successful")
        print("\nThe logistic normal priors are ready for use!")
        
        return {
            "standard_model": model1,
            "structured_model": model2,
            "test_data": (data1, data2, covariates2)
        }
        
    except Exception as e:
        print(f"\n❌ TEST FAILED: {e}")
        import traceback
        traceback.print_exc()
        return None


if __name__ == "__main__":
    results = main()