import numpy as np
from PIL import Image
import config

def preprocess_image(image, target_size=512):
        """
        Preprocess the input image: resize and ensure it's RGB.

        Parameters:
        -----------
        image : PIL.Image or np.ndarray
            Input image
        target_size : int
            Target size for the image (will be square)

        Returns:
        --------
        np.ndarray : Preprocessed image
        """
        if isinstance(image, np.ndarray):
            image = Image.fromarray(image)

        # Convert to RGB if needed
        if image.mode != 'RGB':
            image = image.convert('RGB')

        # Resize to target size (square)
        image = image.resize((target_size, target_size))

        return np.array(image)

def divide_into_grid(image):
    """
    Divide image into a grid of cells.

    Parameters:
    -----------
    image : np.ndarray
        Input image

    Returns:
    --------
    list : List of cell information (row, col, cell_array, avg_color)
    """
    height, width = image.shape[:2]
    grid_h, grid_w = config.grid_size

    cell_height = height // grid_h
    cell_width = width // grid_w

    cells = []

    try : 
        cells = image.reshape(grid_h, cell_height, grid_w, cell_width, 3)
    except Exception: 
         raise TypeError("Image too big! The Grid doesn't divide the resized image evenly due to cascading rounding")
    cells = cells.transpose(0, 2, 1, 3, 4)
    cells = cells.reshape(-1, cell_height, cell_width, 3)
    #may return errors, double check after optimizing compute average colors

    vectorized = np.vectorize(compute_average_color, signature='(i,j,k)->(n)')
    cells = vectorized(cells)
    return cells

def compute_average_color(cell):
    """
    Compute the average color of a cell.

    Parameters:
    -----------
    cell : np.ndarray
        Image cell

    Returns:
    --------
    np.ndarray : Average RGB color
    """

    return np.mean(cell, axis=(0,1))