"""
Image processing utilities using PIL only.

No OpenCV dependency - we use Pillow for all image operations.
"""

from typing import Union, Tuple, Optional
from pathlib import Path
import numpy as np
from PIL import Image
import torch


def load_image(
    source: Union[str, Path, np.ndarray, Image.Image],
    size: Optional[Tuple[int, int]] = None
) -> Image.Image:
    """
    Load image from various sources.
    
    Args:
        source: File path, numpy array, or PIL Image
        size: Optional (width, height) to resize to
        
    Returns:
        PIL Image in RGB mode
    """
    if isinstance(source, (str, Path)):
        img = Image.open(source).convert('RGB')
    elif isinstance(source, np.ndarray):
        img = Image.fromarray(source).convert('RGB')
    elif isinstance(source, Image.Image):
        img = source.convert('RGB')
    else:
        raise TypeError(f"Unsupported image type: {type(source)}")
    
    if size:
        img = img.resize(size, Image.Resampling.BILINEAR)
    
    return img


def save_image(
    image: Union[Image.Image, np.ndarray, torch.Tensor],
    path: Union[str, Path],
    quality: int = 95
) -> None:
    """
    Save image to file.
    
    Args:
        image: PIL Image, numpy array, or torch Tensor
        path: Output file path
        quality: JPEG quality (if applicable)
    """
    if isinstance(image, torch.Tensor):
        # Assume (C, H, W) or (H, W) tensor
        arr = image.detach().cpu().numpy()
        if len(arr.shape) == 3 and arr.shape[0] in (1, 3):
            arr = np.transpose(arr, (1, 2, 0))
        if arr.max() <= 1.0:
            arr = (arr * 255).astype(np.uint8)
        image = Image.fromarray(arr.squeeze())
    elif isinstance(image, np.ndarray):
        if image.max() <= 1.0:
            image = (image * 255).astype(np.uint8)
        image = Image.fromarray(image)
    
    path = Path(path)
    if path.suffix.lower() in ('.jpg', '.jpeg'):
        image.save(path, quality=quality)
    else:
        image.save(path)


def overlay_heatmap(
    image: Image.Image,
    heatmap: Union[np.ndarray, torch.Tensor],
    alpha: float = 0.5,
    colormap: str = 'jet'
) -> Image.Image:
    """
    Overlay heatmap on image.
    
    Args:
        image: Base PIL Image
        heatmap: 2D array with values in [0, 1]
        alpha: Heatmap opacity (0-1)
        colormap: Color scheme ('jet', 'hot', 'viridis')
        
    Returns:
        PIL Image with heatmap overlay
    """
    if isinstance(heatmap, torch.Tensor):
        heatmap = heatmap.detach().cpu().numpy()
    
    # Ensure 2D
    heatmap = np.squeeze(heatmap)
    
    # Normalize to [0, 1]
    if heatmap.max() > 1.0 or heatmap.min() < 0.0:
        heatmap = (heatmap - heatmap.min()) / (heatmap.max() - heatmap.min() + 1e-8)
    
    # Resize heatmap to image size
    heatmap_img = Image.fromarray((heatmap * 255).astype(np.uint8))
    heatmap_img = heatmap_img.resize(image.size, Image.Resampling.BILINEAR)
    heatmap = np.array(heatmap_img) / 255.0
    
    # Apply colormap
    colored = apply_colormap(heatmap, colormap)
    
    # Convert to PIL
    heatmap_rgb = Image.fromarray((colored * 255).astype(np.uint8))
    
    # Blend with original
    return Image.blend(image, heatmap_rgb, alpha)


def apply_colormap(
    data: np.ndarray, 
    colormap: str = 'jet'
) -> np.ndarray:
    """
    Apply colormap to grayscale data.
    
    Args:
        data: 2D array with values in [0, 1]
        colormap: 'jet', 'hot', or 'viridis'
        
    Returns:
        3D RGB array with shape (H, W, 3)
    """
    h, w = data.shape
    rgb = np.zeros((h, w, 3), dtype=np.float32)
    
    if colormap == 'jet':
        # Proper jet colormap: Blue -> Cyan -> Green -> Yellow -> Red
        # More accurate implementation
        r = np.clip(1.5 - np.abs(4 * data - 3), 0, 1)
        g = np.clip(1.5 - np.abs(4 * data - 2), 0, 1)
        b = np.clip(1.5 - np.abs(4 * data - 1), 0, 1)
        
        # Adjust for better visibility
        rgb[:, :, 0] = np.where(data > 0.5, r, data * 0.5)
        rgb[:, :, 1] = g
        rgb[:, :, 2] = np.where(data < 0.5, 1 - data, 0)
        
    elif colormap == 'hot':
        # Black -> Red -> Yellow -> White
        rgb[:, :, 0] = np.clip(data * 3, 0, 1)
        rgb[:, :, 1] = np.clip(data * 3 - 1, 0, 1)
        rgb[:, :, 2] = np.clip(data * 3 - 2, 0, 1)
        
    elif colormap == 'viridis':
        # Approximate viridis
        rgb[:, :, 0] = np.clip(data * 0.5 + 0.3, 0, 1) * data
        rgb[:, :, 1] = np.clip(data * 0.8, 0, 1)
        rgb[:, :, 2] = np.clip(0.9 - data * 0.5, 0, 1)
    
    else:
        # Fallback: grayscale
        rgb[:, :, 0] = data
        rgb[:, :, 1] = data
        rgb[:, :, 2] = data
    
    return rgb


def create_detection_overlay(
    image: Image.Image,
    heatmap: np.ndarray,
    bbox: tuple,
    alpha: float = 0.6,
    show_bbox: bool = True,
    bbox_color: tuple = (0, 255, 0)
) -> Image.Image:
    """
    Create heatmap overlay focused on detection bbox.
    
    This creates a more interpretable visualization by:
    1. Overlaying heatmap only in the bbox region
    2. Drawing the bbox boundary
    3. Keeping the rest of image visible
    
    Args:
        image: Original PIL Image
        heatmap: 2D attribution map
        bbox: (x1, y1, x2, y2) bounding box coordinates
        alpha: Heatmap opacity in bbox region
        show_bbox: Whether to draw bbox outline
        bbox_color: RGB color for bbox
        
    Returns:
        Image with focused heatmap overlay
    """
    from PIL import ImageDraw
    
    # Convert heatmap to proper format
    if hasattr(heatmap, 'cpu'):
        heatmap = heatmap.cpu().numpy()
    heatmap = np.squeeze(heatmap)
    
    # Normalize
    if heatmap.max() > 1.0 or heatmap.min() < 0.0:
        heatmap = (heatmap - heatmap.min()) / (heatmap.max() - heatmap.min() + 1e-8)
    
    # Resize heatmap to image size
    heatmap_img = Image.fromarray((heatmap * 255).astype(np.uint8))
    heatmap_img = heatmap_img.resize(image.size, Image.Resampling.BILINEAR)
    heatmap_arr = np.array(heatmap_img) / 255.0
    
    # Apply colormap
    colored = apply_colormap(heatmap_arr, 'jet')
    
    # Create result image (copy of original)
    result = image.copy().convert('RGB')
    result_arr = np.array(result).astype(np.float32) / 255.0
    
    # Get bbox coordinates (clamp to image bounds)
    x1, y1, x2, y2 = [int(c) for c in bbox]
    x1, y1 = max(0, x1), max(0, y1)
    x2, y2 = min(image.width, x2), min(image.height, y2)
    
    # Blend heatmap only in bbox region
    result_arr[y1:y2, x1:x2] = (
        (1 - alpha) * result_arr[y1:y2, x1:x2] + 
        alpha * colored[y1:y2, x1:x2]
    )
    
    # Also add subtle heatmap outside bbox (lower alpha)
    outside_alpha = 0.2
    mask = np.ones_like(result_arr[:, :, 0], dtype=bool)
    mask[y1:y2, x1:x2] = False
    
    for c in range(3):
        result_arr[:, :, c] = np.where(
            mask,
            (1 - outside_alpha) * result_arr[:, :, c] + outside_alpha * colored[:, :, c],
            result_arr[:, :, c]
        )
    
    # Convert back to PIL
    result = Image.fromarray((result_arr * 255).clip(0, 255).astype(np.uint8))
    
    # Draw bbox
    if show_bbox:
        draw = ImageDraw.Draw(result)
        draw.rectangle([x1, y1, x2, y2], outline=bbox_color, width=3)
    
    return result


def tensor_to_pil(tensor: torch.Tensor) -> Image.Image:
    """
    Convert torch Tensor to PIL Image.
    
    Args:
        tensor: Tensor of shape (C, H, W) or (H, W)
        
    Returns:
        PIL Image
    """
    arr = tensor.detach().cpu().numpy()
    
    if len(arr.shape) == 3:
        if arr.shape[0] in (1, 3):
            arr = np.transpose(arr, (1, 2, 0))
    
    if arr.max() <= 1.0:
        arr = (arr * 255)
    
    arr = arr.clip(0, 255).astype(np.uint8)
    
    if arr.shape[-1] == 1:
        arr = arr.squeeze(-1)
    
    return Image.fromarray(arr)


def pil_to_tensor(image: Image.Image) -> torch.Tensor:
    """
    Convert PIL Image to torch Tensor.
    
    Args:
        image: PIL Image
        
    Returns:
        Tensor of shape (3, H, W) with values in [0, 1]
    """
    arr = np.array(image.convert('RGB'))
    tensor = torch.from_numpy(arr).float() / 255.0
    tensor = tensor.permute(2, 0, 1)  # HWC -> CHW
    return tensor
