import torch

from ..builder import ARCHITECTURES, build_loss, build_submodule
from .base_architecture import BaseArchitecture


@ARCHITECTURES.register_module()
class PoseVAE(BaseArchitecture):

    def __init__(self,
                 encoder=None,
                 decoder=None,
                 loss_recon=None,
                 kl_div_loss_weight=None,
                 init_cfg=None,
                 **kwargs):
        super().__init__(init_cfg=init_cfg, **kwargs)
        self.encoder = build_submodule(encoder)
        self.decoder = build_submodule(decoder)
        self.loss_recon = build_loss(loss_recon)
        self.kl_div_loss_weight = kl_div_loss_weight

    def reparameterize(self, mu, logvar):
        std = torch.exp(logvar / 2)

        eps = std.data.new(std.size()).normal_()
        latent_code = eps.mul(std).add_(mu)
        return latent_code

    def encode(self, pose):
        mu, logvar = self.encoder(pose)
        return mu

    def forward(self, **kwargs):
        motion = kwargs['motion'].float()
        B, T = motion.shape[:2]
        pose = motion.reshape(B * T, -1)
        pose = pose[:, :-4]

        mu, logvar = self.encoder(pose)
        z = self.reparameterize(mu, logvar)
        pred = self.decoder(z)

        loss = dict()
        recon_loss = self.loss_recon(pred, pose, reduction_override='none')
        loss['recon_loss'] = recon_loss
        if self.kl_div_loss_weight is not None:
            loss_kl = -0.5 * torch.sum(1 + logvar - mu.pow(2) - logvar.exp())
            loss['kl_div_loss'] = (loss_kl * self.kl_div_loss_weight)

        return loss


@ARCHITECTURES.register_module()
class MotionVAE(BaseArchitecture):

    def __init__(self,
                 encoder=None,
                 decoder=None,
                 loss_recon=None,
                 kl_div_loss_weight=None,
                 init_cfg=None,
                 **kwargs):
        super().__init__(init_cfg=init_cfg, **kwargs)
        self.encoder = build_submodule(encoder)
        self.decoder = build_submodule(decoder)
        self.loss_recon = build_loss(loss_recon)
        self.kl_div_loss_weight = kl_div_loss_weight

    def sample(self, std=1, latent_code=None):
        if latent_code is not None:
            z = latent_code
        else:
            z = torch.randn(1, 7, self.decoder.latent_dim).cuda() * std
        output = self.decoder(z)
        if self.use_normalization:
            output = output * self.motion_std
            output = output + self.motion_mean
        return output

    def reparameterize(self, mu, logvar):
        std = torch.exp(logvar / 2)

        eps = std.data.new(std.size()).normal_()
        latent_code = eps.mul(std).add_(mu)
        return latent_code

    def encode(self, motion, motion_mask):
        mu, logvar = self.encoder(motion, motion_mask)
        return self.reparameterize(mu, logvar)

    def decode(self, z, motion_mask):
        return self.decoder(z, motion_mask)

    def forward(self, **kwargs):
        motion, motion_mask = kwargs['motion'].float(), kwargs['motion_mask']
        B, T = motion.shape[:2]

        mu, logvar = self.encoder(motion, motion_mask)
        z = self.reparameterize(mu, logvar)
        pred = self.decoder(z, motion_mask)

        loss = dict()
        recon_loss = self.loss_recon(pred, motion, reduction_override='none')
        recon_loss = recon_loss.mean(dim=-1) * motion_mask
        recon_loss = recon_loss.sum() / motion_mask.sum()
        loss['recon_loss'] = recon_loss
        if self.kl_div_loss_weight is not None:
            loss_kl = -0.5 * torch.sum(1 + logvar - mu.pow(2) - logvar.exp())
            loss['kl_div_loss'] = (loss_kl * self.kl_div_loss_weight)

        return loss