"""Deterministic image preprocessing for CuRe."""

import io
from pathlib import Path

import albumentations as A
import cv2
import numpy as np
from albumentations.pytorch import ToTensorV2
from PIL import Image


def preprocess_image(path):
    """Read an image and return a normalized float tensor of shape [3, 336, 336]."""
    path = Path(path)
    with Image.open(path) as source:
        image = source.convert("RGB")
        image.load()
    if path.suffix.lower() == ".png":
        buffer = io.BytesIO()
        image.save(buffer, format="JPEG", quality=96)
        buffer.seek(0)
        with Image.open(buffer) as decoded:
            image = decoded.convert("RGB")
            image.load()
    pipeline = A.Compose([
        A.SmallestMaxSize(max_size=390, interpolation=cv2.INTER_CUBIC, p=1.0),
        A.CenterCrop(height=390, width=390, p=1.0),
        A.Resize(height=336, width=336, interpolation=cv2.INTER_CUBIC, p=1.0),
        A.ImageCompression(quality_range=(70, 70), p=1.0),
        A.Normalize(mean=(0.5, 0.5, 0.5), std=(0.5, 0.5, 0.5)),
        ToTensorV2(),
    ])
    return pipeline(image=np.asarray(image))["image"]