"""Load images for BA-T inference from image paths and/or folders.

Every image goes through the same deterministic preprocessing as the training /
evaluation dataloaders, with the principal point assumed at the image centre:

  1. crop to a window centred on the image centre (a no-op for even sizes);
  2. Lanczos-rescale so that the image just covers the target resolution
     (the short side becomes 224 for the released 224x224 model);
  3. centre-crop to the target resolution;
  4. normalize RGB to [-1, 1] (``ImgNorm``).

A view is a dict ``{'img': (3,H,W) float tensor, 'true_shape': (2,) int32 tensor
(H, W), 'path': str, 'crop': dict}``; ``crop`` records the transform so that
pixels can be mapped back to the original image (``to_original_pixels``).
``collate_views`` stacks views into the batched dicts that ``BAT.forward``
expects.
"""
import os
import re

import numpy as np
import PIL.Image
import PIL.ImageOps
import torch

from .utils.cropping import rescale_image, crop_image, center_crop_bbox
from .utils.image import ImgNorm

IMG_EXTENSIONS = ('.png', '.jpg', '.jpeg', '.bmp', '.tif', '.tiff', '.webp')
DEFAULT_RESOLUTION = (224, 224)  # (W, H) the released model was trained at


def _natural_key(s):
    return [int(t) if t.isdigit() else t.lower() for t in re.split(r'(\d+)', s)]


def collect_image_paths(inputs):
    """Expand a path or a list of paths (image files and/or folders) into a list of image files.

    Folders are scanned non-recursively and sorted in natural order (frame2 < frame10);
    explicitly listed files keep the given order.
    """
    if isinstance(inputs, (str, os.PathLike)):
        inputs = [inputs]
    paths = []
    for p in inputs:
        p = os.fspath(p)
        if os.path.isdir(p):
            files = [os.path.join(p, f) for f in os.listdir(p) if f.lower().endswith(IMG_EXTENSIONS)]
            if not files:
                raise FileNotFoundError(f'no images ({", ".join(IMG_EXTENSIONS)}) found in folder {p}')
            paths += sorted(files, key=_natural_key)
        elif os.path.isfile(p):
            paths.append(p)
        else:
            raise FileNotFoundError(f'{p} is neither an image file nor a folder')
    return paths


def _parse_resolution(resolution):
    if isinstance(resolution, int):
        return resolution, resolution
    width, height = resolution
    assert width % 16 == 0 and height % 16 == 0, f'resolution {resolution} must be a multiple of the patch size 16'
    return int(width), int(height)


def preprocess_image(image, resolution=DEFAULT_RESOLUTION):
    """Crop + resize a PIL image like the training dataloader (deterministic, no augmentation).

    Returns the processed PIL image and a dict describing the transform.
    """
    width, height = _parse_resolution(resolution)
    image = image.convert('RGB')
    W, H = image.size

    # 1. crop centred on the principal point (assumed at the image centre)
    cx, cy = int(W/2), int(H/2)
    min_margin_x = min(cx, W-cx)
    min_margin_y = min(cy, H-cy)
    l1, t1 = cx - min_margin_x, cy - min_margin_y
    image = crop_image(image, (l1, t1, cx + min_margin_x, cy + min_margin_y))

    # the target resolution is landscape; use it transposed for portrait images
    target = (width, height)
    if width != height and image.size[1] > 1.1 * image.size[0]:
        target = (height, width)

    # 2. high-quality Lanczos down-scaling, 3. centre crop
    image, scale = rescale_image(image, target)
    l2, t2, r2, b2 = center_crop_bbox(image.size, target)
    image = crop_image(image, (l2, t2, r2, b2))

    crop = {'orig_size': (W, H), 'offset1': (int(l1), int(t1)), 'scale': float(scale),
            'offset2': (int(l2), int(t2)), 'size': image.size}
    return image, crop


def to_original_pixels(uv, crop):
    """Map pixel coordinates (..., 2) [u=column, v=row] of a preprocessed image back to
    the original image (pixel centres at integer coordinates)."""
    uv = np.asarray(uv, dtype=np.float64)
    off2 = np.asarray(crop['offset2'], dtype=np.float64)
    off1 = np.asarray(crop['offset1'], dtype=np.float64)
    return (uv + off2 + 0.5) / crop['scale'] - 0.5 + off1


def load_image(path, resolution=DEFAULT_RESOLUTION):
    """Read one image file (EXIF orientation applied) and preprocess it into a view dict."""
    with PIL.Image.open(path) as im:
        image = PIL.ImageOps.exif_transpose(im).convert('RGB')
    image, crop = preprocess_image(image, resolution)
    width, height = image.size
    return {'img': ImgNorm(image),
            'true_shape': torch.tensor((height, width), dtype=torch.int32),
            'path': os.fspath(path),
            'crop': crop}


class ImageViews(torch.utils.data.Dataset):
    """Dataset over image paths / folders, yielding preprocessed view dicts."""

    def __init__(self, inputs, resolution=DEFAULT_RESOLUTION):
        self.paths = collect_image_paths(inputs)
        self.resolution = resolution

    def __len__(self):
        return len(self.paths)

    def __getitem__(self, idx):
        return load_image(self.paths[idx], self.resolution)


def collate_views(views):
    """Stack a list of view dicts into one batched view: 'img' (B,3,H,W), 'true_shape' (B,2).

    All views must have the same shape (always true for square resolutions).
    """
    return {'img': torch.stack([v['img'] for v in views]),
            'true_shape': torch.stack([v['true_shape'] for v in views]),
            'path': [v['path'] for v in views],
            'crop': [v['crop'] for v in views]}


def load_views(inputs, resolution=DEFAULT_RESOLUTION):
    """Load images from paths / folders as a list of single-image batched views (B=1)."""
    ds = ImageViews(inputs, resolution)
    return [collate_views([ds[i]]) for i in range(len(ds))]


def denormalize(img):
    """(3,H,W) or (B,3,H,W) tensor in [-1,1] -> uint8 numpy (H,W,3) or (B,H,W,3)."""
    img = ((img.detach().cpu().float() + 1.0) * 127.5).round().clamp(0, 255).to(torch.uint8)
    return img.movedim(-3, -1).numpy()