import torch
import torch.nn as nn
import torch.nn.functional as F
import numpy as np

class ReConsLoss(nn.Module):
    def __init__(self, motion_dim=272):
        super(ReConsLoss, self).__init__()
        self.motion_dim = motion_dim
    
    def softclip(self, tensor, min):
        result_tensor = min + F.softplus(tensor - min)
        return result_tensor
    
    def gaussian_nll(self, mu, log_sigma, x):
        return 0.5 * torch.pow((x - mu) / log_sigma.exp(), 2) + log_sigma + 0.5 * np.log(2 * np.pi)
    
    def forward(self, motion_pred, motion_gt) : 
        """Optimal sigma VAE loss, see https://arxiv.org/pdf/2006.13202 for more details"""
        log_sigma = ((motion_gt[..., :self.motion_dim] - motion_pred[..., :self.motion_dim]) ** 2).mean([0,1,2], keepdim=True).sqrt().log()
        log_sigma = self.softclip(log_sigma, -6)
        loss = self.gaussian_nll(motion_pred[..., :self.motion_dim], log_sigma, motion_gt[..., :self.motion_dim]).sum()
        return loss
    
    
    def forward_KL(self, mu, logvar):
        loss = -0.5 * torch.sum(1 + logvar - mu.pow(2) - logvar.exp(), dim=(1, 2))
        return loss.mean()
    
    def forward_root(self, motion_pred, motion_gt):
        """[..., :8] relate to the root joint"""
        root_log_sigma = ((motion_gt[..., :8] - motion_pred[..., :8]) ** 2).mean([0,1,2], keepdim=True).sqrt().log()
        root_log_sigma = self.softclip(root_log_sigma, -6)
        root_loss = self.gaussian_nll(motion_pred[..., :8], root_log_sigma, motion_gt[..., :8]).sum()           
        return root_loss