import torch
import torch.nn as nn

class Model(nn.Module):
    """
    原始模型：InstanceNorm + Dropout
    """
    def __init__(self, num_features=64, eps=1e-5, affine=True, dropout_p=0.1, track_running_stats=False):
        super(Model, self).__init__()
        self.num_features = num_features
        self.eps = eps
        self.affine = affine
        self.dropout_p = dropout_p
        self.track_running_stats = track_running_stats

        # 创建InstanceNorm层
        self.instance_norm = nn.InstanceNorm2d(
            num_features=num_features,
            eps=eps,
            affine=affine,
            track_running_stats=track_running_stats
        )
        
        # 创建Dropout层
        self.dropout = nn.Dropout(p=dropout_p)

    def forward(self, x: torch.Tensor) -> torch.Tensor:
        """
        原始实现：先InstanceNorm，再Dropout
        
        Args:
            x (torch.Tensor): 输入张量 shape [B, C, H, W]
        
        Returns:
            torch.Tensor: Dropout(InstanceNorm(x))
        """
        x_normalized = self.instance_norm(x)
        return self.dropout(x_normalized)

class ModelNew(torch.nn.Module):
    """
    融合模型：直接实现InstanceNorm + Dropout
    """
    def __init__(self, num_features=64, eps=1e-5, affine=True, dropout_p=0.1, track_running_stats=False):
        super(ModelNew, self).__init__()
        self.num_features = num_features
        self.eps = eps
        self.affine = affine
        self.dropout_p = dropout_p
        self.track_running_stats = False  # 强制为False以支持CUDA实现

        if affine:
            self.weight = torch.nn.Parameter(torch.ones(num_features))
            self.bias = torch.nn.Parameter(torch.zeros(num_features))
        else:
            self.register_parameter('weight', None)
            self.register_parameter('bias', None)

    def forward(self, x: torch.Tensor) -> torch.Tensor:
        """
        融合实现：在CUDA kernel中直接完成InstanceNorm + Dropout
        
        Args:
            x (torch.Tensor): 输入张量 shape [B, C, H, W]
        
        Returns:
            torch.Tensor: Dropout(InstanceNorm(x))
        """
        # 这个将在CUDA中实现
        pass

# 测试参数
batch_size = 128
num_features = 64
height = 128
width = 128
dropout_p = 0.1

def get_inputs():
    """生成测试输入"""
    x = torch.randn(batch_size, num_features, height, width)
    return [x]

def get_init_inputs():
    """获取初始化参数"""
    return [num_features]
