from __future__ import annotations
import numpy as np
from PIL import Image

def pil_to_np(img: Image.Image) -> np.ndarray:
    if img.mode not in ("RGB", "RGBA", "L"):
        img = img.convert("RGB")
    if img.mode == "L":
        img = img.convert("RGB")
    arr = np.asarray(img).astype(np.float32)
    if arr.ndim == 2:
        arr = np.repeat(arr[..., None], 3, axis=2)
    if arr.shape[2] == 4:
        arr = arr[..., :3]
    return arr / 255.0

def np_to_pil(arr: np.ndarray) -> Image.Image:
    return Image.fromarray(np.clip(arr * 255.0, 0, 255).astype(np.uint8))

def resize_and_crop_to_grid(img: Image.Image, width: int, height: int, grid: int) -> Image.Image:
    img = img.convert("RGB").resize((width, height), Image.LANCZOS)
    H, W = img.height, img.width
    H2, W2 = (H // grid) * grid, (W // grid) * grid
    if H2 != H or W2 != W:
        left = (W - W2) // 2
        top = (H - H2) // 2
        img = img.crop((left, top, left + W2, top + H2))
    return img

def block_view(arr: np.ndarray, bh: int, bw: int) -> np.ndarray:
    H, W, C = arr.shape
    assert H % bh == 0 and W % bw == 0, "Dims must be divisible by block."
    shape   = (H//bh, W//bw, bh, bw, C)
    strides = (arr.strides[0]*bh, arr.strides[1]*bw, arr.strides[0], arr.strides[1], arr.strides[2])
    return np.lib.stride_tricks.as_strided(arr, shape=shape, strides=strides)

def cell_means(arr: np.ndarray, grid: int) -> np.ndarray:
    H, W, _ = arr.shape
    bh, bw = H//grid, W//grid
    blocks = block_view(arr, bh, bw)
    
    # Use weighted mean with center bias for better detail preservation
    # Create a weight matrix that emphasizes the center of each block
    center_h, center_w = bh // 2, bw // 2
    weights = np.zeros((bh, bw))
    
    for i in range(bh):
        for j in range(bw):
            # Distance from center (normalized)
            dist_from_center = np.sqrt((i - center_h)**2 + (j - center_w)**2)
            max_dist = np.sqrt(center_h**2 + center_w**2)
            # Higher weight for pixels closer to center
            weights[i, j] = 1.0 - (dist_from_center / max_dist) * 0.5
    
    # Normalize weights
    weights = weights / np.sum(weights)
    
    # Apply weighted mean
    weighted_means = np.zeros((grid, grid, 3))
    for i in range(grid):
        for j in range(grid):
            block = blocks[i, j]
            for c in range(3):
                weighted_means[i, j, c] = np.sum(block[:, :, c] * weights)
    
    return weighted_means