"""Dynamic convolution operation."""

import torch


def _dynconv2d_naive(
    feature: torch.Tensor,  # (B, C, H, W)
    kernels: torch.Tensor,  # (B, K, C)
) -> torch.Tensor:
    batch_size = feature.shape[0]

    feature = feature.flatten(0, 1)  # (B*C, H, W)
    feature = feature.unsqueeze(0)  # (1, B*C, H, W)

    kernels = kernels.unflatten(-1, (-1, 1, 1))  # (B, K, C, 1, 1)
    kernels = kernels.flatten(0, 1)  # (B*K, C, 1, 1)

    results = torch.nn.functional.conv2d(
        feature, kernels, stride=1, padding=0, groups=batch_size
    )  # (1, B*K, H, W)
    results = results.squeeze(0).unflatten(0, (batch_size, -1))  # (B, K, H, W)
    return results.contiguous()


def _dynconv2d_einsum(
    feature: torch.Tensor,  # (B, C, H, W)
    kernels: torch.Tensor,  # (B, K, C)
) -> torch.Tensor:
    return torch.einsum("bkc, bchw -> bkhw", kernels, feature).contiguous()


def _dynconv2d_matmul(
    feature: torch.Tensor,  # (B, C, H, W)
    kernels: torch.Tensor,  # (B, K, C)
) -> torch.Tensor:
    # (K, C) x (C, H*W) -> (K, H*W)
    return (
        torch.bmm(kernels, feature.flatten(2))
        .unflatten(-1, feature.shape[2:])
        .contiguous()
    )


def dynconv2d(feature: torch.Tensor, kernels: torch.Tensor) -> torch.Tensor:
    """Perform dynamic convolution on the features using the flattened kernels."""
    # Use the fastest implementation available
    return _dynconv2d_matmul(feature, kernels)
