"""
ChequeImageProcessor – preprocessing helper for ChequeRegressorModel.

Handles resizing, tensor conversion, and ImageNet normalisation.
Can be used standalone or with HuggingFace pipelines.

Usage::

    from image_processing_cheque import ChequeImageProcessor
    from PIL import Image

    processor = ChequeImageProcessor()                  # or from_pretrained(...)
    inputs = processor(images=Image.open("cheque.tif"))
    pixel_values = inputs["pixel_values"]               # [1, 3, 512, 1024]
"""

from __future__ import annotations

from typing import Dict, List, Optional, Union

import numpy as np
import torch
from PIL import Image
from transformers import BaseImageProcessor, BatchFeature
from transformers.image_utils import ImageInput


class ChequeImageProcessor(BaseImageProcessor):
    """Minimal image processor for :class:`ChequeRegressorModel`.

    Args:
        img_height  : Target height after resizing (default 512).
        img_width   : Target width after resizing (default 1024).
        image_mean  : ImageNet RGB mean for normalisation.
        image_std   : ImageNet RGB std for normalisation.
        do_resize   : Whether to resize images (default True).
        do_normalize: Whether to apply ImageNet normalisation (default True).
    """

    model_input_names = ["pixel_values"]

    def __init__(
        self,
        img_height:   int        = 512,
        img_width:    int        = 1024,
        image_mean:   List[float] = None,
        image_std:    List[float] = None,
        do_resize:    bool       = True,
        do_normalize: bool       = True,
        **kwargs,
    ):
        super().__init__(**kwargs)
        self.img_height   = img_height
        self.img_width    = img_width
        self.image_mean   = image_mean or [0.485, 0.456, 0.406]
        self.image_std    = image_std  or [0.229, 0.224, 0.225]
        self.do_resize    = do_resize
        self.do_normalize = do_normalize

    # ── internal helpers ───────────────────────────────────────────────────────

    def _to_pil(self, image: ImageInput) -> Image.Image:
        if isinstance(image, Image.Image):
            return image.convert("RGB")
        if isinstance(image, np.ndarray):
            return Image.fromarray(image).convert("RGB")
        if isinstance(image, torch.Tensor):
            arr = image.numpy()
            if arr.ndim == 3 and arr.shape[0] in (1, 3, 4):
                arr = arr.transpose(1, 2, 0)
            return Image.fromarray((arr * 255).astype(np.uint8)).convert("RGB")
        raise TypeError(f"Unsupported image type: {type(image)}")

    def _process_one(self, image: ImageInput) -> torch.Tensor:
        img = self._to_pil(image)
        if self.do_resize:
            img = img.resize((self.img_width, self.img_height), Image.BILINEAR)
        arr = np.array(img, dtype=np.float32) / 255.0          # [H, W, 3]  ∈ [0,1]
        arr = arr.transpose(2, 0, 1)                            # [3, H, W]
        t   = torch.from_numpy(arr)
        if self.do_normalize:
            mean = torch.tensor(self.image_mean).view(3, 1, 1)
            std  = torch.tensor(self.image_std).view(3, 1, 1)
            t    = (t - mean) / std
        return t                                                 # [3, H, W]

    # ── public API ─────────────────────────────────────────────────────────────

    def __call__(
        self,
        images: Union[ImageInput, List[ImageInput]],
        return_tensors: Optional[str] = "pt",
        **kwargs,
    ) -> BatchFeature:
        """Preprocess one or more images.

        Args:
            images         : A single image or list of images.
                             Accepts PIL.Image, np.ndarray, or torch.Tensor.
            return_tensors : ``"pt"`` (default) → returns ``torch.Tensor``.

        Returns:
            :class:`BatchFeature` with key ``"pixel_values"``
            of shape ``[N, 3, img_height, img_width]``.
        """
        if not isinstance(images, (list, tuple)):
            images = [images]
        tensors = [self._process_one(img) for img in images]
        pixel_values = torch.stack(tensors)                     # [N, 3, H, W]
        return BatchFeature({"pixel_values": pixel_values})

    def postprocess_boxes(
        self,
        boxes:       torch.Tensor,
        target_sizes: Optional[List[tuple]] = None,
    ) -> List[Dict]:
        """Convert normalised model output to per-image result dicts.

        Args:
            boxes        : Float32 tensor ``[N, num_fields, 4]`` in normalised
                           ``[xmin, ymin, xmax, ymax]`` format.
            target_sizes : Optional list of ``(width, height)`` tuples for the
                           *original* images. When provided, boxes are rescaled
                           to absolute pixel coordinates.

        Returns:
            List of dicts, one per image::

                {
                    "date":   [xmin, ymin, xmax, ymax],   # pixels or normalised
                    "amount": [...],
                    ...
                }
        """
        field_names = ["date", "amount", "ifsc", "acno", "sign", "name"]
        results = []
        for i, image_boxes in enumerate(boxes):
            row = {}
            for f_idx, name in enumerate(field_names):
                box = image_boxes[f_idx].tolist()
                if target_sizes is not None:
                    w, h = target_sizes[i]
                    box = [
                        box[0] * w, box[1] * h,
                        box[2] * w, box[3] * h,
                    ]
                row[name] = [round(v, 2) for v in box]
            results.append(row)
        return results