"""Simple image augmentations for training."""

import numpy as np
import torch


def brightness_augment(
    imgs: torch.Tensor,
    coords: torch.Tensor,
    masks: torch.Tensor,
    *,
    rng: np.random.Generator,
    shift_range: float = 0.1,
) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
    """Random additive brightness shift."""
    shift = rng.uniform(-shift_range, shift_range)
    return imgs + shift, coords, masks


def flip_augment(
    imgs: torch.Tensor,
    coords: torch.Tensor,
    masks: torch.Tensor,
    *,
    rng: np.random.Generator,
) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
    """Random spatial flip while preserving valid node coordinates."""

    flip_mask = rng.random(3) < 0.5
    dims_to_flip = [1 + dim for dim, flip in enumerate(flip_mask) if flip]

    if not dims_to_flip:
        return imgs, coords, masks

    imgs = imgs.flip(dims=dims_to_flip)

    coords = coords.clone()
    shape = imgs.shape[1:]

    if coords.ndim == 2:
        # Actual observed case:
        # coords: (W, 3)
        # masks:  (W, M)
        valid = masks.any(dim=1)

        for dim in range(3):
            if flip_mask[dim]:
                coords[valid, dim] = (
                    shape[dim] - coords[valid, dim] - 1
                )

    elif coords.ndim == 3:
        # General case:
        # coords: (W, M, 3)
        # masks:  (W, M)
        valid = masks.bool()

        for dim in range(3):
            if flip_mask[dim]:
                coords[..., dim] = torch.where(
                    valid,
                    shape[dim] - coords[..., dim] - 1,
                    coords[..., dim],
                )

    else:
        raise ValueError(
            f"Unexpected coords shape: {coords.shape}"
        )

    return imgs, coords, masks


# === extra robustness augments =======================================

import torch.nn.functional as _Fg


def noise_augment(imgs, coords, masks, *, rng, sigma_frac=0.5):
    s = rng.uniform(0.0, sigma_frac) * float(imgs.std())
    return imgs + torch.randn_like(imgs) * s, coords, masks


def contrast_augment(imgs, coords, masks, *, rng, lo=0.6, hi=1.6):
    c = rng.uniform(lo, hi)
    mu = float(imgs.mean())
    return (imgs - mu) * c + mu, coords, masks


def gamma_augment(imgs, coords, masks, *, rng, lo=0.6, hi=1.6):
    g = rng.uniform(lo, hi)
    return imgs.clamp(min=0).pow(g), coords, masks


def blur_augment(imgs, coords, masks, *, rng, max_k=3):
    k = int(rng.integers(1, max_k + 1)) | 1

    if k == 1:
        return imgs, coords, masks

    x = imgs.unsqueeze(1).float()
    p = k // 2

    x = _Fg.avg_pool3d(
        _Fg.pad(
            x,
            (p, p, p, p, p, p),
            mode="replicate"
        ),
        k,
        stride=1,
    )

    return x[:, 0].to(imgs.dtype), coords, masks


generalization_augments = [
    brightness_augment,
    flip_augment,
    noise_augment,
    contrast_augment,
    gamma_augment,
    blur_augment,
]