import torch
import torch.nn as nn

class EraWNet(nn.Module):
    """
    EraW-Net: Enhance-Refine-Align W-Net
    
    双编码器结构：
    - Encoder 1 (Enhance): 增强面部/视线相关特征
    - Decoder (Refine): 精化注意力图
    - Encoder 2 (Align): 对齐场景语义
    """
    
    def __init__(self):
        super().__init__()
        
        # Encoder 1: Enhance
        self.encoder1 = nn.Sequential(
            nn.Conv2d(3, 64, 3, padding=1), nn.ReLU(),
            nn.Conv2d(64, 64, 3, stride=2, padding=1), nn.ReLU(),
            nn.Conv2d(64, 128, 3, padding=1), nn.ReLU(),
            nn.Conv2d(128, 128, 3, stride=2, padding=1), nn.ReLU(),
            nn.Conv2d(128, 256, 3, padding=1), nn.ReLU(),
        )
        
        # Decoder: Refine
        self.decoder = nn.Sequential(
            nn.ConvTranspose2d(256, 128, 3, stride=2, padding=1, output_padding=1), nn.ReLU(),
            nn.ConvTranspose2d(128, 64, 3, stride=2, padding=1, output_padding=1), nn.ReLU(),
            nn.Conv2d(64, 1, 1), nn.Sigmoid()  # 注意力图
        )
        
        # Encoder 2: Align
        self.encoder2 = nn.Sequential(
            nn.Conv2d(1, 32, 3, padding=1), nn.ReLU(),
            nn.Conv2d(32, 64, 3, stride=2, padding=1), nn.ReLU(),
            nn.Conv2d(64, 128, 3, stride=2, padding=1), nn.ReLU(),
            nn.AdaptiveAvgPool2d((1, 1)),
        )
        
        # 场景对齐头
        self.align_head = nn.Sequential(
            nn.Linear(128, 64), nn.ReLU(),
            nn.Linear(64, 8)  # 8 个场景类别
        )
    
    def forward(self, x):
        # Enhance
        feat = self.encoder1(x)
        
        # Refine
        attn_map = self.decoder(feat)  # (B, 1, H, W) 注意力图
        
        # Align
        align_feat = self.encoder2(attn_map)
        align_feat = align_feat.flatten(1)
        scene_logits = self.align_head(align_feat)
        
        return {
            'attention_map': attn_map,
            'scene_class': scene_logits
        }


# 测试
if __name__ == "__main__":
    model = EraWNet()
    x = torch.randn(2, 3, 224, 224)
    out = model(x)
    
    print(f"注意力图: {out['attention_map'].shape}")
    print(f"场景分类: {out['scene_class'].shape}")
    print(f"参数: {sum(p.numel() for p in model.parameters()):,}")