"""
Raster Processor

This module handles raster data loading, tiling, and preprocessing for inference.
It uses GDAL to read geospatial rasters and prepares tiles with padding and
normalization for PyTorch models.

Key features:
- Sliding window tiling with configurable overlap
- Padding to reduce edge artifacts
- Normalization (ImageNet or custom)
- Memory-efficient tile generation (yields one tile at a time)
- Geospatial metadata preservation
"""

import numpy as np
import torch
from osgeo import gdal
from qgis.core import QgsRasterLayer, QgsCoordinateTransform, QgsProject, QgsRectangle
from typing import Dict, Any, Generator, Tuple, Optional


class RasterProcessor:
    """
    Process raster data for deep learning inference.
    
    This class handles:
    - Raster reading with GDAL
    - Sliding window tiling with overlap
    - Padding addition and handling
    - Pixel normalization for neural networks
    - Geospatial coordinate tracking
    
    Args:
        raster_layer: QGIS raster layer to process
        tile_size: Size of tiles for inference (without padding)
        overlap: Overlap between adjacent tiles in pixels
        padding: Padding to add around each tile
        normalization: Dict with 'mean' and 'std' lists for normalization
        extent_mode: Processing extent ('Full Raster', 'Current Viewport', etc.)
        
    Example:
        >>> processor = RasterProcessor(
        ...     raster_layer=layer,
        ...     tile_size=512,
        ...     overlap=256,
        ...     padding=64,
        ...     normalization={'mean': [0.485, 0.456, 0.406], 'std': [0.229, 0.224, 0.225]}
        ... )
        >>> for tile in processor.generate_tiles():
        ...     prediction = model(tile['tensor'])
    """
    
    def __init__(
        self,
        raster_layer: QgsRasterLayer,
        tile_size: int,
        overlap: int,
        padding: int,
        normalization: Dict[str, list],
        extent_mode: str = 'Full Raster',
        selection_extent: Optional[Any] = None,
        divisibility: int = 1
    ):
        """
        Initialize raster processor.
        
        Args:
            raster_layer: QGIS raster layer
            tile_size: Tile size (without padding)
            overlap: Overlap between tiles
            padding: Padding around each tile
            normalization: Normalization parameters
            extent_mode: Extent to process
        """
        self.raster_layer = raster_layer
        self.tile_size = tile_size
        self.overlap = overlap
        self.padding = padding
        self.normalization = normalization
        self.extent_mode = extent_mode
        self.selection_extent = selection_extent
        self.divisibility = divisibility
        
        # Open raster with GDAL
        self.dataset = gdal.Open(raster_layer.source(), gdal.GA_ReadOnly)
        if self.dataset is None:
            raise ValueError(f"Could not open raster: {raster_layer.source()}")
        
        # Get raster dimensions
        self.width = self.dataset.RasterXSize
        self.height = self.dataset.RasterYSize
        self.num_bands = self.dataset.RasterCount
        
        # Get geotransform for coordinate calculations
        self.geo_transform = self.dataset.GetGeoTransform()
        self.projection = self.dataset.GetProjection()
        
        # Calculate stride (distance between tile origins)
        self.stride = tile_size - overlap
        
        # Calculate processing extent
        self._calculate_extent()
        
        # Calculate tile grid
        self._calculate_tile_grid()
    
    def _calculate_extent(self):
        """
        Calculate the extent to process based on extent_mode.
        
        Sets:
            self.extent_x_min, self.extent_y_min: Top-left corner in pixels
            self.extent_width, self.extent_height: Extent dimensions in pixels
        """
        if self.extent_mode == 'Full Raster':
            # Process entire raster
            self.extent_x_min = 0
            self.extent_y_min = 0
            self.extent_width = self.width
            self.extent_height = self.height
        
        elif self.extent_mode == 'Current Viewport':
            # Get map canvas extent and transform to raster pixel coordinates
            from qgis.utils import iface
            
            # Get canvas extent in map coordinates
            canvas = iface.mapCanvas()
            map_extent = canvas.extent()  # QgsRectangle in canvas CRS
            
            # Get canvas CRS and raster CRS
            canvas_crs = canvas.mapSettings().destinationCrs()
            raster_crs = self.raster_layer.crs()
            
            # Transform extent to raster CRS if needed
            if canvas_crs != raster_crs:
                transform = QgsCoordinateTransform(canvas_crs, raster_crs, QgsProject.instance())
                map_extent = transform.transformBoundingBox(map_extent)
            
            # Convert map coordinates to raster pixel coordinates
            # Using inverse geotransform: pixel = (geo - origin) / pixel_size
            gt = self.geo_transform
            
            # Calculate pixel bounds
            pixel_x_min = int((map_extent.xMinimum() - gt[0]) / gt[1])
            pixel_x_max = int((map_extent.xMaximum() - gt[0]) / gt[1])
            pixel_y_min = int((map_extent.yMinimum() - gt[3]) / gt[5])
            pixel_y_max = int((map_extent.yMaximum() - gt[3]) / gt[5])
            
            # Ensure correct order (y is inverted in raster coordinates)
            if pixel_y_min > pixel_y_max:
                pixel_y_min, pixel_y_max = pixel_y_max, pixel_y_min
            
            # Clip to raster bounds
            pixel_x_min = max(0, pixel_x_min)
            pixel_y_min = max(0, pixel_y_min)
            pixel_x_max = min(self.width, pixel_x_max)
            pixel_y_max = min(self.height, pixel_y_max)
            
            # Set extent
            self.extent_x_min = pixel_x_min
            self.extent_y_min = pixel_y_min
            self.extent_width = pixel_x_max - pixel_x_min
            self.extent_height = pixel_y_max - pixel_y_min
            
            # Validate extent
            if self.extent_width <= 0 or self.extent_height <= 0:
                raise ValueError(
                    f"Invalid viewport extent: no overlap with raster. "
                    f"Viewport may be outside raster bounds."
                )
        
        else:  # 'Selection (Draw)'
            if not self.selection_extent:
                raise ValueError("Selection extent not provided for 'Selection (Draw)' mode")
            
            # Use provided selection extent (QgsRectangle in map coordinates)
            map_extent = self.selection_extent
            
            # Get raster CRS
            raster_crs = self.raster_layer.crs()
            
            # Get extent CRS (assume it's in the same CRS as the canvas/project)
            from qgis.utils import iface
            canvas_crs = iface.mapCanvas().mapSettings().destinationCrs()
            
            # Transform extent to raster CRS if needed
            if canvas_crs != raster_crs:
                transform = QgsCoordinateTransform(canvas_crs, raster_crs, QgsProject.instance())
                map_extent = transform.transformBoundingBox(map_extent)
            
            # Convert map coordinates to raster pixel coordinates
            gt = self.geo_transform
            
            pixel_x_min = int((map_extent.xMinimum() - gt[0]) / gt[1])
            pixel_x_max = int((map_extent.xMaximum() - gt[0]) / gt[1])
            pixel_y_min = int((map_extent.yMinimum() - gt[3]) / gt[5])
            pixel_y_max = int((map_extent.yMaximum() - gt[3]) / gt[5])
            
            # Ensure correct order (y is inverted)
            if pixel_y_min > pixel_y_max:
                pixel_y_min, pixel_y_max = pixel_y_max, pixel_y_min

            
            # Clip to raster bounds
            pixel_x_min = max(0, pixel_x_min)
            pixel_y_min = max(0, pixel_y_min)
            pixel_x_max = min(self.width, pixel_x_max)
            pixel_y_max = min(self.height, pixel_y_max)

            
            
            # Set extent
            self.extent_x_min = pixel_x_min
            self.extent_y_min = pixel_y_min
            self.extent_width = pixel_x_max - pixel_x_min
            self.extent_height = pixel_y_max - pixel_y_min
            
            # Validate extent
            if self.extent_width <= 0 or self.extent_height <= 0:
                raise ValueError(
                    f"Invalid selection extent: no overlap with raster. "
                    f"Selection may be outside raster bounds."
                )
    
    def _calculate_tile_grid(self):
        """
        Calculate the grid of tiles to generate.
        
        Sets:
            self.num_tiles_x: Number of tiles horizontally
            self.num_tiles_y: Number of tiles vertically
            self.total_tiles: Total number of tiles
        """
        # Calculate number of tiles needed to cover extent with stride
        self.num_tiles_x = int(np.ceil((self.extent_width - self.tile_size) / self.stride)) + 1
        self.num_tiles_y = int(np.ceil((self.extent_height - self.tile_size) / self.stride)) + 1
        
        # Handle edge case: if extent is smaller than tile_size
        if self.extent_width <= self.tile_size:
            self.num_tiles_x = 1
        if self.extent_height <= self.tile_size:
            self.num_tiles_y = 1
        
        self.total_tiles = self.num_tiles_x * self.num_tiles_y
    
    def get_num_tiles(self) -> int:
        """
        Get total number of tiles to process.
        
        Returns:
            Total number of tiles
        """
        return self.total_tiles
    
    def generate_tiles(self) -> Generator[Dict[str, Any], None, None]:
        """
        Generate tiles one at a time (memory efficient).
        
        Yields:
            Dictionary with tile information:
            {
                'tensor': torch.Tensor,  # Shape (C, H+2*padding, W+2*padding), normalized
                'row': int,              # Row index in grid
                'col': int,              # Column index in grid
                'x_offset': int,         # X position in raster (without padding)
                'y_offset': int,         # Y position in raster (without padding)
                'tile_width': int,       # Actual tile width (may be smaller at edges)
                'tile_height': int,      # Actual tile height (may be smaller at edges)
                'geo_transform': tuple   # Geotransform for this tile
            }
            
        Example:
            >>> for tile_data in processor.generate_tiles():
            ...     tensor = tile_data['tensor']
            ...     prediction = model(tensor.unsqueeze(0))  # Add batch dimension
        """
        for row in range(self.num_tiles_y):
            for col in range(self.num_tiles_x):
                # Calculate tile position in raster (without padding)
                x_offset = self.extent_x_min + col * self.stride
                y_offset = self.extent_y_min + row * self.stride
                
                # Calculate actual tile size (may be smaller at edges)
                tile_width = min(self.tile_size, self.extent_width - col * self.stride)
                tile_height = min(self.tile_size, self.extent_height - row * self.stride)
                
                # Calculate read window with padding
                x_read = max(0, x_offset - self.padding)
                y_read = max(0, y_offset - self.padding)
                
                # Calculate read dimensions
                width_read = min(
                    tile_width + 2 * self.padding,
                    self.width - x_read
                )
                height_read = min(
                    tile_height + 2 * self.padding,
                    self.height - y_read
                )
                
                # Read data from raster
                tile_array = self._read_tile(x_read, y_read, width_read, height_read)
                
                # Handle padding at edges (reflect mode)
                tile_array, actual_pad_left, actual_pad_top = self._apply_edge_padding(
                    tile_array,
                    tile_width, tile_height,
                    x_offset, y_offset,
                    x_read, y_read
                )
                
                # Apply divisibility padding if required by architecture
                tile_array, extra_pad_left, extra_pad_top, extra_pad_right, extra_pad_bottom = self._pad_to_divisibility(
                    tile_array,
                    tile_width,
                    tile_height
                )
                
                # Update actual padding values to include divisibility padding
                total_pad_left = actual_pad_left + extra_pad_left
                total_pad_top = actual_pad_top + extra_pad_top
                
                # Normalize
                tile_tensor = self._normalize_tile(tile_array)
                
                # Calculate geotransform for this tile
                tile_geo_transform = (
                    self.geo_transform[0] + x_offset * self.geo_transform[1], 
                    self.geo_transform[1],  
                    self.geo_transform[2],  
                    self.geo_transform[3] + y_offset * self.geo_transform[5],  
                    self.geo_transform[4],  
                    self.geo_transform[5]  
                )
                
                yield {
                    'tensor': tile_tensor,
                    'row': row,
                    'col': col,
                    'x_offset': x_offset,
                    'y_offset': y_offset,
                    'tile_width': tile_width,
                    'tile_height': tile_height,
                    'geo_transform': tile_geo_transform,
                    'pad_left': total_pad_left,
                    'pad_top': total_pad_top,
                    'extra_pad_right': extra_pad_right,
                    'extra_pad_bottom': extra_pad_bottom
                }
    
    def _read_tile(
        self,
        x_offset: int,
        y_offset: int,
        width: int,
        height: int
    ) -> np.ndarray:
        """
        Read a tile from the raster dataset.
        
        Args:
            x_offset: X position to read from
            y_offset: Y position to read from
            width: Width to read
            height: Height to read
            
        Returns:
            Numpy array of shape (bands, height, width)
        """
        # Read all bands
        bands = []
        for band_idx in range(1, self.num_bands + 1):
            band = self.dataset.GetRasterBand(band_idx)
            
            # Read band data
            data = band.ReadAsArray(
                int(x_offset),
                int(y_offset),
                int(width),
                int(height)
            )
            
            if data is None:
                raise RuntimeError(
                    f"Failed to read band {band_idx} at "
                    f"({x_offset}, {y_offset}, {width}, {height})"
                )
            
            bands.append(data)
        
    # Stack bands: shape (bands, height, width)
        tile_array = np.stack(bands, axis=0)
        
        return tile_array

    def _pad_to_divisibility(
        self,
        tile_array: np.ndarray,
        tile_width: int,
        tile_height: int
    ) -> Tuple[np.ndarray, int, int, int, int]:
        """
        Pad tile to ensure dimensions are divisible by architecture constraint.
        
        Some architectures (UNet++, DeepLabV3+, etc.) require input dimensions
        to be divisible by specific factors (e.g., 32, 16).
        
        Args:
            tile_array: Input array (C, H, W) - already has edge padding
            tile_width: Target tile width (without padding)
            tile_height: Target tile height (without padding)
            
        Returns:
            Tuple of (padded_array, extra_pad_left, extra_pad_top, extra_pad_right, extra_pad_bottom)
        """
        if self.divisibility <= 1:
            # No constraint
            return tile_array, 0, 0, 0, 0
        
        c, h, w = tile_array.shape
        
        # Calculate required padding to meet divisibility constraint
        pad_h = (self.divisibility - (h % self.divisibility)) % self.divisibility
        pad_w = (self.divisibility - (w % self.divisibility)) % self.divisibility
        
        # Distribute padding (prefer symmetric, but adjust if odd)
        extra_pad_top = pad_h // 2
        extra_pad_bottom = pad_h - extra_pad_top
        extra_pad_left = pad_w // 2
        extra_pad_right = pad_w - extra_pad_left
        
        # Apply reflection padding if needed
        if pad_h > 0 or pad_w > 0:
            padded_array = np.pad(
                tile_array,
                ((0, 0), (extra_pad_top, extra_pad_bottom), (extra_pad_left, extra_pad_right)),
                mode='reflect'
            )
        else:
            padded_array = tile_array
        
        return padded_array, extra_pad_left, extra_pad_top, extra_pad_right, extra_pad_bottom

    def _apply_edge_padding(
            self,
            tile_array: np.ndarray,
            tile_width: int,
            tile_height: int,
            x_offset: int,
            y_offset: int,
            x_read: int,
            y_read: int
        ) -> Tuple[np.ndarray, int, int]:
        """
        Apply padding to tile at raster edges using reflection.
        
        For tiles NOT at edges, tile_array already has correct size (read with padding).
        For tiles AT edges, we need to add reflection padding where raster boundary was hit.
        
        Args:
            tile_array: Input array (C, H, W) - already read with padding where possible
            tile_width: Tile width (without padding)
            tile_height: Tile height (without padding)
            x_offset: Tile position in raster (without padding)
            y_offset: Tile position in raster (without padding)
            x_read: Actual X read position
            y_read: Actual Y read position
            
        Returns:
            Tuple of (padded_array, actual_pad_left, actual_pad_top)
            - padded_array: shape (C, tile_height + 2*padding, tile_width + 2*padding)
            - actual_pad_left: Padding présent à gauche dans le array retourné
            - actual_pad_top: Padding présent en haut dans le array retourné
        """
        target_height = tile_height + 2 * self.padding
        target_width = tile_width + 2 * self.padding
        
        # Calculate how much padding we want on each side
        desired_pad_left = self.padding
        desired_pad_top = self.padding
        desired_pad_right = self.padding
        desired_pad_bottom = self.padding
        
        # Calculate how much padding we actually got from raster read
        # (limited by raster boundaries)
        read_pad_left = x_offset - x_read
        read_pad_top = y_offset - y_read
        read_pad_right = (x_read + tile_array.shape[2]) - (x_offset + tile_width)
        read_pad_bottom = (y_read + tile_array.shape[1]) - (y_offset + tile_height)
        
        # Calculate missing padding that needs to be added via reflection
        missing_pad_left = desired_pad_left - read_pad_left
        missing_pad_top = desired_pad_top - read_pad_top
        missing_pad_right = desired_pad_right - read_pad_right
        missing_pad_bottom = desired_pad_bottom - read_pad_bottom
        
        # Ensure no negative values
        missing_pad_left = max(0, missing_pad_left)
        missing_pad_top = max(0, missing_pad_top)
        missing_pad_right = max(0, missing_pad_right)
        missing_pad_bottom = max(0, missing_pad_bottom)
        
        # Apply reflection padding only where needed
        if missing_pad_left > 0 or missing_pad_top > 0 or missing_pad_right > 0 or missing_pad_bottom > 0:
            padded = np.pad(
                tile_array,
                (
                    (0, 0),  # No padding on channel dimension
                    (missing_pad_top, missing_pad_bottom),
                    (missing_pad_left, missing_pad_right)
                ),
                mode='reflect'
            )
        else:
            padded = tile_array
        
        # Return the actual total padding in the final array
        # This is always self.padding (target padding), regardless of raster edges
        actual_pad_left = self.padding
        actual_pad_top = self.padding
        
        return padded, actual_pad_left, actual_pad_top

    
    def _normalize_tile(self, tile_array: np.ndarray) -> torch.Tensor:
        """
        Normalize tile and convert to PyTorch tensor.
        
        Args:
            tile_array: Numpy array of shape (C, H, W) with uint8 or float values
            
        Returns:
            Normalized PyTorch tensor of shape (C, H, W) with float32
        """
        # Convert to float32 and scale to [0, 1]
        tile_float = tile_array.astype(np.float32) / 255.0
        
        # Apply normalization (ImageNet or custom)
        mean = np.array(self.normalization['mean'], dtype=np.float32).reshape(-1, 1, 1)
        std = np.array(self.normalization['std'], dtype=np.float32).reshape(-1, 1, 1)
        
        # Handle case where we have fewer bands than normalization expects
        if tile_float.shape[0] < len(mean):
            mean = mean[:tile_float.shape[0]]
            std = std[:tile_float.shape[0]]
        
        # Normalize: (x - mean) / std
        tile_normalized = (tile_float - mean) / std
        
        # Convert to PyTorch tensor
        tile_tensor = torch.from_numpy(tile_normalized).float()
        
        return tile_tensor
    
    def get_raster_info(self) -> Dict[str, Any]:
        """
        Get information about the raster being processed.
        
        Returns:
            Dictionary with raster metadata
        """
        return {
            'width': self.width,
            'height': self.height,
            'num_bands': self.num_bands,
            'geo_transform': self.geo_transform,
            'projection': self.projection,
            'extent_width': self.extent_width,
            'extent_height': self.extent_height,
            'tile_size': self.tile_size,
            'overlap': self.overlap,
            'padding': self.padding,
            'stride': self.stride,
            'num_tiles_x': self.num_tiles_x,
            'num_tiles_y': self.num_tiles_y,
            'total_tiles': self.total_tiles
        }
    
    def cleanup(self):
        """
        Clean up resources (close GDAL dataset).
        """
        if self.dataset is not None:
            self.dataset = None