"""
Test script for cloud_removal module
Kiểm tra các phương pháp xử lý mây
"""

import numpy as np
import xarray as xr
from cloud_removal import (
    process_cloud_removal,
    get_available_methods,
    compare_methods
)


def create_mock_s2_data():
    """Tạo mock Sentinel-2 data để test"""
    # Create synthetic data: 5 time steps, 100x100 pixels
    np.random.seed(42)
    
    time_steps = 5
    y_size = 100
    x_size = 100
    
    # Create bands
    bands = {}
    for band in ["B02", "B03", "B04", "B08", "B11"]:
        # Random reflectance values
        data = np.random.rand(time_steps, y_size, x_size) * 0.3 + 0.1
        bands[band] = (["time", "y", "x"], data)
    
    # Create SCL (Scene Classification Layer)
    # Mostly vegetation (4), with some clouds
    scl_data = np.full((time_steps, y_size, x_size), 4, dtype=np.uint8)
    
    # Add clouds (class 9) in random locations
    for t in range(time_steps):
        # Random cloud patches
        n_clouds = np.random.randint(5, 15)
        for _ in range(n_clouds):
            y_start = np.random.randint(0, y_size - 20)
            x_start = np.random.randint(0, x_size - 20)
            cloud_height = np.random.randint(10, 20)
            cloud_width = np.random.randint(10, 20)
            scl_data[t, y_start:y_start+cloud_height, x_start:x_start+cloud_width] = 9
    
    bands["SCL"] = (["time", "y", "x"], scl_data)
    
    # Create xarray Dataset
    ds = xr.Dataset(
        bands,
        coords={
            "time": np.arange(time_steps),
            "y": np.arange(y_size),
            "x": np.arange(x_size)
        }
    )
    
    return ds


def test_available_methods():
    """Test lấy danh sách methods"""
    print("=" * 60)
    print("TEST: Get Available Methods")
    print("=" * 60)
    
    methods = get_available_methods()
    print(f"\nFound {len(methods)} methods:")
    for method, description in methods.items():
        print(f"  - {method:20s}: {description}")
    
    print("\n✅ Test passed!")


def test_single_method(method_name="classic"):
    """Test một method cụ thể"""
    print("\n" + "=" * 60)
    print(f"TEST: Cloud Removal Method '{method_name}'")
    print("=" * 60)
    
    # Create mock data
    s2_data = create_mock_s2_data()
    print(f"\nMock data created: {dict(s2_data.dims)}")
    
    # Process clouds
    cleaned_data, metadata = process_cloud_removal(
        s2_data=s2_data,
        method=method_name,
        verbose=True
    )
    
    # Check results
    print(f"\nMetadata:")
    print(f"  - Method: {metadata['method']}")
    print(f"  - Cloud coverage: {metadata['cloud_coverage_percent']:.1f}%")
    print(f"  - Masked pixels: {metadata['masked_pixels']:,}/{metadata['total_pixels']:,}")
    print(f"  - Steps applied: {', '.join(metadata['steps_applied'])}")
    
    # Verify no NaN remaining
    nan_count = 0
    for band in cleaned_data.data_vars:
        if band != "SCL":
            nan_count += np.isnan(cleaned_data[band].values).sum()
    
    print(f"\nRemaining NaN pixels: {nan_count}")
    
    if nan_count == 0:
        print("✅ Test passed - no NaN remaining!")
    else:
        print(f"⚠️  Warning - {nan_count} NaN pixels remaining")


def test_comparison():
    """Test so sánh nhiều methods"""
    print("\n" + "=" * 60)
    print("TEST: Compare Multiple Methods")
    print("=" * 60)
    
    # Create mock data
    s2_data = create_mock_s2_data()
    
    # Compare methods
    methods_to_test = ["classic", "temporal_only", "median_composite", "ml_knn"]
    
    print(f"\nComparing {len(methods_to_test)} methods...")
    results = compare_methods(s2_data, methods=methods_to_test)
    
    # Print summary
    print("\n" + "-" * 60)
    print(f"{'Method':<20} {'Success':<10} {'NaN %':<10} {'Steps'}")
    print("-" * 60)
    
    for method, result in results.items():
        if result['success']:
            nan_pct = result['remaining_nan_percent']
            steps = ', '.join(result['metadata']['steps_applied'][:2])  # First 2 steps
            print(f"{method:<20} {'✅':<10} {nan_pct:>6.2f}%   {steps}")
        else:
            print(f"{method:<20} {'❌':<10} {'ERROR':<10} {result['error']}")
    
    print("-" * 60)
    print("\n✅ Comparison test completed!")


def test_edge_cases():
    """Test các trường hợp đặc biệt"""
    print("\n" + "=" * 60)
    print("TEST: Edge Cases")
    print("=" * 60)
    
    # Case 1: No SCL band
    print("\n1. Testing without SCL band...")
    s2_data = create_mock_s2_data()
    s2_data_no_scl = s2_data.drop_vars("SCL")
    
    cleaned, meta = process_cloud_removal(s2_data_no_scl, method="classic", verbose=False)
    print(f"   Result: {meta.get('warning', 'OK')}")
    
    # Case 2: 100% cloud coverage
    print("\n2. Testing with 100% cloud coverage...")
    s2_data_full_cloud = create_mock_s2_data()
    s2_data_full_cloud["SCL"][:] = 9  # All clouds
    
    cleaned, meta = process_cloud_removal(s2_data_full_cloud, method="classic", verbose=False)
    print(f"   Cloud coverage: {meta['cloud_coverage_percent']:.1f}%")
    
    # Case 3: No clouds
    print("\n3. Testing with no clouds...")
    s2_data_clear = create_mock_s2_data()
    s2_data_clear["SCL"][:] = 4  # All vegetation
    
    cleaned, meta = process_cloud_removal(s2_data_clear, method="classic", verbose=False)
    print(f"   Cloud coverage: {meta['cloud_coverage_percent']:.1f}%")
    
    print("\n✅ Edge case tests passed!")


if __name__ == "__main__":
    print("\n" + "🌥️  CLOUD REMOVAL MODULE TESTS 🌥️ ".center(60, "="))
    print()
    
    # Run tests
    test_available_methods()
    test_single_method("classic")
    test_single_method("hybrid")
    test_comparison()
    test_edge_cases()
    
    print("\n" + "=" * 60)
    print("ALL TESTS COMPLETED!")
    print("=" * 60)
    print("\nModule is ready to use. Available methods:")
    for method, desc in get_available_methods().items():
        print(f"  • {method}")