"""Loss helpers for context-aware inpainting distillation."""

from __future__ import annotations

import torch
import torch.nn.functional as F


def min_snr_weights(
    alphas_cumprod: torch.Tensor,
    timesteps: torch.Tensor,
    *,
    gamma: float,
    prediction_type: str,
) -> torch.Tensor:
    """Return the Min-SNR weighting from diffusion fine-tuning literature."""

    alpha = alphas_cumprod.to(device=timesteps.device, dtype=torch.float32)[timesteps]
    snr = alpha / (1.0 - alpha).clamp_min(1e-8)
    clipped = torch.minimum(snr, torch.full_like(snr, gamma))
    if prediction_type == "epsilon":
        return clipped / snr.clamp_min(1e-8)
    if prediction_type == "v_prediction":
        return clipped / (snr + 1.0)
    raise ValueError(f"Unsupported prediction type: {prediction_type}")


def spatial_loss_weights(
    mask: torch.Tensor,
    *,
    context_weight: float,
    masked_weight: float,
    boundary_weight: float,
    boundary_radius: int = 2,
) -> torch.Tensor:
    """Emphasize regenerated pixels and the context-sensitive mask boundary."""

    if mask.ndim != 4 or mask.shape[1] != 1:
        raise ValueError(f"Expected BCHW single-channel mask, got {tuple(mask.shape)}")
    if boundary_radius < 1:
        raise ValueError("boundary_radius must be positive")
    binary = (mask >= 0.5).to(dtype=torch.float32)
    kernel = boundary_radius * 2 + 1
    dilated = F.max_pool2d(binary, kernel_size=kernel, stride=1, padding=boundary_radius)
    eroded = 1.0 - F.max_pool2d(
        1.0 - binary,
        kernel_size=kernel,
        stride=1,
        padding=boundary_radius,
    )
    boundary = (dilated - eroded).clamp(0.0, 1.0)
    weights = torch.full_like(binary, float(context_weight))
    weights = weights + binary * (float(masked_weight) - float(context_weight))
    weights = weights + boundary * float(boundary_weight)
    return weights / weights.mean(dim=(1, 2, 3), keepdim=True).clamp_min(1e-8)


def weighted_mse(
    prediction: torch.Tensor,
    target: torch.Tensor,
    *,
    spatial_weights: torch.Tensor,
    sample_weights: torch.Tensor,
) -> torch.Tensor:
    """Compute channel-averaged, spatially and per-sample weighted MSE."""

    if prediction.shape != target.shape:
        raise ValueError("prediction and target must have identical shapes")
    per_pixel = (prediction.float() - target.float()).square().mean(dim=1, keepdim=True)
    per_sample = (per_pixel * spatial_weights.float()).mean(dim=(1, 2, 3))
    return (per_sample * sample_weights.float()).mean()