from PIL import Image
import numpy as np


class ImageProcess:
    def __init__(self, size=256):
        self.size = size

    def process(self, image_path):
        with Image.open(image_path) as img:
            image = img.convert('RGB')

        w, h = image.size

        # 1. Match PyTorch's resizing dimensions (cast to int, don't use round)
        if w < h:
            new_w = self.size
            new_h = int(h * self.size / w)
        else:
            new_h = self.size
            new_w = int(w * self.size / h)

        image = image.resize((new_w, new_h), Image.Resampling.BILINEAR)

        # 2. Center crop math
        left = (new_w - self.size) // 2
        top = (new_h - self.size) // 2
        image = image.crop((left, top, left + self.size, top + self.size))

        # 3. Convert to float array AND scale to [0.0, 1.0]
        # CRITICAL: This is what transforms.ToTensor() does under the hood
        image = np.asarray(image, dtype=np.float32) / 255.0

        # 4. Reorder dimensions to CHW and expand to BCHW
        image = np.transpose(image, (2, 0, 1))
        image = np.expand_dims(image, 0)

        return image