"""
Test script for three-covariate functionality (covariates_prior, covariates_decoder, covariates_predictor).
"""

import jax
import jax.numpy as jnp
import sys
sys.path.insert(0, 'c:/Users/Gauthier/Desktop/DeepLatent')

from src2.models import DeepLatent

def test_three_covariates():
    """Test that model works with three separate covariate types."""
    print("=" * 80)
    print("Testing Three-Covariate Functionality")
    print("=" * 80)
    
    # Setup dimensions
    N = 50
    embedding_dim = 20
    prior_cov_dim = 3
    decoder_cov_dim = 2
    predictor_cov_dim = 4
    latent_dim = 2
    
    print(f"\nTest Configuration:")
    print(f"  N samples: {N}")
    print(f"  Prior covariate dim: {prior_cov_dim}")
    print(f"  Decoder covariate dim: {decoder_cov_dim}")
    print(f"  Predictor covariate dim: {predictor_cov_dim}")
    print(f"  Latent dim: {latent_dim}")
    
    # Generate synthetic data
    key = jax.random.PRNGKey(42)
    key, *subkeys = jax.random.split(key, 6)
    
    data = {
        "embeddings": jax.random.normal(subkeys[0], (N, embedding_dim), dtype=jnp.float32)
    }
    
    covariates_prior = jax.random.normal(subkeys[1], (N, prior_cov_dim), dtype=jnp.float32)
    covariates_decoder = jax.random.normal(subkeys[2], (N, decoder_cov_dim), dtype=jnp.float32)
    covariates_predictor = jax.random.normal(subkeys[3], (N, predictor_cov_dim), dtype=jnp.float32)
    
    outcomes = {
        "vote_choice": jax.random.randint(subkeys[4], (N,), 0, 3, dtype=jnp.int32)
    }
    
    # Define modality specs
    modality_specs = {
        "embeddings": {"type": "gaussian", "dim": embedding_dim}
    }
    
    # Define predictor specs
    predictor_specs = {
        "vote_choice": {
            "output_dim": 3,
            "hidden_dims": (32, 16),
            "loss_fn": "cross_entropy"
        }
    }
    
    print("\n" + "=" * 80)
    print("1. Creating Model with Three Covariate Types")
    print("=" * 80)
    
    model = DeepLatent(
        modality_specs=modality_specs,
        latent_dim=latent_dim,
        encoder_hidden_dims=(64, 32),
        decoder_hidden_dims=(32,),
        use_structured_prior=True,
        prior_covariate_dim=prior_cov_dim,
        decoder_covariate_dim=decoder_cov_dim,
        predictor_covariate_dim=predictor_cov_dim,
        predictor_specs=predictor_specs
    )
    
    print("✓ Model created successfully")
    print(f"  - Prior covariate dim: {model.prior_covariate_dim}")
    print(f"  - Decoder covariate dim: {model.decoder_covariate_dim}")
    print(f"  - Predictor covariate dim: {model.predictor_covariate_dim}")
    
    print("\n" + "=" * 80)
    print("2. Training Model")
    print("=" * 80)
    
    try:
        model.fit(
            data,
            covariates_prior=covariates_prior,
            covariates_decoder=covariates_decoder,
            covariates_predictor=covariates_predictor,
            outcomes=outcomes,
            batch_size=25,
            num_steps=50,
            lr=1e-3,
            progress_bar=False,
            seed=42
        )
        print("✓ Model trained successfully")
        print(f"  - Final loss: {model.losses[-1]:.4f}")
    except Exception as e:
        print(f"✗ Training failed: {e}")
        raise
    
    print("\n" + "=" * 80)
    print("3. Sampling Posterior")
    print("=" * 80)
    
    try:
        posterior_samples = model.sample_posterior(
            data,
            covariates_prior=covariates_prior,
            covariates_decoder=covariates_decoder,
            covariates_predictor=covariates_predictor,
            outcomes=outcomes,
            num_samples=10,
            seed=1
        )
        print("✓ Posterior sampling successful")
        print(f"  - Sampled z shape: {posterior_samples['z'].shape}")
        if 'gamma' in posterior_samples:
            print(f"  - Gamma shape: {posterior_samples['gamma'].shape}")
    except Exception as e:
        print(f"✗ Posterior sampling failed: {e}")
        raise
    
    print("\n" + "=" * 80)
    print("4. Generating New Samples")
    print("=" * 80)
    
    try:
        # Generate 10 new samples with new covariates
        key, *subkeys = jax.random.split(key, 4)
        new_cov_prior = jax.random.normal(subkeys[0], (10, prior_cov_dim), dtype=jnp.float32)
        new_cov_decoder = jax.random.normal(subkeys[1], (10, decoder_cov_dim), dtype=jnp.float32)
        new_cov_predictor = jax.random.normal(subkeys[2], (10, predictor_cov_dim), dtype=jnp.float32)
        
        generated = model.generate_samples(
            num_samples=10,
            covariates_prior=new_cov_prior,
            covariates_decoder=new_cov_decoder,
            covariates_predictor=new_cov_predictor,
            seed=2
        )
        print("✓ Sample generation successful")
        print(f"  - Generated embeddings shape: {generated['embeddings'].shape}")
        if 'outcome_vote_choice' in generated:
            print(f"  - Generated outcomes shape: {generated['outcome_vote_choice'].shape}")
    except Exception as e:
        print(f"✗ Sample generation failed: {e}")
        raise
    
    print("\n" + "=" * 80)
    print("5. Testing Without Decoder Covariates (Optional)")
    print("=" * 80)
    
    # Create model without decoder covariates
    model_no_decoder_cov = DeepLatent(
        modality_specs=modality_specs,
        latent_dim=latent_dim,
        use_structured_prior=True,
        prior_covariate_dim=prior_cov_dim,
        decoder_covariate_dim=0,  # No decoder covariates
        predictor_covariate_dim=predictor_cov_dim,
        predictor_specs=predictor_specs
    )
    
    try:
        model_no_decoder_cov.fit(
            data,
            covariates_prior=covariates_prior,
            covariates_decoder=None,  # Not providing decoder covariates
            covariates_predictor=covariates_predictor,
            outcomes=outcomes,
            num_steps=20,
            progress_bar=False
        )
        print("✓ Model without decoder covariates trained successfully")
    except Exception as e:
        print(f"✗ Training without decoder covariates failed: {e}")
        raise
    
    print("\n" + "=" * 80)
    print("ALL TESTS PASSED!")
    print("=" * 80)
    print("\nSummary:")
    print("  ✓ Model creation with three covariate types")
    print("  ✓ Training with all three covariate types")
    print("  ✓ Posterior sampling with covariates")
    print("  ✓ Sample generation with new covariates")
    print("  ✓ Optional decoder covariates")
    print("\nThe three-covariate system is working correctly!")


if __name__ == "__main__":
    test_three_covariates()
