"""Image preprocessing, identical to the original application.

ToTensor -> zero-pad to square (centred) -> resize to 448 x 448, bilinear,
no antialiasing, no mean/std normalisation.

``antialias=False`` reproduces tensor resizing in torchvision 0.13
(PyTorch 1.12.1, the version bundled in toolv9). Newer torchvision versions
default to ``antialias=True``, which changes the input noticeably.
"""
from __future__ import annotations

import torchvision.transforms.functional as F
from PIL import Image
from torchvision import transforms

# Formats the original application accepted (it used imghdr: jpeg, png, tiff).
ACCEPTED_FORMATS = {"JPEG", "PNG", "TIFF", "MPO"}


def square_pad(t):
    """Replica of the app's SquarePad, including its rounding of odd differences."""
    h, w = t.size()[-2:]
    m = max(h, w)
    top = int((m - h) / 2)
    left = int((m - w) / 2)
    # torchvision pad order: (left, top, right, bottom)
    return F.pad(t, (left, top, left, top), 0, "constant")


TRANSFORM = transforms.Compose([
    transforms.ToTensor(),
    square_pad,
    transforms.Resize((448, 448), antialias=False),
])


def app_would_read(path: str) -> bool:
    """Replicates the app's file filter: Python 3.8 ``imghdr.what(path) in ('jpeg', 'png', 'tiff')``.

    Python 3.8 imghdr recognises a JPEG only if bytes 6-10 are 'JFIF' or 'Exif',
    so JPEGs without those headers were silently ignored by the desktop app.
    """
    try:
        with open(path, "rb") as f:
            h = f.read(32)
    except OSError:
        return False
    if h[6:10] in (b"JFIF", b"Exif"):
        return True                                   # jpeg
    if h.startswith(b"\211PNG\r\n\032\n"):
        return True                                   # png
    if h[:6] in (b"GIF87a", b"GIF89a"):
        return False                                  # gif (tested before tiff)
    return h[:2] in (b"MM", b"II")                    # tiff


def open_image(path: str):
    """Open an image; return None if it is not a readable JPEG/PNG/TIFF."""
    try:
        img = Image.open(path)
        if img.format not in ACCEPTED_FORMATS:
            return None
        return img.convert("RGB")
    except Exception:
        return None


def to_tensor(img: Image.Image):
    return TRANSFORM(img.convert("RGB")).unsqueeze(0)