import numpy as np
import cv2
import requests
from PIL import Image
import io
from typing import Dict, List, Tuple, Optional
from dataclasses import dataclass
from sklearn.cluster import KMeans
import matplotlib.pyplot as plt

@dataclass
class ColorDBZMapping:
    """Represents a color to DBZ value mapping."""
    color_rgb: Tuple[int, int, int]
    dbz_value: float
    description: str

class RadarColorScale:
    """Defines standard radar color scales for different countries/organizations."""
    
    # Environment and Climate Change Canada radar color scale
    # Based on official documentation and 14-color palette
    CANADIAN_SCALE = [
        ColorDBZMapping((0, 0, 0, 0), -32, "No precipitation"),  # Transparent
        ColorDBZMapping((102, 102, 102), -20, "Very light"),
        ColorDBZMapping((0, 255, 255), -10, "Light drizzle"),    # Cyan
        ColorDBZMapping((0, 200, 0), 0, "Light rain"),           # Green
        ColorDBZMapping((0, 144, 0), 5, "Light rain"),
        ColorDBZMapping((255, 255, 0), 10, "Light-moderate"),    # Yellow
        ColorDBZMapping((255, 200, 0), 15, "Moderate rain"),     # Orange-yellow
        ColorDBZMapping((255, 144, 0), 20, "Moderate rain"),     # Orange
        ColorDBZMapping((255, 96, 0), 25, "Moderate-heavy"),     # Dark orange
        ColorDBZMapping((255, 0, 0), 30, "Heavy rain"),          # Red
        ColorDBZMapping((215, 0, 0), 35, "Heavy rain"),
        ColorDBZMapping((192, 0, 192), 40, "Very heavy"),        # Magenta
        ColorDBZMapping((148, 0, 211), 50, "Extreme"),           # Dark violet
        ColorDBZMapping((75, 0, 130), 60, "Intense"),            # Indigo
        ColorDBZMapping((255, 255, 255), 70, "Hail/Extreme")     # White
    ]
    
    # US National Weather Service radar color scale (NEXRAD standard)
    # Based on standard NWS/NOAA color scheme
    AMERICAN_SCALE = [
        ColorDBZMapping((0, 0, 0, 0), -32, "No precipitation"),  # Transparent
        ColorDBZMapping((64, 64, 64), -20, "Very light"),
        ColorDBZMapping((30, 144, 255), -10, "Light drizzle"),   # Dodger blue
        ColorDBZMapping((0, 255, 0), 0, "Light rain"),           # Lime
        ColorDBZMapping((0, 200, 0), 5, "Light rain"),           # Green  
        ColorDBZMapping((0, 144, 0), 10, "Light-moderate"),      # Dark green
        ColorDBZMapping((255, 255, 0), 15, "Moderate rain"),     # Yellow
        ColorDBZMapping((229, 255, 0), 20, "Moderate rain"),     # Yellow-green
        ColorDBZMapping((255, 140, 0), 25, "Moderate-heavy"),    # Dark orange
        ColorDBZMapping((255, 0, 0), 30, "Heavy rain"),          # Red
        ColorDBZMapping((255, 0, 255), 35, "Heavy rain"),        # Magenta
        ColorDBZMapping((153, 85, 201), 40, "Very heavy"),       # Medium slate blue
        ColorDBZMapping((99, 0, 99), 50, "Extreme"),             # Dark magenta
        ColorDBZMapping((0, 0, 0), 60, "Intense"),               # Black
        ColorDBZMapping((255, 255, 255), 70, "Hail/Extreme")     # White
    ]

class RadarImageProcessor:
    """Processes radar images for color detection and reclassification."""
    
    def __init__(self):
        self.canadian_scale = RadarColorScale.CANADIAN_SCALE
        self.american_scale = RadarColorScale.AMERICAN_SCALE
        self.color_tolerance = 30  # RGB tolerance for color matching
        
    def fetch_radar_tile(self, wms_url: str, layer: str, bbox: List[float], 
                        width: int = 512, height: int = 512) -> Optional[np.ndarray]:
        """Fetch a radar tile from WMS service."""
        try:
            # Updated parameters for Environment Canada WMS
            params = {
                'SERVICE': 'WMS',
                'VERSION': '1.3.0',
                'REQUEST': 'GetMap',
                'LAYERS': layer,
                'BBOX': f"{bbox[1]},{bbox[0]},{bbox[3]},{bbox[2]}",  # Note: lat,lon order for 1.3.0
                'WIDTH': width,
                'HEIGHT': height,
                'CRS': 'EPSG:4326',  # Updated from 'srs' to 'CRS' for version 1.3.0
                'FORMAT': 'image/png',
                'TRANSPARENT': 'TRUE',
                'STYLES': ''
            }
            
            # Add headers to mimic browser request
            headers = {
                'User-Agent': 'Mozilla/5.0 (Windows NT 10.0; Win64; x64) AppleWebKit/537.36 (KHTML, like Gecko) Chrome/91.0.4472.124 Safari/537.36'
            }
            
            print(f"Fetching radar data from: {wms_url}")
            print(f"Layer: {layer}, Bbox: {bbox}")
            print(f"Request URL: {wms_url}?{'&'.join([f'{k}={v}' for k, v in params.items()])}")
            
            response = requests.get(wms_url, params=params, headers=headers, timeout=30)
            
            print(f"Response status: {response.status_code}")
            print(f"Response headers: {dict(response.headers)}")
            
            if response.status_code != 200:
                print(f"HTTP Error: {response.status_code}")
                print(f"Response content: {response.text[:500]}")
                return None
            
            # Check if response is actually an image
            content_type = response.headers.get('content-type', '')
            if 'image' not in content_type.lower():
                print(f"Unexpected content type: {content_type}")
                print(f"Response content: {response.text[:500]}")
                return None
            
            # Convert to numpy array
            image = Image.open(io.BytesIO(response.content))
            image_array = np.array(image)
            
            print(f"Successfully fetched image: {image_array.shape}")
            return image_array
            
        except requests.exceptions.Timeout:
            print("Timeout error: WMS request took too long")
            return None
        except requests.exceptions.ConnectionError:
            print("Connection error: Could not connect to WMS service")
            return None
        except Exception as e:
            print(f"Error fetching radar tile: {e}")
            import traceback
            traceback.print_exc()
            return None
    
    def create_synthetic_radar_image(self, width: int = 512, height: int = 512) -> np.ndarray:
        """Create a synthetic radar image for testing when WMS fails."""
        print("Creating synthetic radar image for demonstration...")
        
        # Create synthetic radar image with realistic patterns
        synthetic_image = np.zeros((height, width, 4), dtype=np.uint8)
        
        # Add some weather patterns using Canadian colors
        canadian_colors = [
            (0, 255, 255, 255),   # Light drizzle (cyan)
            (0, 200, 0, 255),     # Light rain (green)
            (255, 255, 0, 255),   # Moderate rain (yellow)
            (255, 150, 0, 255),   # Heavy rain (orange)
            (255, 0, 0, 255),     # Very heavy rain (red)
        ]
        
        # Create circular weather patterns
        center_y, center_x = height // 2, width // 2
        
        for i, color in enumerate(canadian_colors):
            # Create concentric circles for different rain intensities
            radius = 50 + i * 30
            y, x = np.ogrid[:height, :width]
            mask = (x - center_x)**2 + (y - center_y)**2 <= radius**2
            
            # Only apply to pixels not already colored (outer rings first)
            alpha_mask = synthetic_image[:, :, 3] == 0
            final_mask = mask & alpha_mask
            
            synthetic_image[final_mask] = color
        
        # Add some scattered precipitation
        np.random.seed(42)  # For reproducible results
        for _ in range(100):
            y = np.random.randint(0, height)
            x = np.random.randint(0, width)
            if synthetic_image[y, x, 3] == 0:  # Only on transparent areas
                color_idx = np.random.randint(0, len(canadian_colors))
                synthetic_image[y, x] = canadian_colors[color_idx]
        
        return synthetic_image
    
    def detect_unique_colors(self, image: np.ndarray, max_colors: int = None) -> List[Tuple[int, int, int]]:
        """Detect unique colors in the radar image using exact pixel values for maximum resolution."""
        # Handle transparency - only process non-transparent pixels
        if image.shape[2] == 4:  # RGBA
            mask = image[:, :, 3] > 0  # Non-transparent pixels
            # Get RGB values of non-transparent pixels
            rgb_pixels = image[mask][:, :3]
        else:  # RGB
            rgb_pixels = image.reshape(-1, 3)
        
        if len(rgb_pixels) == 0:
            return []
        
        # Get unique colors directly (no clustering for max resolution)
        unique_colors = np.unique(rgb_pixels.view(np.dtype((np.void, rgb_pixels.dtype.itemsize * 3))))
        unique_rgb = unique_colors.view(rgb_pixels.dtype).reshape(-1, 3)
        
        print(f"Found {len(unique_rgb)} unique colors in radar image")
        
        # Convert to list of tuples
        return [tuple(color) for color in unique_rgb]
    
    def map_color_to_dbz(self, color: Tuple[int, int, int], 
                        color_scale: List[ColorDBZMapping]) -> float:
        """Map a detected color to its corresponding DBZ value."""
        min_distance = float('inf')
        closest_dbz = -30  # Default to no precipitation
        
        for mapping in color_scale:
            # Calculate Euclidean distance in RGB space
            distance = np.sqrt(sum((c1 - c2) ** 2 for c1, c2 in zip(color, mapping.color_rgb[:3])))
            
            if distance < min_distance:
                min_distance = distance
                closest_dbz = mapping.dbz_value
        
        return closest_dbz
    
    def create_color_mapping(self, detected_colors: List[Tuple[int, int, int]]) -> Dict[Tuple[int, int, int], float]:
        """Create a mapping from detected colors to DBZ values."""
        color_to_dbz = {}
        
        for color in detected_colors:
            dbz_value = self.map_color_to_dbz(color, self.canadian_scale)
            color_to_dbz[color] = dbz_value
        
        return color_to_dbz
    
    def get_american_color_for_dbz(self, dbz_value: float) -> Tuple[int, int, int]:
        """Get the American color scheme color for a given DBZ value."""
        # Find the closest DBZ value in American scale
        min_diff = float('inf')
        closest_color = (0, 0, 0)  # Default to black
        
        for mapping in self.american_scale:
            diff = abs(mapping.dbz_value - dbz_value)
            if diff < min_diff:
                min_diff = diff
                closest_color = mapping.color_rgb[:3]
        
        return closest_color
    
    def reclassify_image(self, image: np.ndarray) -> np.ndarray:
        """Reclassify radar image using color region analysis for WMS compressed data."""
        print(f"Reclassifying WMS compressed image of size: {image.shape}")
        
        # First, analyze the actual colors in the image to understand the data
        if image.shape[2] >= 4:
            mask = image[:, :, 3] > 0  # Non-transparent pixels
            if np.sum(mask) == 0:
                print("No radar data found - all pixels are transparent")
                return image.copy()
            
            non_transparent_pixels = image[mask][:, :3]
            unique_colors = np.unique(non_transparent_pixels.view(np.dtype((np.void, 3))), return_counts=True)
            unique_rgb = unique_colors[0].view(np.uint8).reshape(-1, 3)
            color_counts = unique_colors[1]
            
            # Sort by frequency to see dominant colors
            sorted_indices = np.argsort(color_counts)[::-1]
            top_colors = unique_rgb[sorted_indices][:20]  # Top 20 most common colors
            
            print(f"Top colors in source: {[tuple(c) for c in top_colors[:5]]}")
        
        # Create output image
        output_image = np.zeros_like(image)
        height, width = image.shape[:2]
        pixels_processed = 0
        
        for y in range(height):
            for x in range(width):
                if image.shape[2] == 4:  # RGBA
                    r, g, b, a = image[y, x]
                    if a == 0:  # Skip transparent pixels
                        output_image[y, x] = [0, 0, 0, 0]
                        continue
                else:  # RGB
                    r, g, b = image[y, x]
                    a = 255
                
                # Classify pixel based on color characteristics rather than exact matching
                american_color = self.classify_compressed_radar_pixel(r, g, b)
                
                if american_color is not None:
                    # Set the output pixel to the American color
                    if image.shape[2] == 4:
                        output_image[y, x] = [american_color[0], american_color[1], american_color[2], a]
                    else:
                        output_image[y, x] = american_color
                    
                    pixels_processed += 1
                    
                    # Debug sample mappings (first few pixels only)
                    if pixels_processed <= 5:
                        print(f"Classified: RGB({r},{g},{b}) -> American RGB{american_color}")
                else:
                    # Keep non-radar pixels transparent
                    if image.shape[2] == 4:
                        output_image[y, x] = [0, 0, 0, 0]
                    else:
                        output_image[y, x] = [0, 0, 0]
        
        print(f"Processed {pixels_processed} pixels with American radar colors")
        return output_image
    
    def classify_compressed_radar_pixel(self, r: int, g: int, b: int) -> tuple:
        """Classify a pixel from compressed WMS data into American radar colors."""
        
        # Skip obvious background colors
        if r < 10 and g < 10 and b < 10:  # Black
            return None
        if r > 240 and g > 240 and b > 240:  # White
            return None
        if abs(r-g) < 10 and abs(g-b) < 10 and abs(r-b) < 10 and r > 200:  # Light gray
            return None
            
        # Analyze color characteristics to determine precipitation intensity
        
        # Light blue/cyan region (light precipitation)
        if b > 200 and g > 150 and r < 200:
            return (30, 144, 255)  # American light drizzle blue
        
        # Green region (light to moderate rain)
        if g > r + 50 and g > b + 50:
            if g > 200:
                return (0, 255, 0)    # American lime green (light rain)
            elif g > 150:
                return (0, 200, 0)    # American green (light rain)
            else:
                return (0, 144, 0)    # American dark green (moderate rain)
        
        # Yellow region (moderate rain)
        if r > 200 and g > 200 and b < 100:
            return (255, 255, 0)  # American yellow
        
        # Orange region (moderate-heavy rain)
        if r > 200 and g > 100 and g < 200 and b < 100:
            if g > 150:
                return (229, 255, 0)  # American yellow-green
            else:
                return (255, 140, 0)  # American dark orange
        
        # Red region (heavy rain)
        if r > 200 and g < 100 and b < 100:
            return (255, 0, 0)  # American red
        
        # Purple/magenta region (very heavy/extreme)
        if r > 100 and b > 100 and g < 100:
            if r > 200:
                return (255, 0, 255)     # American magenta
            else:
                return (153, 85, 201)    # American medium slate blue
        
        # If we can't classify it, don't include it
        return None
    
    def estimate_dbz_from_color(self, r: int, g: int, b: int) -> float:
        """Estimate dBZ value from a Canadian radar color using distance-based matching."""
        # Skip if pixel is mostly black/transparent (background)
        if r < 10 and g < 10 and b < 10:
            return -32
        
        # Skip if pixel is mostly white (map background)
        if r > 240 and g > 240 and b > 240:
            return -32
            
        min_distance = float('inf')
        closest_dbz = -32  # Default to no precipitation
        
        for mapping in self.canadian_scale:
            # Calculate Euclidean distance in RGB space with proper type handling
            dr = float(r) - float(mapping.color_rgb[0])
            dg = float(g) - float(mapping.color_rgb[1]) 
            db = float(b) - float(mapping.color_rgb[2])
            distance = np.sqrt(dr*dr + dg*dg + db*db)
            
            if distance < min_distance:
                min_distance = distance
                closest_dbz = mapping.dbz_value
        
        # Much stricter threshold - only very close color matches are considered radar
        if min_distance < 25:  # Stricter threshold for radar color detection
            return closest_dbz
        else:
            return -32  # Not a radar color
    
    def get_precise_american_color(self, dbz_value: float) -> tuple:
        """Get the precise American color for a given dBZ value."""
        # Find the exact matching dBZ in American scale, or closest one
        min_diff = float('inf')
        best_color = (0, 0, 0)  # Default to black
        
        for mapping in self.american_scale:
            diff = abs(mapping.dbz_value - dbz_value)
            if diff < min_diff:
                min_diff = diff
                best_color = mapping.color_rgb[:3]
        
        return best_color
    
    def analyze_color_distribution(self, image: np.ndarray) -> Dict:
        """Analyze the color distribution in the radar image."""
        detected_colors = self.detect_unique_colors(image)
        color_to_dbz = self.create_color_mapping(detected_colors)
        
        analysis = {
            'detected_colors': detected_colors,
            'color_to_dbz': color_to_dbz,
            'num_unique_colors': len(detected_colors),
            'dbz_range': (min(color_to_dbz.values()), max(color_to_dbz.values()))
        }
        
        return analysis
    
    def create_color_legend(self, output_path: str = None) -> plt.Figure:
        """Create a visual comparison of Canadian vs American color scales."""
        fig, (ax1, ax2) = plt.subplots(1, 2, figsize=(12, 8))
        
        # Canadian scale
        canadian_colors = [mapping.color_rgb[:3] for mapping in self.canadian_scale[1:]]  # Skip transparent
        canadian_dbz = [mapping.dbz_value for mapping in self.canadian_scale[1:]]
        canadian_labels = [mapping.description for mapping in self.canadian_scale[1:]]
        
        # American scale
        american_colors = [mapping.color_rgb[:3] for mapping in self.american_scale[1:]]  # Skip transparent
        american_dbz = [mapping.dbz_value for mapping in self.american_scale[1:]]
        american_labels = [mapping.description for mapping in self.american_scale[1:]]
        
        # Normalize colors to 0-1 range for matplotlib
        canadian_colors_norm = [[c/255.0 for c in color] for color in canadian_colors]
        american_colors_norm = [[c/255.0 for c in color] for color in american_colors]
        
        # Plot Canadian scale
        ax1.barh(range(len(canadian_colors)), [1] * len(canadian_colors), 
                color=canadian_colors_norm, edgecolor='black', linewidth=0.5)
        ax1.set_yticks(range(len(canadian_colors)))
        ax1.set_yticklabels([f"{dbz} dBZ" for dbz in canadian_dbz])
        ax1.set_title("Canadian Radar Color Scale", fontsize=14, fontweight='bold')
        ax1.set_xlabel("Color")
        
        # Plot American scale
        ax2.barh(range(len(american_colors)), [1] * len(american_colors), 
                color=american_colors_norm, edgecolor='black', linewidth=0.5)
        ax2.set_yticks(range(len(american_colors)))
        ax2.set_yticklabels([f"{dbz} dBZ" for dbz in american_dbz])
        ax2.set_title("American (NWS) Radar Color Scale", fontsize=14, fontweight='bold')
        ax2.set_xlabel("Color")
        
        plt.tight_layout()
        
        if output_path:
            plt.savefig(output_path, dpi=300, bbox_inches='tight')
        
        return fig

# Example usage and testing
if __name__ == "__main__":
    processor = RadarImageProcessor()
    
    # Create color legend
    fig = processor.create_color_legend("color_scales_comparison.png")
    plt.show()
    
    # Test with a sample WMS request (this would need actual coordinates)
    # bbox = [-75.0, 45.0, -74.0, 46.0]  # Example: Montreal area
    # image = processor.fetch_radar_tile(
    #     "https://geo.weather.gc.ca/geomet", 
    #     "RADAR_1KM_RRAI", 
    #     bbox
    # )
    # 
    # if image is not None:
    #     analysis = processor.analyze_color_distribution(image)
    #     print(f"Detected {analysis['num_unique_colors']} unique colors")
    #     print(f"DBZ range: {analysis['dbz_range']}")
    #     
    #     reclassified = processor.reclassify_image(image)
    #     # Save or display results...