"""
ImagePreprocessor.py - FULLY OPTIMIZED
Minimal overhead, fast operations, no unnecessary prints.
"""

import cv2
import numpy as np
from PIL import Image
from typing import Tuple, Optional
from sklearn.cluster import MiniBatchKMeans
import warnings


class ImagePreprocessor:
    """OPTIMIZED image preprocessing for mosaic generation."""
    
    def __init__(self, target_resolution: Tuple[int, int] = (800, 600), 
                 grid_size: Tuple[int, int] = (20, 15),
                 verbose: bool = False):
        """Initialize preprocessor - MINIMAL overhead."""
        self.target_resolution = target_resolution
        self.grid_size = grid_size
        self.verbose = verbose
        
        # Pre-compute (avoid division later)
        self.tile_width = target_resolution[0] // grid_size[0]
        self.tile_height = target_resolution[1] // grid_size[1]
        self.adjusted_width = self.tile_width * grid_size[0]
        self.adjusted_height = self.tile_height * grid_size[1]
    
    def load_and_preprocess_image(self, image_path: str, 
                                 apply_quantization: bool = False,
                                 n_colors: int = 16) -> Optional[np.ndarray]:
        """Load and preprocess - OPTIMIZED for speed."""
        try:
            # Load and convert
            image = cv2.imread(image_path)
            if image is None:
                raise ValueError(f"Could not load image: {image_path}")
            
            image = cv2.cvtColor(image, cv2.COLOR_BGR2RGB)
            processed_image = self._resize_and_crop(image)
            
            # Only quantize if explicitly requested
            if apply_quantization and n_colors < 256:
                processed_image = self._apply_color_quantization(processed_image, n_colors)
            
            return processed_image
            
        except Exception as e:
            if self.verbose:
                print(f"Error processing {image_path}: {str(e)}")
            return None
    
    def preprocess_numpy_image(self, image: np.ndarray,
                              apply_quantization: bool = False,
                              n_colors: int = 16) -> Optional[np.ndarray]:
        """Preprocess numpy image (for Gradio)."""
        try:
            if len(image.shape) != 3 or image.shape[2] != 3:
                raise ValueError("Image must be RGB (H, W, 3)")
            
            processed_image = self._resize_and_crop(image)
            
            if apply_quantization and n_colors < 256:
                processed_image = self._apply_color_quantization(processed_image, n_colors)
            
            return processed_image
            
        except Exception as e:
            if self.verbose:
                print(f"Error processing image: {str(e)}")
            return None
    
    def _resize_and_crop(self, image: np.ndarray) -> np.ndarray:
        """Resize and crop - OPTIMIZED."""
        h, w = image.shape[:2]
        target_w, target_h = self.adjusted_width, self.adjusted_height
        
        # Scale to fill
        scale = max(target_w / w, target_h / h)
        new_w = int(w * scale)
        new_h = int(h * scale)
        
        # Use INTER_LINEAR for speed (faster than INTER_AREA)
        resized = cv2.resize(image, (new_w, new_h), interpolation=cv2.INTER_LINEAR)
        
        # Center crop
        start_x = (new_w - target_w) // 2
        start_y = (new_h - target_h) // 2
        cropped = resized[start_y:start_y + target_h, start_x:start_x + target_w]
        
        return cropped
    
    def _apply_color_quantization(self, image: np.ndarray, n_colors: int) -> np.ndarray:
        """OPTIMIZED color quantization with aggressive sampling."""
        h, w, c = image.shape
        pixels = image.reshape(-1, c)
        total_pixels = len(pixels)
        
        # Aggressive sampling
        max_sample_size = 10000
        
        if total_pixels > max_sample_size:
            # Sample for fitting
            sample_indices = np.random.randint(0, total_pixels, size=max_sample_size)
            sampled_pixels = pixels[sample_indices].astype(np.float32)
            
            batch_size = min(max(len(sampled_pixels) // 10, 500), 2000)
            
            with warnings.catch_warnings():
                warnings.filterwarnings("ignore")
                
                kmeans = MiniBatchKMeans(
                    n_clusters=n_colors, batch_size=batch_size,
                    random_state=42, n_init=1, max_iter=50
                )
                
                # Fit on sample, predict on full
                kmeans.fit(sampled_pixels)
                labels = kmeans.predict(pixels.astype(np.float32))
        else:
            # Small image
            batch_size = min(max(total_pixels // 100, 500), 5000)
            
            with warnings.catch_warnings():
                warnings.filterwarnings("ignore")
                
                kmeans = MiniBatchKMeans(
                    n_clusters=n_colors, batch_size=batch_size,
                    random_state=42, n_init=1, max_iter=50
                )
                labels = kmeans.fit_predict(pixels.astype(np.float32))
        
        # Apply quantization
        quantized_pixels = kmeans.cluster_centers_[labels]
        quantized_image = quantized_pixels.reshape(h, w, c).astype(np.uint8)
        
        return quantized_image
    
    def save_preprocessed_image(self, image: np.ndarray, output_path: str):
        """Save preprocessed image."""
        try:
            pil_image = Image.fromarray(image)
            pil_image.save(output_path, quality=95)
        except Exception as e:
            if self.verbose:
                print(f"Error saving image: {str(e)}")