#!/usr/bin/env python3
"""
Test script for memory-efficient batch loader.
"""

import sys
from pathlib import Path

# Add src directory to path for imports
sys.path.append(str(Path(__file__).parent / "src"))

from src.batch_loader import MemoryEfficientBatchLoader, BatchProcessor
from src.logger_config import setup_logging

def test_batch_loader():
    """Test the batch loader with sample functionality."""
    
    # Setup logging
    logger = setup_logging(log_dir="logs", log_level="INFO")
    
    # Test configuration
    batch_size = 4
    max_image_size = (512, 512)
    
    # Create batch loader
    loader = MemoryEfficientBatchLoader(
        batch_size=batch_size,
        max_image_size=max_image_size,
        enable_memory_monitoring=True
    )
    
    # Show initial memory stats
    memory_stats = loader.get_memory_stats()
    logger.info(f"Initial memory stats: {memory_stats}")
    
    # Sample image paths (for testing - these would be real image paths)
    sample_image_paths = [
        f"/home/nguyendc/thanh-dev/IQA-Metric-Benchmark/data/task/cni/images/image_{i}.png" 
        for i in range(10)
    ]
    
    # Test batch creation
    batches = list(loader.create_batches(sample_image_paths))
    logger.info(f"Created {len(batches)} batches from {len(sample_image_paths)} images")
    for i, batch in enumerate(batches):
        logger.info(f"Batch {i+1}: {len(batch)} images")
    
    # Mock scoring function for testing
    def mock_scoring_function(images):
        """Mock function that returns random scores."""
        import random
        return [random.uniform(1.0, 5.0) for _ in images]
    
    # Test processing a single batch (with mock data)
    if batches:
        logger.info("Testing batch processing with mock data...")
        # Note: This will fail with actual file paths that don't exist
        # but demonstrates the interface
    
    logger.info("Batch loader test completed!")

if __name__ == "__main__":
    test_batch_loader()