"""
TILA — Image Processor

Single processor that handles the full pipeline:
    raw image (path, numpy, or PIL) → model-ready tensor [1, 3, 448, 448]

Combines:
    1. Medical image preprocessing (windowing, padding removal, resize)
    2. Model transforms (resize, center crop, to tensor, expand channels)

Usage:
    from processor import TILAProcessor

    processor = TILAProcessor()

    # From file path (applies full preprocessing)
    tensor = processor("raw_cxr.png")

    # From PIL image (skips medical preprocessing, applies model transforms only)
    tensor = processor(Image.open("preprocessed.png"))

    # Pair of images for the model
    current = processor("current.png")
    previous = processor("previous.png")
    result = model.get_interval_change_prediction(current, previous)
"""

import cv2
import numpy as np
import torch
from PIL import Image
from torchvision import transforms
from typing import Union

from preprocess import preprocess_image


class TILAProcessor:
    """End-to-end image processor for the TILA model.

    Accepts file paths (str/Path), numpy arrays, or PIL Images.
    - File paths: full pipeline (windowing → crop → resize → model transform)
    - Numpy arrays: treated as raw, full pipeline applied
    - PIL Images: assumed already preprocessed, only model transforms applied

    Args:
        raw_preprocess: Apply medical preprocessing (windowing, padding removal).
                        Set False if images are already preprocessed PNGs.
        width_param: Windowing width parameter (default: 4.0)
        max_size: Resize longest side to this before model transforms (default: 512)
        crop_size: Center crop size for model input (default: 448)
        dtype: Output tensor dtype (default: torch.bfloat16)
        device: Output tensor device (default: "cpu")
    """

    def __init__(
        self,
        raw_preprocess: bool = True,
        width_param: float = 4.0,
        max_size: int = 512,
        crop_size: int = 448,
        dtype: torch.dtype = torch.bfloat16,
        device: str = "cpu",
    ):
        self.raw_preprocess = raw_preprocess
        self.width_param = width_param
        self.max_size = max_size
        self.dtype = dtype
        self.device = device

        self.model_transform = transforms.Compose([
            transforms.Resize(max_size),
            transforms.CenterCrop(crop_size),
            transforms.ToTensor(),
            _ExpandChannels(),
        ])

    def __call__(self, image: Union[str, np.ndarray, Image.Image]) -> torch.Tensor:
        """Process a single image into a model-ready tensor.

        Args:
            image: File path (str), numpy array, or PIL Image

        Returns:
            Tensor of shape [1, 3, 448, 448]
        """
        if isinstance(image, str):
            if self.raw_preprocess:
                img_np = preprocess_image(image, self.width_param, self.max_size)
                pil_img = Image.fromarray(img_np)
            else:
                pil_img = Image.open(image).convert("L")
        elif isinstance(image, np.ndarray):
            if self.raw_preprocess:
                from preprocess import apply_windowing, remove_black_padding, resize_preserve_aspect_ratio
                if len(image.shape) == 3:
                    image = cv2.cvtColor(image, cv2.COLOR_BGR2GRAY)
                image = apply_windowing(image, self.width_param)
                image = (image * 255.0).astype(np.uint8)
                image = remove_black_padding(image)
                image = resize_preserve_aspect_ratio(image, self.max_size)
            pil_img = Image.fromarray(image)
        elif isinstance(image, Image.Image):
            pil_img = image.convert("L")
        else:
            raise TypeError(f"Expected str, np.ndarray, or PIL.Image, got {type(image)}")

        tensor = self.model_transform(pil_img).unsqueeze(0)
        return tensor.to(dtype=self.dtype, device=self.device)


class _ExpandChannels:
    """Expand single-channel tensor to 3 channels."""
    def __call__(self, x: torch.Tensor) -> torch.Tensor:
        if x.shape[0] == 1:
            return x.repeat(3, 1, 1)
        return x