import torch
import torch.nn as nn

class BiFuseNet(nn.Module):
    """
    BiFuseNet: 3D 时空多模态 BAC 估计网络
    
    架构:
    1. 双路 3D ResNet 主干 (RGB + IR)
    2. Bi-directional Hierarchical Fusion (BiHF)
       - Early fusion: 浅层特征交叉注意力
       - Middle fusion: 中层特征融合
       - Late fusion: 深层语义融合
    3. 3D 时空卷积捕获时序动态
    4. 分级头: sober / moderate / severe
    
    输入: (B, C_rgb, T, H, W) + (B, C_ir, T, H, W)
    输出: (B, 3) 三级分类 + (B, 1) BAC 回归
    """
    
    def __init__(self, num_classes=3, dropout=0.3):
        super().__init__()
        
        # RGB 分支 (3D ResNet)
        self.rgb_branch = ResNet3D(
            in_channels=3,  # RGB
            layers=[2, 2, 2, 2],
            base_ch=32
        )
        
        # IR 分支 (3D ResNet)
        self.ir_branch = ResNet3D(
            in_channels=1,  # 红外灰度
            layers=[2, 2, 2, 2],
            base_ch=32
        )
        
        # BiHF: 双向分层融合
        self.bihf = BiHierarchicalFusion(
            channels_list=[32, 64, 128, 256],  # 四个语义层级
            num_heads=4
        )
        
        # 时空注意力
        self.temporal_attn = TemporalAttention(256, num_heads=4)
        
        # 分类头
        self.classifier = nn.Sequential(
            nn.AdaptiveAvgPool3d(1),
            nn.Flatten(),
            nn.Linear(256, 128),
            nn.ReLU(),
            nn.Dropout(dropout),
            nn.Linear(128, num_classes)
        )
        
        # BAC 回归头
        self.regressor = nn.Sequential(
            nn.AdaptiveAvgPool3d(1),
            nn.Flatten(),
            nn.Linear(256, 64),
            nn.ReLU(),
            nn.Dropout(dropout),
            nn.Linear(64, 1),
            nn.Sigmoid()  # 0-1 BAC 归一化
        )
        
    def forward(self, rgb_video, ir_video):
        """
        Args:
            rgb_video: (B, 3, T, H, W) RGB 面部视频
            ir_video: (B, 1, T, H, W) 红外面部视频
            
        Returns:
            logits: (B, 3) 三级分类
            bac: (B, 1) BAC 估计值
        """
        # 双路特征提取
        rgb_feats = self.rgb_branch(rgb_video)  # list of features
        ir_feats = self.ir_branch(ir_video)
        
        # 双向分层融合
        fused = self.bihf(rgb_feats, ir_feats)  # (B, 256, T', H', W')
        
        # 时空注意力
        fused = self.temporal_attn(fused)
        
        # 输出
        logits = self.classifier(fused)   # (B, 3)
        bac = self.regressor(fused)       # (B, 1)
        
        return logits, bac


class ResNet3D(nn.Module):
    """简化版 3D ResNet"""
    
    def __init__(self, in_channels=3, layers=[2,2,2,2], base_ch=32):
        super().__init__()
        
        self.stem = nn.Sequential(
            nn.Conv3d(in_channels, base_ch, kernel_size=(3,7,7),
                     stride=(1,2,2), padding=(1,3,3)),
            nn.BatchNorm3d(base_ch),
            nn.ReLU(),
            nn.MaxPool3d(kernel_size=(1,3,3), stride=(1,2,2))
        )
        
        self.blocks = nn.ModuleList()
        ch = base_ch
        for i, num_blocks in enumerate(layers):
            out_ch = ch * 2 if i > 0 else ch
            block = self._make_layer(ch, out_ch, num_blocks, 
                                     stride=(2,2,2) if i > 0 else (1,2,2))
            self.blocks.append(block)
            ch = out_ch
            
    def _make_layer(self, in_ch, out_ch, num_blocks, stride):
        layers = [ResBlock3D(in_ch, out_ch, stride)]
        for _ in range(1, num_blocks):
            layers.append(ResBlock3D(out_ch, out_ch))
        return nn.Sequential(*layers)
    
    def forward(self, x):
        x = self.stem(x)
        feats = []
        for block in self.blocks:
            x = block(x)
            feats.append(x)
        return feats


class ResBlock3D(nn.Module):
    def __init__(self, in_ch, out_ch, stride=1):
        super().__init__()
        self.conv1 = nn.Conv3d(in_ch, out_ch, 3, stride, 1, bias=False)
        self.bn1 = nn.BatchNorm3d(out_ch)
        self.conv2 = nn.Conv3d(out_ch, out_ch, 3, 1, 1, bias=False)
        self.bn2 = nn.BatchNorm3d(out_ch)
        self.relu = nn.ReLU()
        
        self.downsample = None
        if stride != 1 or in_ch != out_ch:
            self.downsample = nn.Sequential(
                nn.Conv3d(in_ch, out_ch, 1, stride, bias=False),
                nn.BatchNorm3d(out_ch)
            )
    
    def forward(self, x):
        identity = x
        out = self.relu(self.bn1(self.conv1(x)))
        out = self.bn2(self.conv2(out))
        if self.downsample:
            identity = self.downsample(x)
        return self.relu(out + identity)


# 测试
if __name__ == "__main__":
    model = BiFuseNet(num_classes=3)
    rgb = torch.randn(2, 3, 16, 224, 224)  # 2 batch, 3 ch, 16 frames, 224x224
    ir = torch.randn(2, 1, 16, 224, 224)
    
    logits, bac = model(rgb, ir)
    print(f"Input RGB: {rgb.shape}")
    print(f"Input IR: {ir.shape}")
    print(f"Classification: {logits.shape}")  # (2, 3)
    print(f"BAC estimate: {bac.shape}")        # (2, 1)
    print(f"Params: {sum(p.numel() for p in model.parameters()):,}")

class BiHierarchicalFusion(nn.Module):
    """
    BiHF: 双向分层融合
    
    论文核心模块: 在三个语义层级做跨模态交叉注意力
    
    Early fusion:  浅层空间特征 (高分辨率, 细粒度)
    Middle fusion: 中层语义特征 (中等分辨率)
    Late fusion:   深层语义特征 (低分辨率, 高级概念)
    
    每层使用双向交叉注意力:
    - RGB → IR: RGB 特征查询 IR 特征
    - IR → RGB: IR 特征查询 RGB 特征
    """
    
    def __init__(self, channels_list=[32,64,128,256], num_heads=4):
        super().__init__()
        self.num_levels = len(channels_list)
        
        # 每层双向交叉注意力
        self.cross_attn_rgb2ir = nn.ModuleList([
            CrossModalAttention(ch, num_heads) for ch in channels_list
        ])
        self.cross_attn_ir2rgb = nn.ModuleList([
            CrossModalAttention(ch, num_heads) for ch in channels_list
        ])
        
        # 每层融合后投影
        self.fuse_proj = nn.ModuleList([
            nn.Conv3d(ch * 3, ch, 1) for ch in channels_list  # RGB + IR + 融合
        ])
        
    def forward(self, rgb_feats, ir_feats):
        """
        Args:
            rgb_feats: list of (B, C, T, H, W)
            ir_feats: list of (B, C, T, H, W)
            
        Returns:
            fused: (B, C_last, T', H', W')
        """
        fused_list = []
        
        for i in range(self.num_levels):
            rf = rgb_feats[i]
            irf = ir_feats[i]
            
            # 双向交叉注意力
            rgb_enhanced = self.cross_attn_rgb2ir[i](rf, irf)  # RGB 用 IR 增强
            ir_enhanced = self.cross_attn_ir2rgb[i](irf, rf)    # IR 用 RGB 增强
            
            # 拼接融合
            combined = torch.cat([rgb_enhanced, ir_enhanced, rf + irf], dim=1)
            fused = self.fuse_proj[i](combined)
            
            fused_list.append(fused)
        
        # 返回最深层的特征
        return fused_list[-1]


class CrossModalAttention(nn.Module):
    """跨模态交叉注意力"""
    
    def __init__(self, channels, num_heads=4):
        super().__init__()
        self.norm_q = nn.LayerNorm(channels)
        self.norm_kv = nn.LayerNorm(channels)
        self.attn = nn.MultiheadAttention(channels, num_heads, batch_first=True)
        self.gamma = nn.Parameter(torch.zeros(1))
        
    def forward(self, query_feat, key_value_feat):
        """
        Args:
            query_feat: (B, C, T, H, W)
            key_value_feat: (B, C, T, H, W)
        """
        B, C, T, H, W = query_feat.shape
        
        # 展平空间维度
        q = query_feat.flatten(2).transpose(1, 2)  # (B, T*H*W, C)
        kv = key_value_feat.flatten(2).transpose(1, 2)
        
        # 交叉注意力
        q_norm = self.norm_q(q)
        kv_norm = self.norm_kv(kv)
        
        attn_out, _ = self.attn(q_norm, kv_norm, kv_norm)
        
        # 残差连接
        out = q + self.gamma * attn_out
        out = out.transpose(1, 2).reshape(B, C, T, H, W)
        
        return out


class TemporalAttention(nn.Module):
    """时空注意力模块"""
    
    def __init__(self, channels, num_heads=4):
        super().__init__()
        self.temporal_pos = nn.Parameter(torch.randn(1, 1, channels))
        self.attn = nn.MultiheadAttention(channels, num_heads, batch_first=True)
        self.norm = nn.LayerNorm(channels)
        
    def forward(self, x):
        """
        Args:
            x: (B, C, T, H, W)
        """
        B, C, T, H, W = x.shape
        
        # 全局平均池化空间维度
        x_spatial = x.mean(dim=(-2, -1))  # (B, C, T)
        x_spatial = x_spatial.transpose(1, 2)  # (B, T, C)
        
        # 加位置编码
        x_spatial = x_spatial + self.temporal_pos
        
        # 时序自注意力
        normed = self.norm(x_spatial)
        attn_out, _ = self.attn(normed, normed, normed)
        x_spatial = x_spatial + attn_out
        
        # 广播回原始形状
        x_spatial = x_spatial.transpose(1, 2).unsqueeze(-1).unsqueeze(-1)
        x = x + x_spatial  # 残差
        
        return x


if __name__ == "__main__":
    # 测试 BiHF
    bihf = BiHierarchicalFusion([32, 64, 128, 256])
    
    rgb_feats = [torch.randn(2, 32, 16, 56, 56),
                 torch.randn(2, 64, 8, 28, 28),
                 torch.randn(2, 128, 4, 14, 14),
                 torch.randn(2, 256, 2, 7, 7)]
    ir_feats = [f.clone() for f in rgb_feats]
    
    fused = bihf(rgb_feats, ir_feats)
    print(f"Fused output shape: {fused.shape}")

def bifuse_loss(logits, bac_pred, labels, bac_gt,
                cls_weight=1.0, reg_weight=0.5, ordinal_weight=0.3):
    """
    BiFuseNet 联合损失函数
    
    三级分类 + BAC 回归 + 序数约束
    
    Args:
        logits: (B, 3) sober/moderate/severe
        bac_pred: (B, 1) BAC 估计 (0-1)
        labels: (B,) 等级标签
        bac_gt: (B, 1) BAC 真值
    """
    import torch.nn.functional as F
    
    # 分类损失
    cls_loss = F.cross_entropy(logits, labels)
    
    # BAC 回归损失 (MSE)
    reg_loss = F.mse_loss(bac_pred.squeeze(), bac_gt.squeeze())
    
    # 序数约束: sober < moderate < severe
    probs = F.softmax(logits, dim=-1)
    # sober 概率应随 BAC 下降, severe 概率应随 BAC 上升
    ordinal_penalty = torch.mean(
        probs[:, 0] * bac_gt.squeeze() +  # sober 不应有高 BAC
        probs[:, 2] * (1 - bac_gt.squeeze())  # severe 不应有低 BAC
    )
    
    total_loss = (
        cls_weight * cls_loss + 
        reg_weight * reg_loss +
        ordinal_weight * ordinal_penalty
    )
    
    return total_loss, {
        'cls': cls_loss.item(),
        'reg': reg_loss.item(),
        'ordinal': ordinal_penalty.item()
    }


if __name__ == "__main__":
    logits = torch.randn(4, 3)
    bac_pred = torch.rand(4, 1)
    labels = torch.tensor([0, 1, 1, 2])
    bac_gt = torch.rand(4, 1)
    
    loss, components = bifuse_loss(logits, bac_pred, labels, bac_gt)
    print(f"Total loss: {loss.item():.4f}")
    for k, v in components.items():
        print(f"  {k}: {v:.4f}")

def bifusenet_deployment_config():
    """
    BiFuseNet 部署配置建议
    
    模型大小和推理时间估算
    """
    config = {
        "model_params": "~15-20M (3D ResNet-18 级别)",
        "fp32_size": "~60-80 MB",
        "int8_size": "~15-20 MB",
        "inference_time": {
            "QCS8295_NPU": "8-15 ms / 16帧",
            "QCS8255_NPU": "12-25 ms / 16帧",
            "CPU_only": "50-100 ms / 16帧"
        },
        "fps_target": "30 fps (每帧 33ms 预算)",
        "power": "< 2W (NPU 推理)",
        "input_format": {
            "rgb": "16帧 × 224×224 × 3ch",
            "ir": "16帧 × 224×224 × 1ch"
        },
        "output": {
            "classification": "sober/moderate/severe",
            "bac_estimate": "0.00-0.30% 连续值"
        }
    }
    
    for k, v in config.items():
        print(f"{k}: {v}")
    
    return config


if __name__ == "__main__":
    bifusenet_deployment_config()