import torch
import torch.nn.functional as F


def _gaussian_kernel_1d(kernel_size: int, sigma: float, device: torch.device, dtype: torch.dtype) -> torch.Tensor:
    radius = kernel_size // 2
    x = torch.arange(-radius, radius + 1, device=device, dtype=dtype)
    kernel = torch.exp(-(x * x) / (2.0 * sigma * sigma))
    return kernel / kernel.sum()


@torch.no_grad()
def build_layerwise_attention_bias(
    masks: torch.Tensor,
    num_layers: int,
    target_tokens: int,
    base_kernel_size: int = 3,
    kernel_growth: int = 2,
    strength: float = 2.0,
    eps: float = 1e-8,
) -> torch.Tensor:
    if masks.ndim == 3:
        masks = masks.unsqueeze(1)
    if masks.ndim != 4 or masks.shape[1] != 1:
        raise ValueError(f"Expected masks shaped (B,1,H,W) or (B,H,W), got {tuple(masks.shape)}")

    masks = masks.float()
    batch_size = masks.shape[0]
    resized = F.interpolate(masks, size=(32, 32), mode="bilinear", align_corners=False).clamp(0.0, 1.0)

    max_kernel = base_kernel_size + max(num_layers, 0) * kernel_growth
    if max_kernel % 2 == 0:
        max_kernel += 1
    pad = max_kernel // 2

    weight_h = torch.zeros((num_layers, 1, 1, max_kernel), device=resized.device, dtype=resized.dtype)
    weight_v = torch.zeros((num_layers, 1, max_kernel, 1), device=resized.device, dtype=resized.dtype)

    for layer_idx in range(num_layers):
        kernel_size = base_kernel_size + (num_layers - layer_idx) * kernel_growth
        if kernel_size % 2 == 0:
            kernel_size += 1
        sigma = max((kernel_size - 1) / 6.0, 1e-3)
        kernel = _gaussian_kernel_1d(kernel_size, sigma, resized.device, resized.dtype)
        start = (max_kernel - kernel_size) // 2
        end = start + kernel_size
        weight_h[layer_idx, 0, 0, start:end] = kernel
        weight_v[layer_idx, 0, start:end, 0] = kernel

    repeated = resized.expand(batch_size, num_layers, 32, 32).contiguous()
    horizontal = F.conv2d(F.pad(repeated, (pad, pad, 0, 0), mode="reflect"), weight_h, groups=num_layers)
    vertical = F.conv2d(F.pad(horizontal, (0, 0, pad, pad), mode="reflect"), weight_v, groups=num_layers)

    min_vals = vertical.amin(dim=(2, 3), keepdim=True)
    max_vals = vertical.amax(dim=(2, 3), keepdim=True)
    normalized = (vertical - min_vals) / (max_vals - min_vals).clamp_min(eps)

    flat = normalized.view(batch_size, num_layers, -1)
    if flat.shape[-1] != target_tokens:
        flat = F.interpolate(flat, size=target_tokens, mode="linear", align_corners=False)
    layerwise_bias = flat.unsqueeze(-2).expand(-1, -1, target_tokens, -1)
    return torch.tril(layerwise_bias) * strength


@torch.no_grad()
def build_legacy_gaussian_attention_bias(
    masks: torch.Tensor,
    num_layers: int,
    target_query_tokens: int,
    target_key_tokens: int,
    base_kernel_size: int = 3,
    kernel_growth: int = 2,
    strength: float = 1.0,
    eps: float = 1e-8,
) -> torch.Tensor:
    if masks.ndim == 3:
        masks = masks.unsqueeze(1)
    if masks.ndim != 4 or masks.shape[1] != 1:
        raise ValueError(f"Expected masks shaped (B,1,H,W) or (B,H,W), got {tuple(masks.shape)}")

    masks = masks.float()
    batch_size = masks.shape[0]
    xmin = masks.amin(dim=(2, 3), keepdim=True)
    xmax = masks.amax(dim=(2, 3), keepdim=True)
    normalized_masks = (masks - xmin) / (xmax - xmin).clamp_min(eps)
    resized = F.interpolate(normalized_masks, size=(32, 32), mode="bilinear", align_corners=False)

    kernel_sizes = []
    for layer_idx in range(num_layers, 0, -1):
        kernel_size = base_kernel_size + layer_idx * kernel_growth
        if kernel_size % 2 == 0:
            kernel_size += 1
        kernel_sizes.append(max(kernel_size, 1))

    max_kernel = max(kernel_sizes)
    pad = max_kernel // 2
    weight_h = torch.zeros((num_layers, 1, 1, max_kernel), device=resized.device, dtype=resized.dtype)
    weight_v = torch.zeros((num_layers, 1, max_kernel, 1), device=resized.device, dtype=resized.dtype)

    for layer_idx, kernel_size in enumerate(kernel_sizes):
        sigma = max((kernel_size - 1) / 6.0, 1e-3)
        kernel = _gaussian_kernel_1d(kernel_size, sigma, resized.device, resized.dtype)
        start = (max_kernel - kernel_size) // 2
        end = start + kernel_size
        weight_h[layer_idx, 0, 0, start:end] = kernel
        weight_v[layer_idx, 0, start:end, 0] = kernel

    repeated = resized.expand(batch_size, num_layers, 32, 32).contiguous()
    horizontal = F.conv2d(F.pad(repeated, (pad, pad, 0, 0), mode="reflect"), weight_h, groups=num_layers)
    vertical = F.conv2d(F.pad(horizontal, (0, 0, pad, pad), mode="reflect"), weight_v, groups=num_layers)

    min_vals = vertical.amin(dim=(2, 3), keepdim=True)
    max_vals = vertical.amax(dim=(2, 3), keepdim=True)
    normalized = (vertical - min_vals) / (max_vals - min_vals).clamp_min(eps)

    flat = normalized.view(batch_size, num_layers, -1)
    if flat.shape[-1] != target_key_tokens:
        flat = F.interpolate(flat, size=target_key_tokens, mode="linear", align_corners=False)
    return flat.unsqueeze(-2).expand(-1, -1, target_query_tokens, -1) * strength


__all__ = ["build_layerwise_attention_bias", "build_legacy_gaussian_attention_bias"]