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

class TPAFlowModel(nn.Module):
    """
    TPA-Flow rPPG心率估计模型
    
    论文核心架构：
    1. 质量感知生理表示
    2. 张量积注意力(TPA) - 捕获面部区域×心动周期×特征子空间交互
    3. 直接分支 - 稳定波形锚点
    4. 流分支 - 连续传输路径修正
    5. 双路径融合
    """
    
    def __init__(self, num_regions: int = 8, num_cycles: int = 4,
                 feature_dim: int = 128, num_heads: int = 4):
        super().__init__()
        
        # 面部区域特征提取
        self.region_encoder = FaceRegionEncoder(
            num_regions=num_regions,
            out_dim=feature_dim
        )
        
        # 质量评估模块
        self.quality_assessor = QualityAssessor(feature_dim)
        
        # 张量积注意力
        self.tpa = PhysiologicalTensorProductAttention(
            feature_dim=feature_dim,
            num_regions=num_regions,
            num_cycles=num_cycles,
            num_heads=num_heads
        )
        
        # 直接分支（稳定锚点）
        self.direct_branch = DirectDecoder(feature_dim)
        
        # 流分支（校正传输）
        self.flow_branch = RectifiedFlowDecoder(feature_dim)
        
        # 融合
        self.fusion = nn.Sequential(
            nn.Linear(feature_dim * 2, feature_dim),
            nn.GELU(),
            nn.Linear(feature_dim, 1)  # 输出rPPG波形
        )
        
    def forward(self, video: torch.Tensor) -> dict:
        """
        Args:
            video: (B, T, C, H, W) 面部视频
                    B=batch, T=时间步(180=6s@30fps)
        Returns:
            rppg_signal: (B, T) rPPG波形
            quality: (B,) 信号质量评分
        """
        B, T, C, H, W = video.shape
        
        # 1. 面部区域特征
        region_feats = self.region_encoder(video)  # (B, T, num_regions, feature_dim)
        
        # 2. 质量评估
        quality = self.quality_assessor(region_feats)  # (B, T)
        
        # 3. TPA注意力
        attended = self.tpa(region_feats, quality)  # (B, T, feature_dim)
        
        # 4. 直接分支：稳定锚点波形
        direct_anchor = self.direct_branch(attended)  # (B, T)
        
        # 5. 流分支：校正传输
        flow_corrected = self.flow_branch(attended, direct_anchor)  # (B, T)
        
        # 6. 双路径融合
        fused = torch.cat([
            direct_anchor.unsqueeze(-1),
            flow_corrected.unsqueeze(-1)
        ], dim=-1)  # (B, T, 2)
        rppg = self.fusion(fused).squeeze(-1)  # (B, T)
        
        return {
            'rppg_signal': rppg,
            'quality': quality,
            'direct_anchor': direct_anchor,
            'flow_corrected': flow_corrected
        }


class FaceRegionEncoder(nn.Module):
    """面部区域特征编码器"""
    
    def __init__(self, num_regions: int = 8, out_dim: int = 128):
        super().__init__()
        self.num_regions = num_regions
        
        # 轻量级3D CNN
        self.stem = nn.Sequential(
            nn.Conv3d(3, 32, (3, 3, 3), stride=(1, 2, 2), padding=(1, 1, 1)),
            nn.BatchNorm3d(32),
            nn.ReLU(),
        )
        self.layer2 = nn.Sequential(
            nn.Conv3d(32, 64, (3, 3, 3), stride=(1, 2, 2), padding=(1, 1, 1)),
            nn.BatchNorm3d(64),
            nn.ReLU(),
        )
        self.layer3 = nn.Sequential(
            nn.Conv3d(64, 128, (3, 3, 3), stride=(1, 1, 1), padding=(1, 1, 1)),
            nn.BatchNorm3d(128),
            nn.ReLU(),
        )
        
        # 区域注意力（学习面部不同区域权重）
        self.region_attention = nn.MultiheadAttention(
            embed_dim=out_dim, num_heads=4, batch_first=True
        )
        
        # 输出投影
        self.proj = nn.Linear(128, out_dim)
        
    def forward(self, video: torch.Tensor) -> torch.Tensor:
        """
        Args:
            video: (B, T, C, H, W)
        Returns:
            features: (B, T, num_regions, out_dim)
        """
        B, T, C, H, W = video.shape
        
        # 3D CNN提取时空特征
        x = video.permute(0, 2, 1, 3, 4)  # (B, C, T, H, W)
        x = self.stem(x)    # (B, 32, T, H/2, W/2)
        x = self.layer2(x)   # (B, 64, T, H/4, W/4)
        x = self.layer3(x)   # (B, 128, T, H/4, W/4)
        
        # 空间分区域
        _, _, T_, H_, W_ = x.shape
        region_h = H_ // 2
        region_w = W_ // 4
        regions = []
        for i in range(2):
            for j in range(4):
                r = x[:, :, :, i*region_h:(i+1)*region_h,
                              j*region_w:(j+1)*region_w]
                r = F.adaptive_avg_pool3d(r, (T_, 1, 1)).squeeze(-1).squeeze(-1)
                regions.append(self.proj(r.permute(0, 2, 1)))  # (B, T, out_dim)
        
        # 堆叠为区域序列
        region_feats = torch.stack(regions, dim=2)  # (B, T, num_regions, out_dim)
        
        return region_feats


class QualityAssessor(nn.Module):
    """信号质量评估器"""
    def __init__(self, dim: int):
        super().__init__()
        self.mlp = nn.Sequential(
            nn.Linear(dim, dim // 2),
            nn.ReLU(),
            nn.Linear(dim // 2, 1),
            nn.Sigmoid()
        )
    
    def forward(self, x):
        return self.mlp(x.mean(dim=2)).squeeze(-1)


class PhysiologicalTensorProductAttention(nn.Module):
    """
    生理张量积注意力(TPA)
    
    捕获三维交互：
    - 面部区域 × 心动周期 × 特征子空间
    使用低秩张量积避免完整三维注意力的O(n³)复杂度
    """
    
    def __init__(self, feature_dim: int, num_regions: int = 8,
                 num_cycles: int = 4, num_heads: int = 4,
                 rank: int = 16):
        super().__init__()
        self.num_heads = num_heads
        self.rank = rank
        
        # 周期性先验（可学习）
        self.periodic_prior = nn.Parameter(
            torch.randn(num_cycles, feature_dim)
        )
        
        # 质量加权
        self.quality_weight = nn.Linear(1, num_heads)
        
        # 区域注意力
        self.region_attn = nn.MultiheadAttention(
            embed_dim=feature_dim, num_heads=num_heads,
            batch_first=True
        )
        
        # 低秩投影
        self.W_q = nn.Linear(feature_dim, rank)
        self.W_k = nn.Linear(feature_dim, rank)
        self.W_v = nn.Linear(feature_dim, rank)
        
        # 输出投影
        self.out_proj = nn.Linear(feature_dim, feature_dim)
        self.norm = nn.LayerNorm(feature_dim)
        
    def forward(self, region_feats: torch.Tensor,
                quality: torch.Tensor) -> torch.Tensor:
        """
        Args:
            region_feats: (B, T, num_regions, feature_dim)
            quality: (B, T) 质量评分
        Returns:
            attended: (B, T, feature_dim)
        """
        B, T, R, D = region_feats.shape
        
        # 周期先验注入
        prior = self.periodic_prior.unsqueeze(0).unsqueeze(0)  # (1, 1, num_cycles, D)
        prior_expanded = prior.expand(B, T, -1, -1)  # (B, T, num_cycles, D)
        
        # 质量加权
        q_weight = self.quality_weight(quality.unsqueeze(-1))  # (B, T, num_heads)
        
        # 区域间注意力
        flat_regions = region_feats.view(B * T, R, D)
        attn_out, _ = self.region_attn(flat_regions, flat_regions, flat_regions)
        attn_out = attn_out.view(B, T, R, D)
        
        # 低秩张量积
        q = self.W_q(attn_out)  # (B, T, R, rank)
        k = self.W_k(prior_expanded.expand(-1, -1, R, -1))  # (B, T, R, rank)
        v = self.W_v(prior_expanded.expand(-1, -1, R, -1))
        
        # 张量积
        tp_output = torch.einsum('btrd,btrd->btd', q, v)  # (B, T, rank)
        tp_output = self.out_proj(
            F.pad(tp_output, (0, D - self.rank))
        )
        
        # 残差连接
        output = self.norm(region_feats.mean(dim=2) + tp_output)
        
        return output


class DirectDecoder(nn.Module):
    """直接解码分支：输出稳定波形锚点"""
    def __init__(self, dim: int):
        super().__init__()
        self.decoder = nn.Sequential(
            nn.Linear(dim, dim * 2),
            nn.GELU(),
            nn.Linear(dim * 2, dim),
            nn.GELU(),
            nn.Linear(dim, 1)
        )
    
    def forward(self, x):
        return self.decoder(x).squeeze(-1)


class RectifiedFlowDecoder(nn.Module):
    """校正流解码分支：修正锚点波形的相位和形态"""
    def __init__(self, dim: int, num_steps: int = 5):
        super().__init__()
        self.num_steps = num_steps
        
        # 条件编码
        self.cond_encoder = nn.Linear(dim + 1, dim)  # +1 for anchor
        
        # 流网络
        self.flow_net = nn.Sequential(
            nn.Linear(dim * 2, dim * 4),
            nn.GELU(),
            nn.Linear(dim * 4, dim * 2),
            nn.GELU(),
            nn.Linear(dim * 2, dim)
        )
        
    def forward(self, x: torch.Tensor, anchor: torch.Tensor) -> torch.Tensor:
        """
        Args:
            x: (B, T, dim) 条件特征
            anchor: (B, T) 锚点波形
        Returns:
            corrected: (B, T) 校正后波形
        """
        # 条件编码
        cond = self.cond_encoder(
            torch.cat([x, anchor.unsqueeze(-1)], dim=-1)
        )  # (B, T, dim)
        
        # 校正流
        z = anchor.unsqueeze(-1)  # 初始点
        for step in range(self.num_steps):
            flow = self.flow_net(torch.cat([z, cond], dim=-1))
            z = z + flow / self.num_steps  # 线性传输
            
        return z.squeeze(-1)


# ==================== 测试代码 ====================
if __name__ == "__main__":
    model = TPAFlowModel(
        num_regions=8, num_cycles=4,
        feature_dim=128, num_heads=4
    )
    
    # 模拟输入
    batch_size = 2
    T = 180  # 6秒@30fps
    video = torch.randn(batch_size, T, 3, 128, 128)
    
    with torch.no_grad():
        output = model(video)
    
    print("TPA-Flow 模型测试")
    print("=" * 50)
    print(f"rPPG信号形状: {output['rppg_signal'].shape}")
    print(f"质量评分形状: {output['quality'].shape}")
    print(f"直接锚点形状: {output['direct_anchor'].shape}")
    print(f"校正流形状:   {output['flow_corrected'].shape}")
    
    # 心率估计
    import numpy as np
    rppg = output['rppg_signal'][0].numpy()
    fft = np.abs(np.fft.rfft(rppg))
    freqs = np.fft.rfftfreq(len(rppg), 1/30)
    hr_mask = (freqs >= 0.67) & (freqs <= 3.0)
    hr_freq = freqs[hr_mask][np.argmax(fft[hr_mask])]
    print(f"\n估计心率: {hr_freq * 60:.1f} BPM")
    
    param_count = sum(p.numel() for p in model.parameters())
    print(f"参数量: {param_count / 1e6:.2f}M")