#preprocess.py
#This file is used to preprocess the image before it is fed into the model.

import numpy as np
from PIL import Image

def normalize_image(image_input, target_size=None):
    """
    Normalizes an image by scaling pixel values to the range [0, 1].
    Works with both file paths and PIL images.

    Args:
        image_input (str or PIL.Image): The image file path or PIL Image object.
        target_size (tuple, optional): Target size as (width, height) for resizing.

    Returns:
        numpy.ndarray: The normalized image as a NumPy array, or None if the image cannot be processed.
    """
    try:
        # Handle different input types
        if isinstance(image_input, str):
            # Load from file path using PIL
            img_pil = Image.open(image_input)
            img = np.array(img_pil)
        elif isinstance(image_input, Image.Image):
            # Convert PIL image to numpy array
            img = np.array(image_input)
        else:
            print(f"Error: Unsupported image input type: {type(image_input)}")
            return None

        # Resize if target size is specified
        if target_size is not None:
            # Use PIL resize for all cases
            img_pil = Image.fromarray(img)
            img_pil = img_pil.resize(target_size)
            img = np.array(img_pil)

        # Convert the image to float32 for accurate calculations
        img = img.astype(np.float32)

        # Normalize the image to the range [0, 1]
        normalized_img = img / 255.0

        return normalized_img

    except Exception as e:
        print(f"Error processing image: {e}")
        return None

def preprocess_for_model(image_input, model_input_shape):
    """
    Preprocess image specifically for model input.
    
    Args:
        image_input (str or PIL.Image): The image to preprocess
        model_input_shape (tuple): Expected model input shape (height, width, channels)
    
    Returns:
        numpy.ndarray: Preprocessed image ready for model input
    """
    try:
        # Extract target size from model input shape
        if len(model_input_shape) == 3:
            target_size = (model_input_shape[1], model_input_shape[0])  # (width, height)
        else:
            target_size = None
            
        # Normalize image
        normalized_img = normalize_image(image_input, target_size)
        
        if normalized_img is None:
            return None
            
        # Add batch dimension if needed
        if len(model_input_shape) == 3:
            # For image models: (batch, height, width, channels)
            preprocessed = np.expand_dims(normalized_img, axis=0)
        elif len(model_input_shape) == 2:
            # For flattened models: (batch, features)
            preprocessed = normalized_img.reshape(1, -1)
        else:
            raise ValueError(f"Unsupported model input shape: {model_input_shape}")
            
        return preprocessed
        
    except Exception as e:
        print(f"Error preprocessing image for model: {e}")
        return None

def get_image_info(image_input):
    """
    Get basic information about an image.
    
    Args:
        image_input (str or PIL.Image): The image to analyze
    
    Returns:
        dict: Dictionary containing image information
    """
    try:
        if isinstance(image_input, str):
            # Use PIL for file paths
            img_pil = Image.open(image_input)
            width, height = img_pil.size
            channels = len(img_pil.getbands())
            format_type = "File (PIL)"
        elif isinstance(image_input, Image.Image):
            width, height = image_input.size
            channels = len(image_input.getbands())
            format_type = "PIL Image"
        else:
            return None
            
        return {
            "format": format_type,
            "width": width,
            "height": height,
            "channels": channels,
            "aspect_ratio": width / height if height > 0 else 0
        }
        
    except Exception as e:
        print(f"Error getting image info: {e}")
        return None