import torch
import torch.nn as nn


BATCH = 64
CHANNELS = 64
HEIGHT = 128
WIDTH = 128
EPS = 1e-5
MOMENTUM = 0.1
R_MAX = 3.0
D_MAX = 5.0

class Model(nn.Module):
    def __init__(self):
        super().__init__()
        self.gamma = nn.Parameter(torch.ones(CHANNELS, device='cuda'))
        self.beta = nn.Parameter(torch.zeros(CHANNELS, device='cuda'))
        self.register_buffer('running_mean', torch.zeros(CHANNELS, device='cuda'))
        self.register_buffer('running_var', torch.ones(CHANNELS, device='cuda'))
        self.training = True 

    def forward(self, x: torch.Tensor) -> torch.Tensor:
        if not self.training:
            scale = self.gamma / (self.running_var.add(EPS).sqrt())
            bias = self.beta - self.running_mean * scale
            return x * scale.view(1, -1, 1, 1) + bias.view(1, -1, 1, 1)

        batch_mean = x.mean(dim=(0, 2, 3))
        batch_var = x.var(dim=(0, 2, 3), unbiased=False)
        batch_std = (batch_var + EPS).sqrt()


        with torch.no_grad():
            self.running_mean.lerp_(batch_mean, MOMENTUM)
            self.running_var.lerp_(batch_var, MOMENTUM)


        with torch.no_grad():
            running_std = (self.running_var + EPS).sqrt()
            r = (batch_std / running_std)
            d = ((batch_mean - self.running_mean) / running_std)
            r = r.clamp(1.0 / R_MAX, R_MAX)
            d = d.clamp(-D_MAX, D_MAX)
        

        x_hat = (x - batch_mean.view(1, -1, 1, 1)) / batch_std.view(1, -1, 1, 1)
        y = x_hat * r.view(1, -1, 1, 1) + d.view(1, -1, 1, 1)
        return self.gamma.view(1, -1, 1, 1) * y + self.beta.view(1, -1, 1, 1)

def get_inputs():
    x = torch.randn(BATCH, CHANNELS, HEIGHT, WIDTH, device='cuda', dtype=torch.float32)
    return [x]

def get_init_inputs():
    return []