# switchablenorm_torch.py 
import torch
import torch.nn as nn
import torch.nn.functional as F

# 定义输入尺寸 (N, C, H, W)
N, C, H, W = 16, 64, 32, 32
EPS = 1e-5

class SN(nn.Module):

    def __init__(self, num_channels, eps):
        super().__init__()
        self.eps = eps
        
        self.weight = nn.Parameter(torch.ones(1, num_channels, 1, 1)) # gamma
        self.bias = nn.Parameter(torch.zeros(1, num_channels, 1, 1))   # beta

        self.w_in = nn.Parameter(torch.ones(num_channels))
        self.w_ln = nn.Parameter(torch.ones(num_channels))
        self.w_bn = nn.Parameter(torch.ones(num_channels))
        
        self.register_buffer('running_mean', torch.zeros(num_channels))
        self.register_buffer('running_var', torch.ones(num_channels))

    def forward(self, x: torch.Tensor) -> torch.Tensor:
    
        ln_mean = x.mean(dim=[1, 2, 3], keepdim=True)
        ln_var = x.var(dim=[1, 2, 3], keepdim=True)  
        

        in_mean = x.mean(dim=[2, 3], keepdim=True) 
        in_var = x.var(dim=[2, 3], keepdim=True)  
        

        bn_mean = self.running_mean.view(1, C, 1, 1)
        bn_var = self.running_var.view(1, C, 1, 1)

        w_sum = self.w_in.abs() + self.w_ln.abs() + self.w_bn.abs()
        w_in_norm = (self.w_in.abs() / w_sum).view(1, C, 1, 1)
        w_ln_norm = (self.w_ln.abs() / w_sum).view(1, C, 1, 1)
        w_bn_norm = (self.w_bn.abs() / w_sum).view(1, C, 1, 1)

        mean = w_in_norm * in_mean + w_ln_norm * ln_mean + w_bn_norm * bn_mean
        

        var_in_M2 = in_var + in_mean.pow(2)
        var_ln_M2 = ln_var + ln_mean.pow(2)
        var_bn_M2 = bn_var + bn_mean.pow(2)
        
        aggregated_var_M2 = w_in_norm * var_in_M2 + w_ln_norm * var_ln_M2 + w_bn_norm * var_bn_M2
        var = aggregated_var_M2 - mean.pow(2)
        

        x_norm = (x - mean) / torch.sqrt(var + self.eps)
        return x_norm * self.weight + self.bias


class Model(nn.Module):
    def __init__(self, weight, bias, w_in, w_ln, w_bn):
        super().__init__()
        self.sn = SN(C, EPS)
        
        with torch.no_grad():
            self.sn.weight.data.copy_(weight)
            self.sn.bias.data.copy_(bias)
            self.sn.w_in.data.copy_(w_in.squeeze())
            self.sn.w_ln.data.copy_(w_ln.squeeze())
            self.sn.w_bn.data.copy_(w_bn.squeeze())
            
    def forward(self, x: torch.Tensor) -> torch.Tensor:
        return self.sn(x)


def get_inputs():

    x = torch.randn(N, C, H, W, dtype=torch.float32)
    return [x]


def get_init_inputs():

    w_in = torch.ones(C)
    w_ln = torch.ones(C)
    w_bn = torch.ones(C)
    weight = torch.ones(1, C, 1, 1)
    bias = torch.zeros(1, C, 1, 1)
    return [weight, bias, w_in, w_ln, w_bn]