"""Fit a whole photo into the model's square input instead of cropping it.

The model takes 448x448. `letterbox` scales the photo so its LONG side is 448
and fills the rest of the square by mirroring the photo's own edge (the model
never saw flat-coloured bars, and the test-time rotation already feeds it
reflected borders). After the model runs, `crop_to_box` cuts the padding off
the maps again, so nothing of the photo is lost and nothing invented is shown.
"""
import numpy as np
import torch
import torch.nn.functional as F
from PIL import Image


def letterbox(img: Image.Image, size: int):
    """Return (square PIL image, (left, top, w, h) of the real photo in it)."""
    img = img.convert("RGB")
    w, h = img.size
    scale = size / float(max(w, h))
    nw, nh = max(1, int(round(w * scale))), max(1, int(round(h * scale)))
    arr = np.asarray(img.resize((nw, nh), Image.BICUBIC))
    left, top = (size - nw) // 2, (size - nh) // 2
    pad = ((top, size - nh - top), (left, size - nw - left), (0, 0))
    # 'symmetric' (edge pixel included) is safe for any pad width, even when
    # the padding is longer than the photo itself (a very elongated panorama).
    out = np.pad(arr, pad, mode="symmetric")
    return Image.fromarray(out), (left, top, nw, nh)


def crop_to_box(maps: torch.Tensor, box) -> torch.Tensor:
    """Cut [B, C, S, S] maps down to the photo's own pixels."""
    left, top, w, h = box
    return maps[:, :, top:top + h, left:left + w]


def resize_maps(maps: torch.Tensor, width: int, height: int) -> torch.Tensor:
    """Resize 7-channel maps for display and keep the normals unit length."""
    out = F.interpolate(maps.float(), size=(height, width), mode="bicubic",
                        align_corners=False).clamp(-1.0, 1.0)
    n = out[:, 4:7]
    out[:, 4:7] = n / n.norm(dim=1, keepdim=True).clamp_min(1e-6)
    return out


def display_size(w: int, h: int, longest: int = 1024):
    """Output size: the photo's own shape, long side at most `longest`."""
    s = min(1.0, longest / float(max(w, h)))
    return max(1, int(round(w * s))), max(1, int(round(h * s)))