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

class KeypointGuidedSampling(nn.Module):
    """
    关键点引导采样模块（KGS）
    
    基于人体关键点（肩、胸、腰）定位安全带可能区域
    安全带通常从肩部斜跨胸部到腰部
    """
    
    # 人体关键点中与安全带相关的索引（COCO 17点）
    SHOULDER_LEFT = 5    # 左肩
    SHOULDER_RIGHT = 6   # 右肩
    ELBOW_LEFT = 7
    ELBOW_RIGHT = 8
    HIP_LEFT = 11        # 左腰
    HIP_RIGHT = 12       # 右腰
    
    def __init__(self, img_size: int = 640, crop_size: int = 256):
        super().__init__()
        self.img_size = img_size
        self.crop_size = crop_size
        
    def forward(self, image: torch.Tensor, 
                keypoints: torch.Tensor) -> dict:
        """
        Args:
            image: (B, 3, H, W) 原始图像
            keypoints: (B, 17, 2) COCO格式关键点坐标
            
        Returns:
            crops: dict containing local and global crops
        """
        B = image.shape[0]
        crops = {
            'local_shoulder': [],  # 肩部区域
            'local_chest': [],     # 胸部区域
            'local_hip': [],       # 腰部区域
            'global': image         # 整体图像
        }
        
        for b in range(B):
            kp = keypoints[b]  # (17, 2)
            
            # 肩部区域采样
            shoulder_center = (kp[self.SHOULDER_LEFT] + 
                              kp[self.SHOULDER_RIGHT]) / 2
            shoulder_crop = self._crop_region(
                image[b], shoulder_center, self.crop_size
            )
            crops['local_shoulder'].append(shoulder_crop)
            
            # 胸部区域采样（肩到腰的中点）
            chest_center = (kp[self.SHOULDER_LEFT] + kp[self.SHOULDER_RIGHT] +
                           kp[self.HIP_LEFT] + kp[self.HIP_RIGHT]) / 4
            chest_crop = self._crop_region(
                image[b], chest_center, self.crop_size
            )
            crops['local_chest'].append(chest_crop)
            
            # 腰部区域采样
            hip_center = (kp[self.HIP_LEFT] + kp[self.HIP_RIGHT]) / 2
            hip_crop = self._crop_region(
                image[b], hip_center, self.crop_size
            )
            crops['local_hip'].append(hip_crop)
        
        # 转为tensor
        for key in ['local_shoulder', 'local_chest', 'local_hip']:
            crops[key] = torch.stack(crops[key])
        
        return crops
    
    def _crop_region(self, img: torch.Tensor, center: torch.Tensor, 
                    size: int) -> torch.Tensor:
        """以center为中心裁剪size×size区域"""
        cx, cy = int(center[0].item()), int(center[1].item())
        half = size // 2
        
        # 边界检查
        _, H, W = img.shape
        x1 = max(0, cx - half)
        y1 = max(0, cy - half)
        x2 = min(W, x1 + size)
        y2 = min(H, y1 + size)
        
        crop = img[:, y1:y2, x1:x2]
        # Resize到统一尺寸
        if crop.shape[1] != size or crop.shape[2] != size:
            crop = F.interpolate(
                crop.unsqueeze(0), size=(size, size), 
                mode='bilinear', align_corners=False
            ).squeeze(0)
        
        return crop


class LocalGlobalAttention(nn.Module):
    """
    局部-全局注意力融合模块（LGA）
    
    将局部裁剪区域的安全带细粒度特征
    与全局图像的上下文特征融合
    """
    def __init__(self, local_dim: int = 256, global_dim: int = 512,
                 fusion_dim: int = 256, nhead: int = 4):
        super().__init__()
        
        # 局部特征编码器（3个区域共享权重）
        self.local_encoder = nn.Sequential(
            nn.Conv2d(3, 32, 3, stride=2, padding=1),
            nn.BatchNorm2d(32),
            nn.ReLU(),
            nn.Conv2d(32, 64, 3, stride=2, padding=1),
            nn.BatchNorm2d(64),
            nn.ReLU(),
            nn.Conv2d(64, 128, 3, stride=2, padding=1),
            nn.BatchNorm2d(128),
            nn.ReLU(),
            nn.AdaptiveAvgPool2d(1)
        )
        self.local_proj = nn.Linear(128, local_dim)
        
        # 全局特征编码器
        self.global_encoder = nn.Sequential(
            nn.Conv2d(3, 32, 3, stride=2, padding=1),
            nn.BatchNorm2d(32),
            nn.ReLU(),
            nn.Conv2d(32, 64, 3, stride=2, padding=1),
            nn.BatchNorm2d(64),
            nn.ReLU(),
            nn.Conv2d(64, 128, 3, stride=2, padding=1),
            nn.BatchNorm2d(128),
            nn.ReLU(),
            nn.Conv2d(128, 256, 3, stride=2, padding=1),
            nn.BatchNorm2d(256),
            nn.ReLU(),
            nn.AdaptiveAvgPool2d(1)
        )
        self.global_proj = nn.Linear(256, global_dim)
        
        # 跨模态注意力
        self.cross_attn = nn.MultiheadAttention(
            embed_dim=fusion_dim,
            num_heads=nhead,
            batch_first=True
        )
        
        # 将local和global投影到统一维度
        self.local_to_fusion = nn.Linear(local_dim, fusion_dim)
        self.global_to_fusion = nn.Linear(global_dim, fusion_dim)
        
        # 分类头
        self.classifier = nn.Sequential(
            nn.Linear(fusion_dim * 4, 256),  # 3 local + 1 global
            nn.ReLU(),
            nn.Dropout(0.3),
            nn.Linear(256, 3)  # 正确/错误/未佩戴
        )
        
    def forward(self, crops: dict) -> torch.Tensor:
        """
        Args:
            crops: dict with local_shoulder, local_chest, 
                   local_hip, global
        Returns:
            logits: (B, 3)
        """
        features = []
        
        # 局部特征
        for key in ['local_shoulder', 'local_chest', 'local_hip']:
            x = crops[key]  # (B, 3, crop_size, crop_size)
            feat = self.local_encoder(x)  # (B, 128, 1, 1)
            feat = feat.flatten(1)  # (B, 128)
            feat = self.local_proj(feat)  # (B, local_dim)
            feat = self.local_to_fusion(feat)  # (B, fusion_dim)
            features.append(feat)
        
        # 全局特征
        global_x = crops['global']  # (B, 3, H, W)
        global_feat = self.global_encoder(global_x)  # (B, 256, 1, 1)
        global_feat = global_feat.flatten(1)  # (B, 256)
        global_feat = self.global_proj(global_feat)  # (B, global_dim)
        global_feat = self.global_to_fusion(global_feat)  # (B, fusion_dim)
        features.append(global_feat)
        
        # 堆叠为序列: (B, 4, fusion_dim)
        feat_seq = torch.stack(features, dim=1)
        
        # 自注意力
        attn_out, _ = self.cross_attn(feat_seq, feat_seq, feat_seq)
        
        # 池化 + 分类
        pooled = attn_out.mean(dim=1)  # (B, fusion_dim)
        # 拼接所有特征
        all_feats = torch.cat([pooled] + 
                             [f for f in features], dim=-1)  # (B, fusion_dim*4)
        logits = self.classifier(all_feats)
        
        return logits


class KPGBeltNet(nn.Module):
    """
    KPGBeltNet 完整模型
    
    论文：The Visual Computer, 2026
    DOI: 10.1007/s00371-026-04572-1
    
    管道：
    1. 人体关键点检测（SPPE）
    2. 关键点引导采样（KGS）
    3. 局部-全局注意力融合（LGA）
    4. 三分类：正确佩戴/错误佩戴/未佩戴
    """
    def __init__(self, img_size: int = 640, crop_size: int = 128):
        super().__init__()
        
        # 关键点检测器（使用预训练SPPE）
        self.keypoint_detector = self._load_sppe()
        
        # 关键点引导采样
        self.kgs = KeypointGuidedSampling(
            img_size=img_size, 
            crop_size=crop_size
        )
        
        # 局部-全局注意力
        self.lga = LocalGlobalAttention(
            local_dim=256,
            global_dim=512,
            fusion_dim=256,
            nhead=4
        )
        
    def forward(self, image: torch.Tensor) -> torch.Tensor:
        """
        Args:
            image: (B, 3, H, W)
        Returns:
            logits: (B, 3) [正确, 错误, 未佩戴]
        """
        # 1. 关键点检测
        with torch.no_grad():
            keypoints = self.keypoint_detector(image)
        
        # 2. 引导采样
        crops = self.kgs(image, keypoints)
        
        # 3. 局部-全局注意力 + 分类
        logits = self.lga(crops)
        
        return logits
    
    def _load_sppe(self):
        """加载预训练单人姿态估计模型"""
        # 简化实现：实际使用FastPose或HRNet
        model = nn.Sequential(
            nn.Conv2d(3, 64, 3, stride=2, padding=1),
            nn.BatchNorm2d(64),
            nn.ReLU(),
            nn.Conv2d(64, 128, 3, stride=2, padding=1),
            nn.BatchNorm2d(128),
            nn.ReLU(),
            nn.Conv2d(128, 256, 3, stride=2, padding=1),
            nn.BatchNorm2d(256),
            nn.ReLU(),
            nn.AdaptiveAvgPool2d((4, 4)),
            nn.Flatten(),
            nn.Linear(256 * 16, 17 * 2)  # 17个关键点, xy
        )
        return model


# ==================== 测试代码 ====================
if __name__ == "__main__":
    model = KPGBeltNet(img_size=640, crop_size=128)
    
    # 模拟输入
    batch_size = 4
    image = torch.randn(batch_size, 3, 640, 640)
    
    with torch.no_grad():
        logits = model(image)
        probs = torch.softmax(logits, dim=1)
    
    classes = ['正确佩戴', '错误佩戴', '未佩戴']
    print(f"输出形状: {logits.shape}")
    for i in range(batch_size):
        pred = classes[probs[i].argmax()]
        print(f"  样本{i}: {pred} (置信度={probs[i].max():.2%})")
    print(f"参数量: {sum(p.numel() for p in model.parameters()) / 1e6:.1f}M")