import torch
import numpy as np
from PIL import Image
from transformers.image_processing_utils import BaseImageProcessor, BatchFeature
from transformers.utils import logging

logger = logging.get_logger(__name__)

class LaTeXOCRImageProcessor(BaseImageProcessor):
    model_type = "latex_ocr"

    def __init__(
        self,
        image_height=64,
        max_image_width=1024,
        patch_size=16,
        **kwargs
    ):
        super().__init__(**kwargs)
        self.image_height = image_height
        self.max_image_width = max_image_width
        self.patch_size = patch_size

    def preprocess(self, images, **kwargs) -> BatchFeature:
        if not isinstance(images, list):
            images = [images]
        
        processed_images = []
        for img in images:
            if img.mode != "RGB":
                img = img.convert("RGB")

            w, h = img.size
            new_w = int(round(w * self.image_height / max(h, 1)))
            new_w = min(new_w, self.max_image_width)
            new_w = max((new_w // self.patch_size) * self.patch_size, self.patch_size)
            
            if (w, h) != (new_w, self.image_height):
                img = img.resize((new_w, self.image_height), Image.BILINEAR)

            img_array = np.array(img).astype(np.float32) / 255.0
            img_array = np.transpose(img_array, (2, 0, 1))
            processed_images.append(img_array)

        return BatchFeature(data={"pixel_values": processed_images}, tensor_type="pt")