# 示意代码，不可直接运行（论文未开源，全文未核对到具体网络结构，U-Net细节与超参为示意，以原文为准）
import torch
import torch.nn as nn
import torch.nn.functional as F

class PPGInpaintingGenerator(nn.Module):
    """
    U-Net生成器：修复rPPG信号退化段（示意）

    摘要确认论文为自定义 Wasserstein GAN + 梯度惩罚；
    具体生成器结构原文未开放，此处的U-Net实现为示意
    """
    
    def __init__(self, in_channels: int = 2,  # 信号+掩码
                 out_channels: int = 1,
                 base_channels: int = 64):
        super().__init__()
        
        # 编码器（下采样）
        self.enc1 = self._double_conv(in_channels, base_channels)
        self.enc2 = self._double_conv(base_channels, base_channels * 2)
        self.enc3 = self._double_conv(base_channels * 2, base_channels * 4)
        self.enc4 = self._double_conv(base_channels * 4, base_channels * 8)
        
        # 瓶颈层
        self.bottleneck = self._double_conv(
            base_channels * 8, base_channels * 16
        )
        
        # 解码器（上采样+跳连接）
        self.upconv4 = nn.ConvTranspose1d(
            base_channels * 16, base_channels * 8, 2, stride=2
        )
        self.dec4 = self._double_conv(
            base_channels * 16, base_channels * 8
        )
        
        self.upconv3 = nn.ConvTranspose1d(
            base_channels * 8, base_channels * 4, 2, stride=2
        )
        self.dec3 = self._double_conv(
            base_channels * 8, base_channels * 4
        )
        
        self.upconv2 = nn.ConvTranspose1d(
            base_channels * 4, base_channels * 2, 2, stride=2
        )
        self.dec2 = self._double_conv(
            base_channels * 4, base_channels * 2
        )
        
        self.upconv1 = nn.ConvTranspose1d(
            base_channels * 2, base_channels, 2, stride=2
        )
        self.dec1 = self._double_conv(
            base_channels * 2, base_channels
        )
        
        # 输出层
        self.out_conv = nn.Conv1d(base_channels, out_channels, 1)
        self.out_act = nn.Tanh()
        
    def _double_conv(self, in_ch, out_ch):
        return nn.Sequential(
            nn.Conv1d(in_ch, out_ch, 3, padding=1),
            nn.BatchNorm1d(out_ch),
            nn.LeakyReLU(0.2, inplace=True),
            nn.Conv1d(out_ch, out_ch, 3, padding=1),
            nn.BatchNorm1d(out_ch),
            nn.LeakyReLU(0.2, inplace=True),
        )
    
    def forward(self, x: torch.Tensor, mask: torch.Tensor) -> torch.Tensor:
        """
        Args:
            x: [B, 1, L] 退化rPPG信号
            mask: [B, 1, L] 二值掩码（1=有效，0=退化段）
            
        Returns:
            restored: [B, 1, L] 修复后信号
        """
        # 拼接信号和掩码
        inp = torch.cat([x, mask], dim=1)  # [B, 2, L]
        
        # 编码器
        e1 = self.enc1(inp)           # [B, 64, L]
        e2 = self.enc2(F.max_pool1d(e1, 2))  # [B, 128, L/2]
        e3 = self.enc3(F.max_pool1d(e2, 2))  # [B, 256, L/4]
        e4 = self.enc4(F.max_pool1d(e3, 2))  # [B, 512, L/8]
        
        # 瓶颈
        b = self.bottleneck(F.max_pool1d(e4, 2))  # [B, 1024, L/16]
        
        # 解码器+跳连接
        d4 = self.upconv4(b)
        d4 = self.dec4(torch.cat([d4, e4], dim=1))
        
        d3 = self.upconv3(d4)
        d3 = self.dec3(torch.cat([d3, e3], dim=1))
        
        d2 = self.upconv2(d3)
        d2 = self.dec2(torch.cat([d2, e2], dim=1))
        
        d1 = self.upconv1(d2)
        d1 = self.dec1(torch.cat([d1, e1], dim=1))
        
        # 输出
        out = self.out_conv(d1)
        out = self.out_act(out)
        
        # 仅替换退化段，保留有效段原始信号
        restored = x * mask + out * (1 - mask)
        
        return restored


class WGANCritic(nn.Module):
    """
    WGAN判别器（Critic）
    
    使用Wasserstein距离而非JS散度，
    梯度惩罚确保Lipschitz连续性
    """
    
    def __init__(self, in_channels: int = 1):
        super().__init__()
        self.model = nn.Sequential(
            nn.Conv1d(in_channels, 64, 3, 2, 1),
            nn.LeakyReLU(0.2),
            
            nn.Conv1d(64, 128, 3, 2, 1),
            nn.BatchNorm1d(128),
            nn.LeakyReLU(0.2),
            
            nn.Conv1d(128, 256, 3, 2, 1),
            nn.BatchNorm1d(256),
            nn.LeakyReLU(0.2),
            
            nn.Conv1d(256, 512, 3, 2, 1),
            nn.BatchNorm1d(512),
            nn.LeakyReLU(0.2),
            
            nn.AdaptiveAvgPool1d(1),
            nn.Flatten(),
            nn.Linear(512, 1),
        )
    
    def forward(self, x: torch.Tensor) -> torch.Tensor:
        return self.model(x)


def gradient_penalty(critic, real, fake, mask, device='cpu'):
    """
    梯度惩罚（WGAN-GP）
    
    确保判别器满足Lipschitz连续性约束
    """
    batch_size = real.size(0)
    
    # 随机插值因子
    alpha = torch.rand(batch_size, 1, 1, device=device)
    
    # 仅在退化段计算梯度惩罚
    interpolated = (alpha * real + (1 - alpha) * fake).requires_grad_(True)
    
    critic_input = interpolated * (1 - mask) + real * mask
    critic_output = critic(critic_input)
    
    gradients = torch.autograd.grad(
        outputs=critic_output,
        inputs=interpolated,
        grad_outputs=torch.ones_like(critic_output),
        create_graph=True,
        retain_graph=True,
    )[0]
    
    gradients = gradients.view(batch_size, -1)
    gradient_norm = gradients.norm(2, dim=1)
    penalty = ((gradient_norm - 1) ** 2).mean()
    
    return penalty


# 训练循环
def train_wgan_inpainting(
    generator, critic,
    dataloader, epochs=100,
    lr_g=1e-4, lr_c=1e-4,
    lambda_gp=10,
    device='cpu'
):
    """
    WGAN修复模型训练（超参为示意）

    摘要确认论文使用 Wasserstein GAN + 梯度惩罚、
    选择性重建退化段；具体优化器超参以原文为准
    """
    opt_g = torch.optim.Adam(
        generator.parameters(), lr=lr_g, betas=(0.5, 0.9)
    )
    opt_c = torch.optim.Adam(
        critic.parameters(), lr=lr_c, betas=(0.5, 0.9)
    )
    
    for epoch in range(epochs):
        for batch in dataloader:
            clean, degraded, mask = batch
            clean = clean.to(device)
            degraded = degraded.to(device)
            mask = mask.to(device)
            
            # 1. 训练Critic（5次）
            for _ in range(5):
                opt_c.zero_grad()
                
                fake = generator(degraded, mask).detach()
                
                # Wasserstein距离
                real_score = critic(clean * (1 - mask) + degraded * mask)
                fake_score = critic(fake * (1 - mask) + degraded * mask)
                
                gp = gradient_penalty(
                    critic, clean, fake, mask, device
                )
                
                loss_c = -(real_score.mean() - fake_score.mean()) + lambda_gp * gp
                loss_c.backward()
                opt_c.step()
            
            # 2. 训练Generator
            opt_g.zero_grad()
            
            fake = generator(degraded, mask)
            fake_score = critic(fake * (1 - mask) + degraded * mask)
            
            # 对抗损失 + 重建损失
            loss_adv = -fake_score.mean()
            loss_recon = F.l1_loss(
                fake * (1 - mask), clean * (1 - mask)
            )
            
            loss_g = loss_adv + 100 * loss_recon  # 重建权重
            loss_g.backward()
            opt_g.step()
        
        if (epoch + 1) % 10 == 0:
            print(f"Epoch {epoch+1}: G_loss={loss_g:.4f}, "
                  f"C_loss={loss_c:.4f}, "
                  f"recon={loss_recon:.4f}")


# 合成数据生成
def generate_synthetic_ppg(
    n_samples=10000, length=1800, fs=30
):
    """
    生成合成PPG信号用于训练（示意）

    摘要确认论文用大规模合成数据集，
    心率范围 30-180 bpm，含多样伪影画像
    """
    np.random.seed(42)
    signals = []
    masks = []
    
    for i in range(n_samples):
        # 随机心率
        hr = np.random.uniform(40, 180)
        freq = hr / 60  # Hz
        
        # 时间轴
        t = np.arange(length) / fs
        
        # 基础脉搏波（高斯导数模型）
        pulse = np.exp(-((t * freq % 1) - 0.3) ** 2 / 0.01)
        pulse -= 0.5 * np.exp(-((t * freq % 1) - 0.5) ** 2 / 0.02)
        
        # 呼吸调制
        resp = np.sin(2 * np.pi * 0.25 * t) * 0.1
        
        # 基线漂移
        baseline = np.cumsum(np.random.randn(length)) * 0.001
        
        # 噪声
        noise = np.random.randn(length) * 0.05
        
        signal = pulse + resp + baseline + noise
        
        # 随机生成退化段
        mask = np.ones(length)
        n_artifacts = np.random.randint(1, 5)
        for _ in range(n_artifacts):
            start = np.random.randint(0, length - 100)
            duration = np.random.randint(30, 200)
            mask[start:start+duration] = 0
            # 在退化段添加大噪声
            signal[start:start+duration] += np.random.randn(duration) * 2.0
        
        signals.append(signal)
        masks.append(mask)
    
    return np.array(signals, dtype=np.float32), \
           np.array(masks, dtype=np.float32)


if __name__ == "__main__":
    import numpy as np
    
    # 生成合成数据
    print("生成合成PPG数据...")
    signals, masks = generate_synthetic_ppg(n_samples=100, length=600)
    
    # 创建退化信号
    degraded = signals * masks + np.random.randn(*signals.shape) * 0.5 * (1 - masks)
    
    # 初始化模型
    generator = PPGInpaintingGenerator()
    critic = WGANCritic()
    
    print(f"生成器参数: {sum(p.numel() for p in generator.parameters()):,}")
    print(f"判别器参数: {sum(p.numel() for p in critic.parameters()):,}")
    
    # 前向测试
    x = torch.from_numpy(degraded[:2]).unsqueeze(1)  # [2, 1, 600]
    m = torch.from_numpy(masks[:2]).unsqueeze(1)
    
    restored = generator(x, m)
    print(f"输入: {x.shape}")
    print(f"掩码: {m.shape}")
    print(f"修复: {restored.shape}")
    print(f"退化段MAE: {F.l1_loss(restored[~m.bool()], x[~m.bool()]):.4f}")

# 示意代码，不可直接运行（退化段检测+修复管道为工程示意，预训练权重不存在）
class rPPGQualityAssessment:
    """
    rPPG信号质量评估+修复管道
    
    1. 实时检测信号退化段
    2. WGAN修复退化段
    3. 输出连续心率估计
    """
    
    def __init__(self, threshold=0.3):
        self.quality_threshold = threshold
        self.inpaintor = PPGInpaintingGenerator()
        # 加载预训练权重
        # self.inpaintor.load_state_dict(...)
        
    def assess_quality(self, rppg_signal):
        """
        评估rPPG信号质量
        
        返回二值掩码：1=有效, 0=退化
        """
        # 基于SNR的退化检测
        window = 30  # 1秒@30fps
        mask = np.ones(len(rppg_signal))
        
        for i in range(0, len(rppg_signal) - window, window//2):
            segment = rppg_signal[i:i+window]
            snr = self._compute_snr(segment)
            if snr < self.quality_threshold:
                mask[i:i+window] = 0
        
        return mask
    
    def _compute_snr(self, segment):
        """计算信噪比"""
        # FFT找到心率峰值
        fft = np.abs(np.fft.rfft(segment))
        freqs = np.fft.rfftfreq(len(segment), d=1/30)
        
        # 心率频段（0.7-4Hz）
        hr_mask = (freqs >= 0.7) & (freqs <= 4.0)
        signal_power = np.max(fft[hr_mask])
        noise_power = np.median(fft[~hr_mask])
        
        return signal_power / (noise_power + 1e-8)
    
    def restore(self, rppg_signal):
        """修复退化信号"""
        mask = self.assess_quality(rppg_signal)
        
        if (mask == 0).sum() == 0:
            return rppg_signal  # 无退化
        
        # WGAN修复
        x = torch.from_numpy(rppg_signal).float().unsqueeze(0).unsqueeze(0)
        m = torch.from_numpy(mask).float().unsqueeze(0).unsqueeze(0)
        
        with torch.no_grad():
            restored = self.inpaintor(x, m)
        
        return restored.squeeze().numpy()