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

# 定义维度常量
N, C, H, W = 32, 64, 56, 56
EPS = 1e-6


class EvoNormS0(nn.Module):
    """
    EvoNorm-S0: Evolving Normalization-Activation Layers (Sample-based, no batch dependency)

    公式:
    v = Var(x) = mean(x^2) - mean(x)^2
    y = x / sqrt(v + eps) * gamma + beta
    y = y * sigmoid(x * w)

    其中 gamma, beta, w 是可学习参数
    """

    def __init__(self, num_channels, eps, nonlinear=True):
        super().__init__()
        self.eps = eps
        self.nonlinear = nonlinear  # 是否使用非线性激活

        # 可学习的缩放和偏移参数（类似 BatchNorm）
        self.gamma = nn.Parameter(torch.ones(1, num_channels, 1, 1))
        self.beta = nn.Parameter(torch.zeros(1, num_channels, 1, 1))

        # 非线性门控参数
        if self.nonlinear:
            self.v = nn.Parameter(torch.ones(1, num_channels, 1, 1))

    def forward(self, x: torch.Tensor) -> torch.Tensor:
        # 1. 计算实例级方差
        # var = E[x^2] - E[x]^2
        x_sq_mean = torch.mean(x * x, dim=[2, 3], keepdim=True)
        x_mean = torch.mean(x, dim=[2, 3], keepdim=True)
        var = x_sq_mean - x_mean * x_mean

        # 2. 归一化
        x_normalized = x / torch.sqrt(var + self.eps)

        # 3. 仿射变换
        y = x_normalized * self.gamma + self.beta

        # 4. 非线性门控（可选）
        if self.nonlinear:
            y = y * torch.sigmoid(x * self.v)

        return y


class EvoNormB0(nn.Module):
    """
    EvoNorm-B0: Evolving Normalization-Activation Layers (Batch-based)

    公式:
    Instance Norm: x_in = (x - mean(x)) / sqrt(var(x) + eps)
    Batch Norm stats: rolling_var = momentum * rolling_var + (1-momentum) * batch_var
    y = x_in * gamma + beta
    y = y * sigmoid(x * w)
    """

    def __init__(self, num_channels, eps, momentum=0.1, nonlinear=True):
        super().__init__()
        self.eps = eps
        self.momentum = momentum
        self.nonlinear = nonlinear

        # 可学习参数
        self.gamma = nn.Parameter(torch.ones(1, num_channels, 1, 1))
        self.beta = nn.Parameter(torch.zeros(1, num_channels, 1, 1))

        # 非线性门控参数
        if self.nonlinear:
            self.v = nn.Parameter(torch.ones(1, num_channels, 1, 1))

        # 运行时统计量（用于推理）
        self.register_buffer('running_var', torch.ones(1, num_channels, 1, 1))
        self.register_buffer('num_batches_tracked', torch.tensor(0, dtype=torch.long))

    def forward(self, x: torch.Tensor) -> torch.Tensor:
        if self.training:
            # 训练模式：计算当前批次的统计量
            # 1. 实例归一化
            x_mean = torch.mean(x, dim=[2, 3], keepdim=True)
            x_var = torch.var(x, dim=[2, 3], keepdim=True, unbiased=False)

            # 2. 更新运行统计量（跨批次的方差）
            batch_var = torch.mean(x_var, dim=0, keepdim=True)
            with torch.no_grad():
                self.running_var = (1 - self.momentum) * self.running_var + self.momentum * batch_var
                self.num_batches_tracked += 1

            # 3. 归一化
            x_normalized = (x - x_mean) / torch.sqrt(x_var + self.eps)
        else:
            # 推理模式：使用运行统计量
            x_mean = torch.mean(x, dim=[2, 3], keepdim=True)
            x_normalized = (x - x_mean) / torch.sqrt(self.running_var + self.eps)

        # 4. 仿射变换
        y = x_normalized * self.gamma + self.beta

        # 5. 非线性门控
        if self.nonlinear:
            y = y * torch.sigmoid(x * self.v)

        return y


class Model(nn.Module):
    """
    EvoNorm 模型包装器
    默认使用 EvoNorm-S0（无批次依赖，更适合小批量）
    """

    def __init__(self, evonorm_gamma, evonorm_beta, evonorm_v=None, use_b0=False):
        super().__init__()

        # 选择 EvoNorm 变体
        if use_b0:
            self.evonorm = EvoNormB0(C, EPS, nonlinear=(evonorm_v is not None))
        else:
            self.evonorm = EvoNormS0(C, EPS, nonlinear=(evonorm_v is not None))

        # 初始化参数
        with torch.no_grad():
            self.evonorm.gamma.data.copy_(evonorm_gamma.view(1, C, 1, 1))
            self.evonorm.beta.data.copy_(evonorm_beta.view(1, C, 1, 1))

            if evonorm_v is not None and self.evonorm.nonlinear:
                self.evonorm.v.data.copy_(evonorm_v.view(1, C, 1, 1))

    def forward(self, x: torch.Tensor) -> torch.Tensor:
        return self.evonorm(x)


def get_inputs():
    """生成测试输入"""
    x = torch.randn(N, C, H, W, dtype=torch.float32)
    return [x]


def get_init_inputs():
    """
    生成初始化参数
    返回 [gamma, beta, v]
    """
    evonorm_gamma = torch.ones(1, C, 1, 1)
    evonorm_beta = torch.zeros(1, C, 1, 1)
    evonorm_v = torch.ones(1, C, 1, 1)  # 门控参数
    return [evonorm_gamma, evonorm_beta, evonorm_v]