"""In-plane rotation of photos and maps, for test-time rotation averaging.

Conventions (match the training augmentation, `MarigoldV2Dataset.__getitem__`):
  * a positive angle rotates the picture counter-clockwise (like PIL's rotate);
  * the maps are `[0:3] albedo | [3] depth | [4:7] normal`, with the normal in the
    camera frame x right, y UP, z toward the viewer. Albedo and depth are
    invariant under a roll about the optical axis, so only their positions move;
    the normal's (nx, ny) rotate with the picture by R(+theta) =
    [[c, -s], [s, c]].

`rotate_image` resamples with a reflected border (no empty corners, which the
model never saw); `unrotate_maps` returns the maps back in the original frame
together with the mask of pixels the rotated canvas actually covered.
"""
from __future__ import annotations

import math

import torch
import torch.nn.functional as F


def _theta(angle_deg: float, device, dtype) -> torch.Tensor:
    """affine_grid matrix [1, 2, 3]: output pixel -> input sample position, so
    that the picture content turns counter-clockwise by `angle_deg` (image
    coordinates: x right, y down)."""
    a = math.radians(angle_deg)
    c, s = math.cos(a), math.sin(a)
    return torch.tensor([[[c, -s, 0.0], [s, c, 0.0]]], device=device,
                        dtype=dtype)


def rotate_image(x: torch.Tensor, angle_deg: float, mode: str = "bicubic",
                 padding: str = "reflection") -> torch.Tensor:
    """Rotate [B, C, H, W] counter-clockwise by `angle_deg` about the centre.

    Sampled in float32 with autocast off: the sampling grid in bf16 has only ~3
    significant digits, which at 448 px is a shift of a pixel or more (found
    when the same code under `torch.autocast` gave outputs 0.5 off)."""
    b = x.shape[0]
    with torch.autocast(device_type=x.device.type, enabled=False):
        xf = x.float()
        theta = _theta(angle_deg, x.device, torch.float32).expand(b, -1, -1)
        grid = F.affine_grid(theta, list(xf.shape), align_corners=False)
        out = F.grid_sample(xf, grid, mode=mode, padding_mode=padding,
                            align_corners=False)
    return out.to(x.dtype)


def _reorient_normals(maps: torch.Tensor, angle_deg: float) -> torch.Tensor:
    """Apply R(angle) to the normal's (nx, ny); nz is unchanged."""
    a = math.radians(angle_deg)
    c, s = math.cos(a), math.sin(a)
    out = maps.clone()
    nx, ny = maps[:, 4], maps[:, 5]
    out[:, 4] = c * nx - s * ny
    out[:, 5] = s * nx + c * ny
    return out


def rotate_maps(maps: torch.Tensor, angle_deg: float):
    """Rotate 7ch maps counter-clockwise and reorient the normals.

    Returns (maps, valid) where `valid` [B, 1, H, W] marks canvas pixels that
    came from inside the original frame (the rest are zero)."""
    ones = torch.ones_like(maps[:, :1])
    r = rotate_image(maps, angle_deg, "bilinear", "zeros")
    valid = rotate_image(ones, angle_deg, "bilinear", "zeros") > 0.999
    return _reorient_normals(r, angle_deg), valid


def unrotate_maps(maps: torch.Tensor, angle_deg: float):
    """Maps predicted on a photo rotated by `angle_deg`, brought back to the
    original frame (rotate by -angle, normals reoriented by R(-angle)).

    Returns (maps, valid): `valid` marks original-frame pixels the rotated
    canvas covered (the corners are lost for rotations of a few degrees)."""
    return rotate_maps(maps, -angle_deg)