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

class SpatialEncoder(nn.Module):
    """空间编码器：从单帧提取 AU 空间分布"""
    
    def __init__(self, n_aus=12):
        super().__init__()
        
        # 面部特征提取（轻量级 CNN）
        self.backbone = 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((4, 4)),
        )
        
        # AU 回归头
        self.au_head = nn.Sequential(
            nn.Linear(128 * 16, 256),
            nn.ReLU(),
            nn.Dropout(0.3),
            nn.Linear(256, n_aus),
            nn.Sigmoid()  # AU 强度 0-1
        )
        
    def forward(self, x):
        """
        Args:
            x: (B, C, H, W) 单帧面部图像
            
        Returns:
            aus: (B, n_aus) AU 强度
            feat: (B, 2048) 特征向量
        """
        feat = self.backbone(x)
        feat_flat = feat.flatten(1)
        aus = self.au_head(feat_flat)
        return aus, feat_flat


class TemporalEncoder(nn.Module):
    """时间编码器：建模 AU 序列动态"""
    
    def __init__(self, n_aus=12, hidden_dim=128, n_layers=2):
        super().__init__()
        
        # GRU 建模时序
        self.gru = nn.GRU(
            input_size=n_aus,
            hidden_dim=hidden_dim,
            num_layers=n_layers,
            batch_first=True,
            bidirectional=True
        )
        
        # AU 时序预测头
        self.temporal_head = nn.Sequential(
            nn.Linear(hidden_dim * 2, n_aus),
            nn.Sigmoid()
        )
        
    def forward(self, au_sequence):
        """
        Args:
            au_sequence: (B, T, n_aus) AU 时间序列
            
        Returns:
            au_pred: (B, T, n_aus) 预测的下一帧 AU
            hidden: (B, hidden_dim*2) 时序特征
        """
        out, hidden = self.gru(au_sequence)
        au_pred = self.temporal_head(out)
        return au_pred, hidden


class TwinCycleAutoencoder(nn.Module):
    """
    双循环自编码器
    
    空间循环：图像 → AU → 重建图像
    时间循环：AU序列 → 预测下一帧 → 重建序列
    
    两个循环互相约束：
    - 空间编码的 AU 必须能重建面部图像
    - 时间编码的 AU 必须能预测未来 AU
    """
    
    def __init__(self, n_aus=12):
        super().__init__()
        self.spatial = SpatialEncoder(n_aus=n_aus)
        self.temporal = TemporalEncoder(n_aus=n_aus)
        
        # 图像解码器（空间循环重建）
        self.image_decoder = nn.Sequential(
            nn.Linear(n_aus, 256),
            nn.ReLU(),
            nn.Linear(256, 128 * 4 * 4),
            nn.ReLU(),
        )
        self.deconv = nn.Sequential(
            nn.ConvTranspose2d(128, 64, 3, stride=2, padding=1, output_padding=1),
            nn.BatchNorm2d(64),
            nn.ReLU(),
            nn.ConvTranspose2d(64, 32, 3, stride=2, padding=1, output_padding=1),
            nn.BatchNorm2d(32),
            nn.ReLU(),
            nn.ConvTranspose2d(32, 3, 3, stride=2, padding=1, output_padding=1),
            nn.Sigmoid()
        )
    
    def forward(self, image_sequence):
        """
        Args:
            image_sequence: (B, T, C, H, W) 面部视频序列
            
        Returns:
            dict: 空间 AU、时序 AU、重建图像、预测 AU
        """
        B, T, C, H, W = image_sequence.shape
        
        # 1. 空间编码：逐帧提取 AU
        frames = image_sequence.reshape(B*T, C, H, W)
        au_spatial, _ = self.spatial(frames)
        au_spatial = au_spatial.reshape(B, T, -1)  # (B, T, n_aus)
        
        # 2. 时间编码：建模 AU 动态
        au_temporal, hidden = self.temporal(au_spatial)
        
        # 3. 空间循环重建：AU → 图像
        au_flat = au_spatial.reshape(B*T, -1)
        recon_feat = self.image_decoder(au_flat)
        recon_feat = recon_feat.reshape(B*T, 128, 4, 4)
        recon_images = self.deconv(recon_feat)
        recon_images = recon_images.reshape(B, T, C, H, W)
        
        # 4. 时间循环预测：当前 AU → 预测下一帧 AU
        # au_temporal 已经是预测序列
        
        return {
            'au_spatial': au_spatial,    # (B, T, n_aus) 检测的 AU
            'au_temporal': au_temporal,   # (B, T, n_aus) 预测的 AU
            'recon_images': recon_images, # (B, T, C, H, W) 重建图像
            'hidden': hidden,             # 时序特征
        }
    
    def compute_loss(self, outputs, targets):
        """
        计算双循环一致性损失
        
        Args:
            outputs: forward() 返回值
            targets: {
                'images': (B, T, C, H, W) 原始图像,
                'au_labels': (B, T, n_aus) AU 标签（可选）
            }
        """
        # 1. 空间循环损失：图像重建
        recon_loss = F.mse_loss(
            outputs['recon_images'], 
            targets['images']
        )
        
        # 2. 时间循环损失：AU 预测一致性
        temporal_loss = F.l1_loss(
            outputs['au_temporal'][:, :-1],  # 预测 t
            outputs['au_spatial'][:, 1:].detach()  # 实际 t+1（stop gradient）
        )
        
        # 3. AU 检测损失（如果有标签）
        if 'au_labels' in targets:
            detection_loss = F.binary_cross_entropy(
                outputs['au_spatial'],
                targets['au_labels']
            )
        else:
            detection_loss = torch.tensor(0.0)
        
        # 总损失
        total_loss = (
            1.0 * recon_loss + 
            0.5 * temporal_loss + 
            1.0 * detection_loss
        )
        
        return {
            'total': total_loss,
            'recon': recon_loss,
            'temporal': temporal_loss,
            'detection': detection_loss
        }


# 测试
if __name__ == "__main__":
    model = TwinCycleAutoencoder(n_aus=12)
    
    # 模拟输入：(batch=2, frames=16, channels=3, H=64, W=64)
    x = torch.randn(2, 16, 3, 64, 64)
    x = torch.sigmoid(x)  # 归一化到 0-1
    
    outputs = model(x)
    
    print("Twin Cycle Autoencoder 输出:")
    print(f"  AU 空间: {outputs['au_spatial'].shape}")
    print(f"  AU 时序: {outputs['au_temporal'].shape}")
    print(f"  重建图像: {outputs['recon_images'].shape}")
    print(f"  隐藏特征: {outputs['hidden'].shape}")
    
    # 损失计算
    targets = {'images': x}
    losses = model.compute_loss(outputs, targets)
    
    print(f"\n损失:")
    for k, v in losses.items():
        print(f"  {k}: {v.item():.4f}")
    
    print(f"\n参数总数: {sum(p.numel() for p in model.parameters()):,}")

class DriverEmotionClassifier:
    """
    基于 AU 检测结果的驾驶员情绪分类器
    
    使用 FACS 规则+时序模式
    """
    
    EMOTION_RULES = {
        'anger': {
            'required_aus': [('AU4', 0.5), ('AU23', 0.3)],
            'optional_aus': ['AU7', 'AU17'],
            'min_duration_frames': 5,  # ~1s @ 5fps
        },
        'fear_surprise': {
            'required_aus': [('AU1', 0.4), ('AU2', 0.4), ('AU5', 0.3)],
            'optional_aus': ['AU20'],
            'min_duration_frames': 3,
        },
        'happiness': {
            'required_aus': [('AU6', 0.4), ('AU12', 0.4)],
            'optional_aus': [],
            'min_duration_frames': 5,
        },
        'sadness': {
            'required_aus': [('AU15', 0.3), ('AU4', 0.2)],
            'optional_aus': ['AU17'],
            'min_duration_frames': 10,
        },
        'disgust': {
            'required_aus': [('AU9', 0.3), ('AU7', 0.3)],
            'optional_aus': ['AU15'],
            'min_duration_frames': 5,
        },
        'concentration': {
            'required_aus': [('AU4', 0.3)],
            'optional_aus': ['AU7'],
            'min_duration_frames': 15,  # 需要持续 3s
        },
    }
    
    # AU 索引映射（与模型输出对齐）
    AU_INDICES = {
        'AU1': 0, 'AU2': 1, 'AU4': 2, 'AU5': 3, 'AU6': 4,
        'AU7': 5, 'AU9': 6, 'AU12': 7, 'AU15': 8, 'AU17': 9,
        'AU20': 10, 'AU23': 11
    }
    
    def classify(self, au_sequence):
        """
        分类情绪状态
        
        Args:
            au_sequence: (T, 12) AU 强度时间序列
            
        Returns:
            emotion: 情绪标签
            confidence: 置信度
        """
        T = len(au_sequence)
        scores = {}
        
        for emotion, rule in self.EMOTION_RULES.items():
            # 检查必要 AU 是否满足阈值
            required_met = 0
            for au_name, threshold in rule['required_aus']:
                idx = self.AU_INDICES[au_name]
                # 检查最近 N 帧是否持续超过阈值
                n_frames = min(T, rule['min_duration_frames'])
                recent = au_sequence[-n_frames:, idx]
                if np.mean(recent > threshold) > 0.6:  # 60% 帧满足
                    required_met += 1
            
            # 检查可选 AU
            optional_met = 0
            for au_name in rule['optional_aus']:
                idx = self.AU_INDICES[au_name]
                recent = au_sequence[-min(T, rule['min_duration_frames']):, idx]
                if np.mean(recent > 0.3) > 0.5:
                    optional_met += 1
            
            # 综合评分
            total_required = len(rule['required_aus'])
            total_optional = len(rule['optional_aus'])
            
            score = (required_met / total_required) * 0.7
            if total_optional > 0:
                score += (optional_met / total_optional) * 0.3
            
            scores[emotion] = score
        
        # 选最高分
        emotion = max(scores, key=scores.get)
        confidence = scores[emotion]
        
        if confidence < 0.3:
            return 'neutral', confidence
        
        return emotion, confidence


# 测试
if __name__ == "__main__":
    import numpy as np
    
    clf = DriverEmotionClassifier()
    
    # 模拟愤怒序列（AU4 + AU23 持续高强度）
    np.random.seed(42)
    anger_seq = np.zeros((20, 12))
    anger_seq[:, 2] = 0.6  # AU4
    anger_seq[:, 11] = 0.5  # AU23
    anger_seq[:, 5] = 0.4   # AU7
    
    r = clf.classify(anger_seq)
    print(f"模拟愤怒: {r[0]} (置信度: {r[1]:.2f})")
    
    # 模拟愉悦序列
    happy_seq = np.zeros((20, 12))
    happy_seq[:, 4] = 0.6   # AU6
    happy_seq[:, 7] = 0.7   # AU12
    
    r = clf.classify(happy_seq)
    print(f"模拟愉悦: {r[0]} (置信度: {r[1]:.2f})")
    
    # 模拟中性
    neutral_seq = np.random.rand(20, 12) * 0.1
    
    r = clf.classify(neutral_seq)
    print(f"模拟中性: {r[0]} (置信度: {r[1]:.2f})")