# image_processing.py
"""
This module provides functions for loading, processing, and manipulating images.
It includes functionalities for adjusting contrast, brightness, and applying color palettes.
"""

import cv2
import numpy as np
import os
import logging
import config # Import the config module

def load_image(path, expected_dtype):
    """
    Loads an image from the given path and ensures the correct data type.
    If the image is not found or cannot be loaded, a placeholder image is returned.

    Args:
        path (str): The path to the image file.
        expected_dtype (numpy.dtype): The expected data type of the image.

    Returns:
        numpy.ndarray: The loaded image as a NumPy array.
    """
    if not os.path.exists(path):
        print(f"Error: Image file not found at {path}")
        logging.error(f"Image file not found at {path}")
        # Create a placeholder image instead
        if expected_dtype == np.uint8:
            return np.zeros((config.MFD_HEIGHT, config.MFD_WIDTH), dtype=np.uint8)
        else:
            return np.zeros((config.SAR_HEIGHT, config.SAR_WIDTH), dtype=config.SAR_DATA_TYPE)

    img = cv2.imread(path, cv2.IMREAD_ANYDEPTH)  # Load as is

    if img is None:
        print(f"Error: Could not load image at {path}")
        logging.error(f"Could not load image at {path}")
        # Create a placeholder image instead
        if expected_dtype == np.uint8:
            return np.zeros((config.MFD_HEIGHT, config.MFD_WIDTH), dtype=np.uint8)
        else:
            return np.zeros((config.SAR_HEIGHT, config.SAR_WIDTH), dtype=config.SAR_DATA_TYPE)

    if img.dtype != expected_dtype:
        print(f"Warning: Converting image at {path} from {img.dtype} to {expected_dtype}")
        logging.warning(f"Converting image at {path} from {img.dtype} to {expected_dtype}")
        img = img.astype(expected_dtype)
    return img

def apply_brightness_contrast(image, brightness=0, contrast=1.0):
    """
    Adjusts the brightness and contrast of an image.

    Args:
        image (numpy.ndarray): The input image.
        brightness (int): The brightness value (default: 0).
        contrast (float): The contrast value (default: 1.0).

    Returns:
        numpy.ndarray: The adjusted image.
    """
    brightness_contrast_lut = create_brightness_contrast_lut(brightness, contrast)
    return cv2.LUT(image, brightness_contrast_lut)

def create_brightness_contrast_lut(brightness=0, contrast=1.0):
    """
    Creates a Look-Up Table (LUT) for brightness and contrast adjustment.

    Args:
        brightness (int): The brightness value (default: 0).
        contrast (float): The contrast value (default: 1.0).

    Returns:
        numpy.ndarray: The LUT as a NumPy array.
    """
    lut = np.array([np.clip(i * contrast + brightness, 0, 255).astype(np.uint8) for i in range(256)])
    return lut

def apply_color_palette(image, palette):
    """
    Applies the given color palette to the image.

    Args:
        image (numpy.ndarray): The input image.
        palette (str): The name of the color palette (e.g., "GRAY", "JET", "RAINBOW").

    Returns:
        numpy.ndarray: The colorized image.
    """
    try:
        colormap = getattr(cv2, f"COLORMAP_{palette.upper()}")
        colorized_image = cv2.applyColorMap(image, colormap)
        return colorized_image
    except AttributeError:
        print(f"Error: Colormap '{palette}' not found in OpenCV.")
        logging.error(f"Colormap '{palette}' not found in OpenCV.")
        return image  # Return the original image if the colormap is invalid
    except Exception as e:
        print(f"Error applying colormap: {e}")
        logging.exception(f"Error applying colormap: {e}")
        return image

def normalize_image(image, target_type=cv2.CV_8U):
    """
    Normalizes the image to the specified data type and range (0-255).

    Args:
        image (numpy.ndarray): The input image.
        target_type (int): The target data type (default: cv2.CV_8U).

    Returns:
        numpy.ndarray: The normalized image.
    """
    return cv2.normalize(image, None, 0, 255, cv2.NORM_MINMAX, target_type)

def resize_image(image, width, height, interpolation=cv2.INTER_AREA):
    """
    Resizes the image to the specified width and height.

    Args:
        image (numpy.ndarray): The input image.
        width (int): The target width.
        height (int): The target height.
        interpolation (int): The interpolation method (default: cv2.INTER_AREA).

    Returns:
        numpy.ndarray: The resized image.
    """
    return cv2.resize(image, (width, height), interpolation=interpolation)