"""Efficient image rendering utilities for PyQt6."""

import numpy as np
from PyQt6.QtGui import QImage, QPixmap, QColor
from PyQt6.QtCore import Qt
from typing import Optional, Tuple
import logging

logger = logging.getLogger(__name__)


def numpy_to_qimage(
    data: np.ndarray,
    colormap: Optional[Tuple[int, int, int]] = None,
    normalize: bool = True,
) -> QImage:
    """
    Convert NumPy array to QImage efficiently.
    
    Args:
        data: 2D NumPy array (grayscale) or 3D array (RGB)
        colormap: Optional (R, G, B) tuple for single-channel tinting
        normalize: If True, normalize to [0, 255]
    
    Returns:
        QImage ready for display
    """
    if data is None:
        return QImage()
    
    # Ensure proper dtype
    if data.dtype != np.uint8:
        if normalize:
            # Normalize to 0-255
            if data.size == 0:
                data = np.zeros((1, 1), dtype=np.uint8)
            else:
                data_min = np.min(data)
                data_max = np.max(data)
                if data_max > data_min:
                    data = ((data - data_min) / (data_max - data_min) * 255).astype(np.uint8)
                else:
                    # If all values are the same, use mid-range gray
                    data = np.full_like(data, 128, dtype=np.uint8)
        else:
            # Clip to 0-255
            data = np.clip(data, 0, 255).astype(np.uint8)
    
    if data.ndim == 2:
        # Grayscale: create RGB by duplicating
        h, w = data.shape
        if colormap is not None:
            # Apply color tint by multiplying normalized data by color values
            # First normalize data to 0-1 range
            if data.dtype != np.uint8:
                data_norm = data.astype(np.float32)
                data_min = np.min(data_norm)
                data_max = np.max(data_norm)
                if data_max > data_min:
                    data_norm = (data_norm - data_min) / (data_max - data_min)
                else:
                    data_norm = np.ones_like(data_norm)
            else:
                data_norm = data.astype(np.float32) / 255.0
            
            # Apply color by scaling normalized data
            rgb_data = np.zeros((h, w, 3), dtype=np.uint8)
            rgb_data[:, :, 0] = (data_norm * colormap[0]).astype(np.uint8)
            rgb_data[:, :, 1] = (data_norm * colormap[1]).astype(np.uint8)
            rgb_data[:, :, 2] = (data_norm * colormap[2]).astype(np.uint8)
            data = rgb_data
        else:
            # Standard grayscale to RGB
            data = np.stack([data, data, data], axis=2)
    
    if data.ndim != 3 or data.shape[2] not in (3, 4):
        logger.warning(f"Unexpected array shape: {data.shape}")
        return QImage()
    
    # Ensure C-contiguous for QImage
    if not data.flags['C_CONTIGUOUS']:
        data = np.ascontiguousarray(data)
    
    h, w = data.shape[:2]
    bytes_per_line = 3 * w if data.shape[2] == 3 else 4 * w
    
    if data.shape[2] == 3:
        # Convert RGB to RGBA for QImage
        rgba = np.zeros((h, w, 4), dtype=np.uint8)
        rgba[:, :, :3] = data
        rgba[:, :, 3] = 255
        bytes_per_line = 4 * w
        data = rgba
    
    # Create QImage from buffer
    qimg = QImage(
        data.tobytes(),
        w,
        h,
        bytes_per_line,
        QImage.Format.Format_RGBA8888
    )
    
    return qimg


def numpy_to_qpixmap(
    data: np.ndarray,
    colormap: Optional[Tuple[int, int, int]] = None,
    normalize: bool = True,
    scale_to: Optional[Tuple[int, int]] = None,
) -> QPixmap:
    """
    Convert NumPy array to QPixmap with optional scaling.
    
    Args:
        data: 2D NumPy array
        colormap: Optional color tint (R, G, B)
        normalize: Whether to normalize to 0-255
        scale_to: Optional (width, height) for scaling
    
    Returns:
        QPixmap ready for display
    """
    qimg = numpy_to_qimage(data, colormap, normalize)
    
    if scale_to is not None:
        qimg = qimg.scaled(
            scale_to[0],
            scale_to[1],
            Qt.AspectRatioMode.KeepAspectRatio,
            Qt.TransformationMode.SmoothTransformation
        )
    
    return QPixmap.fromImage(qimg)


def apply_colormap(
    data: np.ndarray,
    color: Tuple[int, int, int],
) -> np.ndarray:
    """
    Apply a single color tint to grayscale data.
    
    Args:
        data: 2D grayscale array
        color: (R, G, B) tuple with values 0-255
    
    Returns:
        RGB array with color applied
    """
    if data.ndim != 2:
        raise ValueError("Data must be 2D")
    
    data = data.astype(np.float32) / 255.0  # Normalize to 0-1
    
    rgb = np.zeros((*data.shape, 3), dtype=np.uint8)
    rgb[:, :, 0] = (data * color[0]).astype(np.uint8)
    rgb[:, :, 1] = (data * color[1]).astype(np.uint8)
    rgb[:, :, 2] = (data * color[2]).astype(np.uint8)
    
    return rgb


def composite_channels(
    channels: dict[int, np.ndarray],
    colors: dict[int, Tuple[int, int, int]],
) -> np.ndarray:
    """
    Composite multiple grayscale channels into RGB with specified colors.
    
    Args:
        channels: Dict of channel_idx -> 2D array
        colors: Dict of channel_idx -> (R, G, B) tuple
    
    Returns:
        RGB composite image
    """
    if not channels:
        return np.zeros((1, 1, 3), dtype=np.uint8)
    
    # Get image dimensions from first channel
    first_channel = next(iter(channels.values()))
    h, w = first_channel.shape
    
    composite = np.zeros((h, w, 3), dtype=np.float32)
    
    for ch_idx, data in channels.items():
        if ch_idx not in colors:
            continue
        
        # Normalize channel data
        if data.size > 0:
            data_min = np.min(data)
            data_max = np.max(data)
            if data_max > data_min:
                normalized = (data - data_min) / (data_max - data_min)
            else:
                normalized = np.zeros_like(data, dtype=np.float32)
        else:
            normalized = np.zeros_like(data, dtype=np.float32)
        
        # Apply color
        r, g, b = colors[ch_idx]
        composite[:, :, 0] += normalized * (r / 255.0)
        composite[:, :, 1] += normalized * (g / 255.0)
        composite[:, :, 2] += normalized * (b / 255.0)
    
    # Normalize to 0-255
    composite = np.clip(composite * 255, 0, 255).astype(np.uint8)
    
    return composite


def hex_to_rgb(hex_color: str) -> Tuple[int, int, int]:
    """Convert hex color string to RGB tuple."""
    hex_color = hex_color.lstrip('#')
    return tuple(int(hex_color[i:i+2], 16) for i in (0, 2, 4))


def rgb_to_hex(rgb: Tuple[int, int, int]) -> str:
    """Convert RGB tuple to hex color string."""
    return '#{:02x}{:02x}{:02x}'.format(*rgb)


def create_histogram_image(
    data: np.ndarray,
    width: int = 256,
    height: int = 100,
) -> QPixmap:
    """
    Create a histogram visualization as a QPixmap.
    
    Args:
        data: 1D or 2D array
        width: Width of histogram image
        height: Height of histogram image
    
    Returns:
        QPixmap with histogram
    """
    # Flatten if needed
    if data.ndim > 1:
        data = data.flatten()
    
    # Handle empty data
    if data.size == 0:
        return QPixmap(width, height)
    
    # Create histogram
    try:
        # Use 256 bins for typical uint8 range
        max_val = max(256, int(np.max(data)) + 1)
        hist, _ = np.histogram(data, bins=256, range=(0, max_val))
    except Exception as e:
        logger.error(f"Error computing histogram: {e}")
        return QPixmap(width, height)
    
    # Normalize to image height
    hist_max = np.max(hist)
    if hist_max > 0:
        hist = (hist / hist_max * (height - 1)).astype(np.uint8)
    
    # Create image
    img = np.zeros((height, width, 3), dtype=np.uint8)
    img[:, :] = [50, 50, 50]  # Dark background
    
    # Draw histogram
    bin_width = width / 256
    for i, h in enumerate(hist):
        x_start = int(i * bin_width)
        x_end = int((i + 1) * bin_width)
        y_start = height - 1 - h
        
        if x_end > x_start:
            img[y_start:, x_start:x_end] = [100, 200, 100]  # Green histogram
    
    return numpy_to_qpixmap(img, normalize=False)
