"""Test-time averaging for ShadeNet-5, copied from the repository's
`detail_diffuser/inference/core.py` (flip_average / tta_average / unit_normals)
so the Space renders exactly what the model card reports.

`refine_fn(x)` takes a [1, 3, H, W] photo in [-1, 1] and returns the 7 maps
[1, 7, H, W] in [-1, 1]: albedo 0:3, depth 3, normal 4:7.
"""
import itertools

import torch

from rotation import rotate_image, unrotate_maps

NORMAL_X = 4          # first channel of the normal (x right, y up, z to viewer)


def unit_normals(maps, eps=1e-6):
    out = maps.clone()
    n = out[:, NORMAL_X:NORMAL_X + 3]
    out[:, NORMAL_X:NORMAL_X + 3] = n / n.norm(dim=1, keepdim=True).clamp_min(eps)
    return out


@torch.no_grad()
def tta_average(refine_fn, x, flip=False, rot_deg=0.0, progress=None):
    """Per-pixel mean over the mirror (optional) x rotations (0, -rot, +rot).

    Each pass: mirror -> rotate the photo (reflected border) -> predict ->
    rotate the maps back, reorienting the normals -> un-mirror. A pass only
    counts at the pixels its rotated canvas covers.
    """
    angles = (0.0, -float(rot_deg), float(rot_deg)) if rot_deg else (0.0,)
    flips = (False, True) if flip else (False,)
    passes = list(itertools.product(flips, angles))
    num = den = None
    for i, (f, a) in enumerate(passes):
        if progress is not None:
            progress(i, len(passes))
        c = x.flip(-1) if f else x
        if a:
            c = rotate_image(c, a, "bicubic", "reflection").clamp(-1.0, 1.0)
        p = refine_fn(c).float()
        valid = torch.ones_like(p[:, :1])
        if a:
            p, v = unrotate_maps(p, a)
            valid = v.to(p.dtype)
        if f:
            p = p.flip(-1).clone()
            p[:, NORMAL_X] = -p[:, NORMAL_X]
            valid = valid.flip(-1)
        num = p * valid if num is None else num + p * valid
        den = valid if den is None else den + valid
    return num / den.clamp_min(1.0)