#!/usr/bin/env python3
"""
Test script for radar reclassification functionality.
"""

import numpy as np
from radar_processor import RadarImageProcessor, RadarColorScale
import matplotlib.pyplot as plt

def test_color_detection():
    """Test the color detection functionality with a synthetic radar image."""
    print("Testing color detection...")
    
    # Create a synthetic radar image with known colors
    height, width = 100, 100
    synthetic_image = np.zeros((height, width, 4), dtype=np.uint8)
    
    # Add some test colors from the Canadian scale
    canadian_colors = [
        (0, 255, 255),   # Light drizzle
        (0, 200, 0),     # Light rain
        (255, 255, 0),   # Moderate rain
        (255, 0, 0),     # Very heavy rain
    ]
    
    # Fill quadrants with different colors
    synthetic_image[0:50, 0:50] = [*canadian_colors[0], 255]  # Top-left
    synthetic_image[0:50, 50:100] = [*canadian_colors[1], 255]  # Top-right
    synthetic_image[50:100, 0:50] = [*canadian_colors[2], 255]  # Bottom-left
    synthetic_image[50:100, 50:100] = [*canadian_colors[3], 255]  # Bottom-right
    
    processor = RadarImageProcessor()
    
    # Test color detection
    detected_colors = processor.detect_unique_colors(synthetic_image, max_colors=10)
    print(f"Detected {len(detected_colors)} unique colors:")
    for i, color in enumerate(detected_colors):
        print(f"  Color {i+1}: RGB{color}")
    
    # Test color to DBZ mapping
    color_to_dbz = processor.create_color_mapping(detected_colors)
    print("\nColor to DBZ mappings:")
    for color, dbz in color_to_dbz.items():
        print(f"  RGB{color} → {dbz:.1f} dBZ")
    
    # Test reclassification
    reclassified = processor.reclassify_image(synthetic_image)
    print(f"\nReclassified image shape: {reclassified.shape}")
    
    return True

def test_color_scales():
    """Test the predefined color scales."""
    print("\nTesting color scales...")
    
    canadian_scale = RadarColorScale.CANADIAN_SCALE
    american_scale = RadarColorScale.AMERICAN_SCALE
    
    print(f"Canadian scale has {len(canadian_scale)} color levels")
    print(f"American scale has {len(american_scale)} color levels")
    
    # Check DBZ ranges
    canadian_dbz = [mapping.dbz_value for mapping in canadian_scale]
    american_dbz = [mapping.dbz_value for mapping in american_scale]
    
    print(f"Canadian DBZ range: {min(canadian_dbz)} to {max(canadian_dbz)}")
    print(f"American DBZ range: {min(american_dbz)} to {max(american_dbz)}")
    
    return True

def test_color_legend():
    """Test the color legend generation."""
    print("\nTesting color legend generation...")
    
    processor = RadarImageProcessor()
    
    try:
        fig = processor.create_color_legend("test_color_legend.png")
        print("Color legend generated successfully: test_color_legend.png")
        plt.close(fig)
        return True
    except Exception as e:
        print(f"Error generating color legend: {e}")
        return False

def test_dbz_mapping():
    """Test DBZ to color mapping accuracy."""
    print("\nTesting DBZ mapping accuracy...")
    
    processor = RadarImageProcessor()
    
    # Test specific DBZ values
    test_dbz_values = [-20, -10, 0, 10, 20, 30, 40, 50, 60, 70]
    
    print("DBZ → American Color mappings:")
    for dbz in test_dbz_values:
        color = processor.get_american_color_for_dbz(dbz)
        print(f"  {dbz:3d} dBZ → RGB{color}")
    
    return True

def run_all_tests():
    """Run all tests."""
    print("=" * 60)
    print("RADAR RECLASSIFICATION SYSTEM TESTS")
    print("=" * 60)
    
    tests = [
        ("Color Detection", test_color_detection),
        ("Color Scales", test_color_scales),
        ("Color Legend", test_color_legend),
        ("DBZ Mapping", test_dbz_mapping),
    ]
    
    results = []
    for test_name, test_func in tests:
        try:
            result = test_func()
            results.append((test_name, result))
            status = "PASS" if result else "FAIL"
            print(f"\n[{status}] {test_name}")
        except Exception as e:
            results.append((test_name, False))
            print(f"\n[ERROR] {test_name}: {e}")
    
    print("\n" + "=" * 60)
    print("TEST SUMMARY")
    print("=" * 60)
    
    passed = sum(1 for _, result in results if result)
    total = len(results)
    
    for test_name, result in results:
        status = "✅ PASS" if result else "❌ FAIL"
        print(f"{status} {test_name}")
    
    print(f"\nTotal: {passed}/{total} tests passed")
    
    if passed == total:
        print("\n🎉 All tests passed! The radar reclassification system is ready.")
    else:
        print(f"\n⚠️  {total - passed} test(s) failed. Please check the implementation.")
    
    return passed == total

if __name__ == "__main__":
    run_all_tests()