import torch
import torch.nn as nn
import torch.nn.functional as F

BATCH = 16
CHANNELS = 32
HEIGHT = 256
WIDTH = 256


KERNEL_SIZE = 5
SIGMA_COLOR = 10.0
SIGMA_SPACE = 5.0

class Model(nn.Module):
    def __init__(self):
        super().__init__()
        self.k = KERNEL_SIZE
        self.pad = KERNEL_SIZE // 2
        self.sigma_c = SIGMA_COLOR
        self.sigma_s = SIGMA_SPACE

    def forward(self, x: torch.Tensor) -> torch.Tensor:
        B, C, H, W = x.shape
        x_d = x.double()

        dy = torch.arange(-self.pad, self.pad+1, dtype=torch.float64, device=x.device)
        dx = torch.arange(-self.pad, self.pad+1, dtype=torch.float64, device=x.device)
        mesh_y, mesh_x = torch.meshgrid(dy, dx, indexing='ij')
        spatial_dist_sq = mesh_y**2 + mesh_x**2
        spatial_weight = torch.exp(-spatial_dist_sq / (2 * self.sigma_s**2))
        spatial_weight = spatial_weight.view(1, 1, self.k, self.k, 1, 1)
        
 
        x_pad = F.pad(x_d, (self.pad, self.pad, self.pad, self.pad), mode='replicate')
        
        patches = F.unfold(x_pad, kernel_size=self.k)
        patches = patches.view(B, C, self.k, self.k, H, W)

        center = x_d.unsqueeze(2).unsqueeze(2)
        
        color_diff_sq = (patches - center) ** 2
        color_weight = torch.exp(-color_diff_sq / (2 * self.sigma_c**2))

        total_weight = spatial_weight * color_weight
        weighted_sum = (patches * total_weight).sum(dim=(2, 3))
        weight_sum = total_weight.sum(dim=(2, 3))
        
        out = weighted_sum / weight_sum
        
        return out.float()

def get_inputs():
    x = torch.randint(0, 256, size=(BATCH, CHANNELS, HEIGHT, WIDTH), device='cuda').float()
    return [x]

def get_init_inputs():
    return []