"""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_passes(refine_fn, x, flip=False, rot_deg=0.0, progress=None):
    """The passes of the mirror (optional) x rotations (0, -rot, +rot) recipe,
    each brought back to the photo's frame, as a list of (maps, valid).

    Each pass: mirror -> rotate the photo (reflected border) -> predict ->
    rotate the maps back, reorienting the normals -> un-mirror. `valid` marks
    the pixels the pass covers (a rotated canvas loses the corners).
    """
    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))
    out = []
    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)
        out.append((p, valid))
    return out


def mean_of(passes):
    """Per-pixel mean over the passes that cover the pixel (the card's recipe)."""
    num = den = None
    for p, valid in passes:
        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)


def median_of(passes):
    """Per-pixel median over the passes that cover the pixel (the average of the
    two middle values when an even number cover it). The unrotated pass covers
    everything, so every pixel has at least one value."""
    P = torch.stack([p for p, _ in passes])                   # [n, B, C, H, W]
    V = torch.stack([v for _, v in passes]) > 0.5             # [n, B, 1, H, W]
    P = torch.where(V.expand_as(P), P, torch.full_like(P, float("nan")))
    return torch.nanquantile(P, 0.5, dim=0)


@torch.no_grad()
def tta_average(refine_fn, x, flip=False, rot_deg=0.0, progress=None):
    """Per-pixel mean over the recipe's passes (see `tta_passes`)."""
    return mean_of(tta_passes(refine_fn, x, flip, rot_deg, progress))