test_dataset_refactor.py

"""Test script to verify the refactored dataset class works correctly."""

import numpy as np
import pandas as pd

# Test 1: Verify imports work
try:
    from latent_price_dataset import LatentForecastDataset, LatentPriceDataset
    print("✅ Test 1 PASSED: Successfully imported LatentForecastDataset and LatentPriceDataset")
except ImportError as e:
    print(f"❌ Test 1 FAILED: {e}")
    exit(1)

# Test 2: Verify backward compatibility class exists
try:
    assert hasattr(LatentPriceDataset, '__init__')
    print("✅ Test 2 PASSED: LatentPriceDataset (backward compatibility) exists")
except AssertionError:
    print("❌ Test 2 FAILED: LatentPriceDataset missing")
    exit(1)

# Test 3: Verify class signature
try:
    import inspect
    sig = inspect.signature(LatentForecastDataset.__init__)
    params = list(sig.parameters.keys())
    
    required_params = ['self', 'latent_zarr_path', 'entsoe_csv_path', 'target_variable', 
                      'start_date', 'end_date']
    for param in required_params:
        assert param in params, f"Missing parameter: {param}"
    
    print("✅ Test 3 PASSED: LatentForecastDataset has correct signature")
    print(f"   Parameters: {params}")
except AssertionError as e:
    print(f"❌ Test 3 FAILED: {e}")
    exit(1)

# Test 4: Verify target_variable handling logic
try:
    # Simulate the logic from __init__
    
    # Test single string target
    target_variable = "price_eur_mwh"
    if isinstance(target_variable, str):
        target_variables = [target_variable]
        single_target = True
    else:
        target_variables = list(target_variable)
        single_target = False
    
    assert target_variables == ["price_eur_mwh"]
    assert single_target == True
    
    # Test list of targets
    target_variable = ["price_eur_mwh", "load_actual_mw"]
    if isinstance(target_variable, str):
        target_variables = [target_variable]
        single_target = True
    else:
        target_variables = list(target_variable)
        single_target = False
    
    assert target_variables == ["price_eur_mwh", "load_actual_mw"]
    assert single_target == False
    
    print("✅ Test 4 PASSED: Target variable handling logic works correctly")
except AssertionError as e:
    print(f"❌ Test 4 FAILED: {e}")
    exit(1)

# Test 5: Verify normalization logic
try:
    # Simulate target data
    target_data = np.random.randn(48, 2) * 100 + 50  # (timesteps, num_targets)
    target_variables = ["price_eur_mwh", "load_actual_mw"]
    
    # Simulate target stats
    target_stats = {
        "price_eur_mwh": {"mean": 50.0, "std": 100.0},
        "load_actual_mw": {"mean": 50.0, "std": 100.0}
    }
    
    # Normalize
    normalized_targets = []
    for i, var in enumerate(target_variables):
        mean = target_stats[var]["mean"]
        std = target_stats[var]["std"]
        normalized_targets.append((target_data[:, i] - mean) / std)
    
    normalized_data = np.column_stack(normalized_targets)
    
    assert normalized_data.shape == (48, 2)
    # Check that mean is approximately 0 and std is approximately 1
    assert abs(normalized_data.mean()) < 0.5, f"Mean should be ~0, got {normalized_data.mean()}"
    
    print("✅ Test 5 PASSED: Multi-target normalization logic works correctly")
except AssertionError as e:
    print(f"❌ Test 5 FAILED: {e}")
    exit(1)

# Test 6: Verify single target squeeze logic
try:
    # Single target case
    target_data = np.random.randn(48, 1)  # (timesteps, 1)
    single_target = True
    
    # Simulate the squeeze logic
    if single_target and target_data.ndim == 2:
        target_data_squeezed = target_data.squeeze(-1)
    else:
        target_data_squeezed = target_data
    
    assert target_data_squeezed.shape == (48,), f"Expected (48,), got {target_data_squeezed.shape}"
    
    # Multi-target case
    target_data = np.random.randn(48, 3)  # (timesteps, 3)
    single_target = False
    
    if single_target and target_data.ndim == 2:
        target_data_squeezed = target_data.squeeze(-1)
    else:
        target_data_squeezed = target_data
    
    assert target_data_squeezed.shape == (48, 3), f"Expected (48, 3), got {target_data_squeezed.shape}"
    
    print("✅ Test 6 PASSED: Single vs multi-target output shape handling works correctly")
except AssertionError as e:
    print(f"❌ Test 6 FAILED: {e}")
    exit(1)

print("\n" + "=" * 60)
print("🎉 ALL TESTS PASSED!")
print("=" * 60)
print("\n✨ The refactored dataset class is ready to use!")
print("\nUsage examples:")
print("  # Single target:")
print('  dataset = LatentForecastDataset(..., target_variable="price_eur_mwh")')
print("\n  # Multiple targets:")
print('  dataset = LatentForecastDataset(..., target_variable=["price_eur_mwh", "load_actual_mw"])')
print("\n  # Backward compatible:")
print('  dataset = LatentPriceDataset(...)  # Still works, but deprecated')