import torch
import torch.nn as nn

class GELUDropoutTorchModel(nn.Module):
    """PyTorch原生GELU+Dropout实现"""
    
    def __init__(self, p=0.1):
        super(GELUDropoutTorchModel, self).__init__()
        self.gelu = nn.GELU()
        self.dropout = nn.Dropout(p=p)
    
    def forward(self, x):
        x = self.gelu(x)
        x = self.dropout(x)
        return x

def get_init_inputs():
    """获取模型初始化参数"""
    return [0.1]  # dropout概率

def get_inputs():
    """获取模型输入数据"""
    torch.manual_seed(42)
    return [torch.randn(128, 256)]

def test_gelu_dropout():
    """测试GELU+Dropout功能"""
    print("=" * 60)
    print("GELU+Dropout功能测试")
    print("=" * 60)
    
    # 创建测试数据
    torch.manual_seed(42)
    batch_size, hidden_size = 128, 256
    input_tensor = torch.randn(batch_size, hidden_size)
    
    # 创建模型
    model = GELUDropoutTorchModel(dropout_prob=0.1)
    
    # 训练模式测试
    model.train()
    train_output = model(input_tensor)
    print(f"训练模式输出形状: {train_output.shape}")
    print(f"训练模式输出范围: [{train_output.min().item():.3f}, {train_output.max().item():.3f}]")
    
    # 推理模式测试
    model.eval()
    eval_output = model(input_tensor)
    print(f"推理模式输出形状: {eval_output.shape}")
    print(f"推理模式输出范围: [{eval_output.min().item():.3f}, {eval_output.max().item():.3f}]")
    
    # 验证Dropout效果
    zero_count_train = (train_output == 0).sum().item()
    zero_count_eval = (eval_output == 0).sum().item()
    
    print(f"训练模式零值比例: {zero_count_train / train_output.numel():.3f}")
    print(f"推理模式零值比例: {zero_count_eval / eval_output.numel():.3f}")
    
    # 验证GELU激活函数
    gelu_only = 0.5 * input_tensor * (1 + torch.tanh(0.7978845608028654 * (input_tensor + 0.044715 * torch.pow(input_tensor, 3))))
    print(f"纯GELU输出范围: [{gelu_only.min().item():.3f}, {gelu_only.max().item():.3f}]")
    
    return train_output, eval_output

def main():
    """主测试函数"""
    try:
        train_output, eval_output = test_gelu_dropout()
        
        print("\n" + "=" * 60)
        print("测试总结")
        print("=" * 60)
        print("✅ GELU+Dropout功能验证通过")
        print("✅ 训练/推理模式切换正常")
        print("✅ Dropout效果符合预期")
        
    except Exception as e:
        print(f"测试过程中出现错误: {e}")

if __name__ == "__main__":
    main()