"""
Test script for batching functionality in MultiModalVAE.

This script tests the batching implementation to ensure:
1. Backward compatibility with existing code
2. Correct functionality of mini-batch training
3. Memory-efficient sampling and generation
4. Proper handling of covariates and outcomes
"""

import jax.numpy as jnp
from jax import random
import sys
import os

# Add the current directory to the path to import models
sys.path.append(os.path.dirname(os.path.abspath(__file__)))

from models import MultiModalVAE, create_batches, get_batch_size_recommendation


def test_data_batching_utilities():
    """Test the basic batching utility functions."""
    print("\n=== Testing Data Batching Utilities ===")
    
    # Create test data
    rng_key = random.PRNGKey(42)
    n_samples = 100
    
    data = {
        "embeddings": random.normal(rng_key, (n_samples, 50)),
        "votes": random.randint(rng_key, (n_samples, 10), 0, 3)
    }
    
    covariates = random.normal(rng_key, (n_samples, 5))
    outcomes = {
        "labels": random.randint(rng_key, (n_samples, 1), 0, 2)
    }
    
    # Test batch creation without shuffling
    batches = create_batches(data, covariates, outcomes, batch_size=25, shuffle=False)
    assert len(batches) == 4, f"Expected 4 batches, got {len(batches)}"
    
    # Check batch sizes
    for i, (batch_data, batch_covariates, batch_outcomes) in enumerate(batches):
        expected_size = 25 if i < 3 else 25  # Last batch might be smaller
        actual_size = next(iter(batch_data.values())).shape[0]
        assert actual_size == expected_size, f"Batch {i} size mismatch: {actual_size} vs {expected_size}"
        
        # Check covariate alignment
        assert batch_covariates.shape[0] == actual_size
        assert batch_outcomes["labels"].shape[0] == actual_size
    
    # Test with shuffling
    batches_shuffled = create_batches(data, covariates, outcomes, batch_size=25, shuffle=True, seed=123)
    
    # The data should be different when shuffled
    first_batch_orig = batches[0][0]["embeddings"]
    first_batch_shuffled = batches_shuffled[0][0]["embeddings"]
    assert not jnp.allclose(first_batch_orig, first_batch_shuffled), "Shuffling didn't work"
    
    print("✓ Data batching utilities work correctly")
    return True


def test_batch_size_recommendation():
    """Test the batch size recommendation function."""
    print("\n=== Testing Batch Size Recommendation ===")
    
    # Create test data of varying sizes
    small_data = {"embeddings": jnp.zeros((100, 50))}
    large_data = {"embeddings": jnp.zeros((1000, 2048))}  # Much larger data
    
    small_batch_size = get_batch_size_recommendation(small_data, target_memory_gb=1.0)
    large_batch_size = get_batch_size_recommendation(large_data, target_memory_gb=1.0)
    
    assert isinstance(small_batch_size, int)
    assert isinstance(large_batch_size, int)
    assert small_batch_size >= 16  # Should be at least minimum
    assert large_batch_size >= 16
    
    # For very different data sizes, we should get different recommendations
    # (though both might hit the upper bound)
    print(f"✓ Batch size recommendations: small_data={small_batch_size}, large_data={large_batch_size}")
    return True


def test_backward_compatibility():
    """Test that existing code still works without batching."""
    print("\n=== Testing Backward Compatibility ===")
    
    # Create test data
    rng_key = random.PRNGKey(42)
    n_samples = 50  # Small for fast testing
    embedding_dim = 32
    latent_dim = 3
    
    test_data = {
        "embeddings": random.normal(rng_key, (n_samples, embedding_dim))
    }
    
    modality_specs = {
        "embeddings": {"type": "gaussian", "dim": embedding_dim}
    }
    
    # Create model and fit without specifying batch_size (should default to full-batch)
    model = MultiModalVAE(
        modality_specs=modality_specs,
        latent_dim=latent_dim,
        encoder_hidden_dims=(16,),
        decoder_hidden_dims=(16,)
    )
    
    try:
        model.fit(test_data, num_steps=10, progress_bar=False)
        print("✓ Full-batch training (backward compatibility) works")
    except Exception as e:
        print(f"✗ Full-batch training failed: {e}")
        return False
    
    # Test sampling without batch_size
    try:
        samples = model.sample_posterior(test_data, num_samples=10)
        assert "z" in samples
        print("✓ Full-batch sampling (backward compatibility) works")
    except Exception as e:
        print(f"✗ Full-batch sampling failed: {e}")
        return False
    
    return True


def test_mini_batch_training():
    """Test mini-batch training functionality."""
    print("\n=== Testing Mini-Batch Training ===")
    
    # Create test data
    rng_key = random.PRNGKey(123)
    n_samples = 80
    embedding_dim = 32
    latent_dim = 3
    
    test_data = {
        "embeddings": random.normal(rng_key, (n_samples, embedding_dim))
    }
    
    modality_specs = {
        "embeddings": {"type": "gaussian", "dim": embedding_dim}
    }
    
    # Test with different batch sizes
    for batch_size in [16, 32, 100]:  # Last one should default to full-batch
        try:
            model = MultiModalVAE(
                modality_specs=modality_specs,
                latent_dim=latent_dim,
                encoder_hidden_dims=(16,),
                decoder_hidden_dims=(16,)
            )
            
            model.fit(test_data, batch_size=batch_size, num_steps=5, progress_bar=False, shuffle=True)
            
            # Check that losses are tracked
            assert hasattr(model, 'losses')
            assert len(model.losses) > 0
            
            print(f"✓ Mini-batch training with batch_size={batch_size} works")
            
        except Exception as e:
            print(f"✗ Mini-batch training with batch_size={batch_size} failed: {e}")
            return False
    
    return True


def test_batched_sampling():
    """Test batched posterior sampling."""
    print("\n=== Testing Batched Sampling ===")
    
    # Create and train a model
    rng_key = random.PRNGKey(456)
    n_samples = 60
    embedding_dim = 32
    latent_dim = 3
    
    test_data = {
        "embeddings": random.normal(rng_key, (n_samples, embedding_dim))
    }
    
    modality_specs = {
        "embeddings": {"type": "gaussian", "dim": embedding_dim}
    }
    
    model = MultiModalVAE(
        modality_specs=modality_specs,
        latent_dim=latent_dim,
        encoder_hidden_dims=(16,),
        decoder_hidden_dims=(16,)
    )
    
    model.fit(test_data, num_steps=5, progress_bar=False)
    
    # Test batched sampling
    try:
        samples_full = model.sample_posterior(test_data, num_samples=10)
        samples_batched = model.sample_posterior(test_data, batch_size=20, num_samples=10)
        
        # Check shapes match
        assert samples_full["z"].shape == samples_batched["z"].shape
        print("✓ Batched posterior sampling works")
        
    except Exception as e:
        print(f"✗ Batched posterior sampling failed: {e}")
        return False
    
    return True


def test_batched_generation():
    """Test batched sample generation."""
    print("\n=== Testing Batched Generation ===")
    
    # Create and train a model
    rng_key = random.PRNGKey(789)
    n_samples = 40
    embedding_dim = 32
    latent_dim = 3
    
    test_data = {
        "embeddings": random.normal(rng_key, (n_samples, embedding_dim))
    }
    
    modality_specs = {
        "embeddings": {"type": "gaussian", "dim": embedding_dim}
    }
    
    model = MultiModalVAE(
        modality_specs=modality_specs,
        latent_dim=latent_dim,
        encoder_hidden_dims=(16,),
        decoder_hidden_dims=(16,)
    )
    
    model.fit(test_data, num_steps=5, progress_bar=False)
    
    # Test batched generation
    try:
        gen_full = model.generate_samples(num_samples=50)
        gen_batched = model.generate_samples(num_samples=50, batch_size=15)
        
        # Check shapes match
        assert gen_full["embeddings"].shape == gen_batched["embeddings"].shape
        assert gen_full["embeddings"].shape == (50, embedding_dim)
        print("✓ Batched sample generation works")
        
    except Exception as e:
        print(f"✗ Batched sample generation failed: {e}")
        return False
    
    return True


def test_with_structured_prior():
    """Test batching with structured prior and covariates."""
    print("\n=== Testing Batching with Structured Prior ===")
    
    # Create test data with covariates
    rng_key = random.PRNGKey(999)
    n_samples = 60
    embedding_dim = 32
    latent_dim = 3
    covariate_dim = 4
    
    test_data = {
        "embeddings": random.normal(rng_key, (n_samples, embedding_dim))
    }
    
    covariates = random.normal(rng_key, (n_samples, covariate_dim))
    
    modality_specs = {
        "embeddings": {"type": "gaussian", "dim": embedding_dim}
    }
    
    # Create model with structured prior
    model = MultiModalVAE(
        modality_specs=modality_specs,
        latent_dim=latent_dim,
        encoder_hidden_dims=(16,),
        decoder_hidden_dims=(16,),
        use_structured_prior=True,
        covariate_dim=covariate_dim
    )
    
    try:
        # Test training with batching
        model.fit(test_data, covariates=covariates, batch_size=20, num_steps=5, progress_bar=False)
        
        # Test sampling with batching
        samples = model.sample_posterior(test_data, covariates=covariates, batch_size=25, num_samples=5)
        
        assert "z" in samples
        assert "gamma" in samples  # Should have structured prior parameters
        print("✓ Batching with structured prior works")
        
    except Exception as e:
        print(f"✗ Batching with structured prior failed: {e}")
        return False
    
    return True


def run_all_tests():
    """Run all batching tests."""
    print("=" * 60)
    print("RUNNING BATCHING FUNCTIONALITY TESTS")
    print("=" * 60)
    
    tests = [
        test_data_batching_utilities,
        test_batch_size_recommendation,
        test_backward_compatibility,
        test_mini_batch_training,
        test_batched_sampling,
        test_batched_generation,
        test_with_structured_prior
    ]
    
    passed = 0
    total = len(tests)
    
    for test in tests:
        try:
            if test():
                passed += 1
            else:
                print(f"✗ {test.__name__} failed")
        except Exception as e:
            print(f"✗ {test.__name__} crashed: {e}")
    
    print("\n" + "=" * 60)
    print(f"RESULTS: {passed}/{total} tests passed")
    
    if passed == total:
        print("🎉 All batching tests passed! The implementation is working correctly.")
        print("\nYou can now use batching in your models:")
        print("  - model.fit(data, batch_size=256)")
        print("  - model.sample_posterior(data, batch_size=512)")
        print("  - model.generate_samples(num_samples=1000, batch_size=128)")
        print("  - recommended_batch_size = model.recommend_batch_size(data)")
    else:
        print("❌ Some tests failed. Please check the implementation.")
    
    print("=" * 60)
    
    return passed == total


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