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

class SMPL2PressureLoss(nn.Module):
    """
    Combined loss for SMPL2Pressure cVAE.
    - Pressure Map Reconstruction (PMR): MSE Loss
    - Point Cloud Reconstruction (PCR): MSE Loss (from condition features)
    - KL Divergence: Regularization of latent space
    """
    def __init__(self, cfg):
        super(SMPL2PressureLoss, self).__init__()
        self.cfg_loss = cfg['training']['loss']
        
        # 权重配置
        self.pmr_weight = self.cfg_loss.get('pmr_weight', 10.0)
        self.pcr_weight = self.cfg_loss.get('pcr_weight', 6.0)
        self.kl_weight = self.cfg_loss.get('kl_weight', 2.0)

    def forward(self, outputs, target_pressure, target_vertices):
        """
        Args:
            outputs: cVAE模型的输出字典
            target_pressure: 真值压力图 (B, H, W) 或 (B, 1, H, W)
            target_vertices: 真值SMPL顶点 (B, 6890, 3)
        """
        # 1. 压力图重建损失 (PMR)
        recon_pressure = outputs['recon_pressure']
        # 统一维度: 确保 target 也是 (B, 1, H, W)
        if recon_pressure.dim() == 3:
            recon_pressure = recon_pressure.unsqueeze(1)

        if target_pressure.dim() == 3:
            target_pressure = target_pressure.unsqueeze(1)
        
        loss_pmr = F.mse_loss(recon_pressure, target_pressure, reduction='mean')

        # 2. 点云重建损失 (PCR)
        # 根据你的提醒，这里的 recon_vertices 应该是从 cond_features 解码出来的
        recon_vertices = outputs['recon_vertices']
        loss_pcr = F.mse_loss(recon_vertices, target_vertices, reduction='mean')

        # 3. KL 散度损失
        mu = outputs['mu']
        log_var = outputs['log_var']
        # KL = -0.5 * sum(1 + log(sigma^2) - mu^2 - sigma^2)
        loss_kl = -0.5 * torch.mean(torch.sum(1 + log_var - mu.pow(2) - log_var.exp(), dim=1))

        # 4. 总损失加权
        total_loss = (self.pmr_weight * loss_pmr + 
                      self.pcr_weight * loss_pcr + 
                      self.kl_weight * loss_kl)

        return {
            'loss': total_loss,
            'loss_pmr': loss_pmr,
            'loss_pcr': loss_pcr,
            'loss_kl': loss_kl
        }