import math
import torch
from torch import nn
from torch.nn import functional as F
import torch.distributions as dist
import numpy as np
import copy
from modules.audio2motion.flow_base import Glow, WN, ResidualCouplingBlock
from modules.audio2motion.transformer_base import Embedding

from utils.commons.pitch_utils import f0_to_coarse
from utils.commons.hparams import hparams


class LambdaLayer(nn.Module):
    def __init__(self, lambd):
        super(LambdaLayer, self).__init__()
        self.lambd = lambd

    def forward(self, x):
        return self.lambd(x)


def make_positions(tensor, padding_idx):
    """Replace non-padding symbols with their position numbers.

    Position numbers begin at padding_idx+1. Padding symbols are ignored.
    """
    # The series of casts and type-conversions here are carefully
    # balanced to both work with ONNX export and XLA. In particular XLA
    # prefers ints, cumsum defaults to output longs, and ONNX doesn't know
    # how to handle the dtype kwarg in cumsum.
    mask = tensor.ne(padding_idx).int()
    return (
                   torch.cumsum(mask, dim=1).type_as(mask) * mask
           ).long() + padding_idx

class SinusoidalPositionalEmbedding(nn.Module):
    """This module produces sinusoidal positional embeddings of any length.

    Padding symbols are ignored.
    """

    def __init__(self, embedding_dim, padding_idx, init_size=1024):
        super().__init__()
        self.embedding_dim = embedding_dim
        self.padding_idx = padding_idx
        self.weights = SinusoidalPositionalEmbedding.get_embedding(
            init_size,
            embedding_dim,
            padding_idx,
        )
        self.register_buffer('_float_tensor', torch.FloatTensor(1))

    @staticmethod
    def get_embedding(num_embeddings, embedding_dim, padding_idx=None):
        """Build sinusoidal embeddings.

        This matches the implementation in tensor2tensor, but differs slightly
        from the description in Section 3.5 of "Attention Is All You Need".
        """
        half_dim = embedding_dim // 2
        emb = math.log(10000) / (half_dim - 1)
        emb = torch.exp(torch.arange(half_dim, dtype=torch.float) * -emb)
        emb = torch.arange(num_embeddings, dtype=torch.float).unsqueeze(1) * emb.unsqueeze(0)
        emb = torch.cat([torch.sin(emb), torch.cos(emb)], dim=1).view(num_embeddings, -1)
        if embedding_dim % 2 == 1:
            # zero pad
            emb = torch.cat([emb, torch.zeros(num_embeddings, 1)], dim=1)
        if padding_idx is not None:
            emb[padding_idx, :] = 0
        return emb

    def forward(self, input, incremental_state=None, timestep=None, positions=None, **kwargs):
        """Input is expected to be of size [bsz x seqlen]."""
        bsz, seq_len = input.shape[:2]
        max_pos = self.padding_idx + 1 + seq_len
        if self.weights is None or max_pos > self.weights.size(0):
            # recompute/expand embeddings if needed
            self.weights = SinusoidalPositionalEmbedding.get_embedding(
                max_pos,
                self.embedding_dim,
                self.padding_idx,
            )
        self.weights = self.weights.to(self._float_tensor)

        if incremental_state is not None:
            # positions is the same for every token when decoding a single step
            pos = timestep.view(-1)[0] + 1 if timestep is not None else seq_len
            return self.weights[self.padding_idx + pos, :].expand(bsz, 1, -1)

        positions = make_positions(input, self.padding_idx) if positions is None else positions
        return self.weights.index_select(0, positions.view(-1)).view(bsz, seq_len, -1).detach()

    def max_positions(self):
        """Maximum number of supported positions."""
        return int(1e4)  # an arbitrary large number

class FVAEEncoder(nn.Module):
    def __init__(self, in_channels, hidden_channels, latent_channels, kernel_size,
                 n_layers, gin_channels=0, p_dropout=0, strides=[4]):
        super().__init__()
        self.strides = strides
        self.hidden_size = hidden_channels
        self.pre_net = nn.Sequential(*[
            nn.Conv1d(in_channels, hidden_channels, kernel_size=s * 2, stride=s, padding=s // 2)
            if i == 0 else
            nn.Conv1d(hidden_channels, hidden_channels, kernel_size=s * 2, stride=s, padding=s // 2)
            for i, s in enumerate(strides)
        ])
        self.wn = WN(hidden_channels, kernel_size, 1, n_layers, gin_channels, p_dropout)
        self.out_proj = nn.Conv1d(hidden_channels, latent_channels * 2, 1)

        self.latent_channels = latent_channels

    def forward(self, x, x_mask, g):
        x = self.pre_net(x)
        x_mask = x_mask[:, :, ::np.prod(self.strides)][:, :, :x.shape[-1]]
        x = x * x_mask
        x = self.wn(x, x_mask, g) * x_mask
        x = self.out_proj(x)
        m, logs = torch.split(x, self.latent_channels, dim=1)
        z = (m + torch.randn_like(m) * torch.exp(logs))
        return z, m, logs, x_mask


class FVAEDecoder(nn.Module):
    def __init__(self, latent_channels, hidden_channels, out_channels, kernel_size,
                 n_layers, gin_channels=0, p_dropout=0,
                 strides=[4]):
        super().__init__()
        self.strides = strides
        self.hidden_size = hidden_channels
        self.pre_net = nn.Sequential(*[
            nn.ConvTranspose1d(latent_channels, hidden_channels, kernel_size=s, stride=s)
            if i == 0 else
            nn.ConvTranspose1d(hidden_channels, hidden_channels, kernel_size=s, stride=s)
            for i, s in enumerate(strides)
        ])
        self.wn = WN(hidden_channels, kernel_size, 1, n_layers, gin_channels, p_dropout)
        self.out_proj = nn.Conv1d(hidden_channels, out_channels, 1)

    def forward(self, x, x_mask, g):
        x = self.pre_net(x)
        x = x * x_mask
        x = self.wn(x, x_mask, g) * x_mask
        x = self.out_proj(x)
        return x

class FVAE(nn.Module):
    def __init__(self,
                 in_out_channels=64, hidden_channels=256, latent_size=16,
                 kernel_size=3, enc_n_layers=5, dec_n_layers=5, gin_channels=80, strides=[4,],
                 use_prior_glow=True, glow_hidden=256, glow_kernel_size=3, glow_n_blocks=5,
                 sqz_prior=False, use_pos_emb=False):
        super(FVAE, self).__init__()
        self.in_out_channels = in_out_channels
        self.strides = strides
        self.hidden_size = hidden_channels
        self.latent_size = latent_size
        self.use_prior_glow = use_prior_glow
        self.sqz_prior = sqz_prior
        self.g_pre_net = nn.Sequential(*[
            nn.Conv1d(gin_channels, gin_channels, kernel_size=s * 2, stride=s, padding=s // 2)
            for i, s in enumerate(strides)
        ])
        self.encoder = FVAEEncoder(in_out_channels, hidden_channels, latent_size, kernel_size,
                                   enc_n_layers, gin_channels, strides=strides)
        if use_prior_glow:
            self.prior_flow = ResidualCouplingBlock(
                latent_size, glow_hidden, glow_kernel_size, 1, glow_n_blocks, 4, gin_channels=gin_channels)
        self.use_pos_embed = use_pos_emb
        if sqz_prior:
            self.query_proj = nn.Linear(latent_size, latent_size)
            self.key_proj = nn.Linear(latent_size, latent_size)
            self.value_proj = nn.Linear(latent_size, hidden_channels)
            if self.in_out_channels in [7, 64]:
                self.decoder = FVAEDecoder(hidden_channels, hidden_channels, in_out_channels, kernel_size,
                                    dec_n_layers, gin_channels, strides=strides)
            elif self.in_out_channels == 71:
                self.exp_decoder = FVAEDecoder(hidden_channels, hidden_channels, 64, kernel_size,
                                    dec_n_layers, gin_channels, strides=strides)
                self.pose_decoder = FVAEDecoder(hidden_channels, hidden_channels, 7, kernel_size,
                                    dec_n_layers, gin_channels, strides=strides)
            if self.use_pos_embed:
                self.embed_positions = SinusoidalPositionalEmbedding(self.latent_size, 0,init_size=2000+1,)
        else:
            self.decoder = FVAEDecoder(latent_size, hidden_channels, in_out_channels, kernel_size,
                                    dec_n_layers, gin_channels, strides=strides)

        self.prior_dist = dist.Normal(0, 1)

    def forward(self, x=None, x_mask=None, g=None, infer=False, temperature=1. , **kwargs):
        """

        :param x: [B, T,  C_in_out]
        :param x_mask: [B, T]
        :param g: [B, T, C_g]
        :return:
        """
        x_mask = x_mask[:, None, :] # [B, 1, T]
        g = g.transpose(1,2) # [B, C_g, T]
        g_for_sqz = g

        g_sqz = self.g_pre_net(g_for_sqz)

        if not infer:
            x = x.transpose(1,2) # [B, C, T]
            z_q, m_q, logs_q, x_mask_sqz = self.encoder(x, x_mask, g_sqz)
            if self.sqz_prior:
                z = z_q
                if self.use_pos_embed:
                    position = self.embed_positions(z.transpose(1,2).abs().sum(-1)).transpose(1,2)
                    z = z + position
                q = self.query_proj(z.mean(dim=-1,keepdim=True).transpose(1,2)) # [B, 1, C=16]
                k = self.key_proj(z.transpose(1,2)) # [B, T, C=16]
                v = self.value_proj(z.transpose(1,2)) # [B, T, C=256]
                attn = torch.bmm(q,k.transpose(1,2)) # [B, 1, T]
                attn = F.softmax(attn, dim=-1)
                out = torch.bmm(attn, v) # [B, 1, C=256]
                style_encoding = out.repeat([1,z_q.shape[-1],1]).transpose(1,2) # [B, C=256, T]
                if self.in_out_channels == 71:
                    x_recon = torch.cat([self.exp_decoder(style_encoding, x_mask, g), self.pose_decoder(style_encoding, x_mask, g)], dim=1)
                else:
                    x_recon = self.decoder(style_encoding, x_mask, g)
            else:
                if self.in_out_channels == 71:
                    x_recon = torch.cat([self.exp_decoder(z_q, x_mask, g), self.pose_decoder(z_q, x_mask, g)], dim=1)
                else:
                    x_recon = self.decoder(z_q, x_mask, g)
            q_dist = dist.Normal(m_q, logs_q.exp())
            if self.use_prior_glow:
                logqx = q_dist.log_prob(z_q)
                z_p = self.prior_flow(z_q, x_mask_sqz, g_sqz)
                logpx = self.prior_dist.log_prob(z_p)
                loss_kl = ((logqx - logpx) * x_mask_sqz).sum() / x_mask_sqz.sum() / logqx.shape[1]
            else:
                loss_kl = torch.distributions.kl_divergence(q_dist, self.prior_dist)
                loss_kl = (loss_kl * x_mask_sqz).sum() / x_mask_sqz.sum() / z_q.shape[1]
                z_p = z_q
            return x_recon.transpose(1,2), loss_kl, z_p.transpose(1,2), m_q.transpose(1,2), logs_q.transpose(1,2)
        else:
            latent_shape = [g_sqz.shape[0], self.latent_size, g_sqz.shape[2]]
            z_p = self.prior_dist.sample(latent_shape).to(g.device) * temperature # [B, latent_size, T_sqz]
            if self.use_prior_glow:
                z_p = self.prior_flow(z_p, 1, g_sqz, reverse=True)
            if self.sqz_prior:
                z = z_p
                if self.use_pos_embed:
                    position = self.embed_positions(z.abs().sum(-1))
                    z += position
                q = self.query_proj(z.mean(dim=-1,keepdim=True).transpose(1,2)) # [B, 1, C=16]
                k = self.key_proj(z.transpose(1,2)) # [B, T, C=16]
                v = self.value_proj(z.transpose(1,2)) # [B, T, C=256]
                attn = torch.bmm(q,k.transpose(1,2)) # [B, 1, T]
                attn = F.softmax(attn, dim=-1)
                out = torch.bmm(attn, v) # [B, 1, C=256]
                style_encoding = out.repeat([1,z_p.shape[-1],1]).transpose(1,2) # [B, C=256, T]
                x_recon = self.decoder(style_encoding, 1, g)
                if self.in_out_channels == 71:
                    x_recon = torch.cat([self.exp_decoder(style_encoding, 1, g), self.pose_decoder(style_encoding, 1, g)], dim=1)
                else:
                    x_recon = self.decoder(style_encoding, 1, g)
            else:
                if self.in_out_channels == 71:
                    x_recon = torch.cat([self.exp_decoder(z_p, 1, g), self.pose_decoder(z_p, 1, g)], dim=1)
                else:
                    x_recon = self.decoder(z_p, 1, g)
            return x_recon.transpose(1,2), z_p.transpose(1,2)


class VAEModel(nn.Module):
    def __init__(self, in_out_dim=64, audio_in_dim=1024, sqz_prior=False, cond_drop=False, use_prior_flow=True):
        super().__init__()
        feat_dim = 64
        self.blink_embed = nn.Embedding(2, feat_dim)
        self.audio_in_dim = audio_in_dim
        cond_dim = feat_dim
        self.mel_encoder = nn.Sequential(*[
                nn.Conv1d(audio_in_dim, 64, 3, 1, 1, bias=False),
                nn.BatchNorm1d(64),
                nn.GELU(),
                nn.Conv1d(64, feat_dim, 3, 1, 1, bias=False)
            ]) 
        self.cond_drop = cond_drop
        if self.cond_drop:
            self.dropout = nn.Dropout(0.5)

        self.in_dim, self.out_dim = in_out_dim, in_out_dim
        self.sqz_prior = sqz_prior
        self.use_prior_flow = use_prior_flow
        self.vae = FVAE(in_out_channels=in_out_dim, hidden_channels=256, latent_size=16, kernel_size=5,
            enc_n_layers=8, dec_n_layers=4, gin_channels=cond_dim, strides=[4,],
            use_prior_glow=self.use_prior_flow, glow_hidden=64, glow_kernel_size=3, glow_n_blocks=4,sqz_prior=sqz_prior)
        self.downsampler = LambdaLayer(lambda x: F.interpolate(x.transpose(1,2), scale_factor=0.5, mode='linear').transpose(1,2))
        # self.downsampler = LambdaLayer(lambda x: F.interpolate(x.transpose(1,2), scale_factor=0.5, mode='nearest').transpose(1,2))

    def num_params(self, model, print_out=True, model_name="model"):
        parameters = filter(lambda p: p.requires_grad, model.parameters())
        parameters = sum([np.prod(p.size()) for p in parameters]) / 1_000_000
        if print_out:
            print(f'| {model_name} Trainable Parameters: %.3fM' % parameters)
        return parameters
    
    @property
    def device(self):
        return self.vae.parameters().__next__().device

    def forward(self, batch, ret, train=True, return_latent=False, temperature=1.):
        infer = not train
        mask = batch['y_mask'].to(self.device)
        mel = batch['audio'].to(self.device)
        mel = self.downsampler(mel)
        cond_feat = self.mel_encoder(mel.transpose(1,2)).transpose(1,2)

        if self.cond_drop:
            cond_feat = self.dropout(cond_feat)
        
        if not infer:
            exp = batch['y'].to(self.device)
            x = exp
            x_recon, loss_kl, z_p, m_q, logs_q = self.vae(x=x, x_mask=mask, g=cond_feat, infer=False)
            x_recon = x_recon * mask.unsqueeze(-1)
            ret['pred'] = x_recon
            ret['mask'] = mask
            ret['loss_kl'] = loss_kl
            if return_latent:
                ret['m_q'] = m_q
                ret['z_p'] = z_p
            return x_recon, loss_kl, m_q, logs_q
        else:
            x_recon, z_p = self.vae(x=None, x_mask=mask, g=cond_feat, infer=True, temperature=temperature)
            x_recon = x_recon * mask.unsqueeze(-1)
            ret['pred'] = x_recon
            ret['mask'] = mask

            return x_recon


class PitchContourVAEModel(nn.Module):
    def __init__(self, hparams, in_out_dim=64, audio_in_dim=1024, sqz_prior=False, cond_drop=False, use_prior_flow=True):
        super().__init__()
        self.hparams = copy.deepcopy(hparams)
        feat_dim = 128
        self.audio_in_dim = audio_in_dim
        self.blink_embed = nn.Embedding(2, feat_dim)
        
        self.mel_encoder = nn.Sequential(*[
                nn.Conv1d(audio_in_dim, feat_dim, 3, 1, 1, bias=False),
                nn.BatchNorm1d(feat_dim ),
                nn.GELU(),
                nn.Conv1d(feat_dim , feat_dim, 3, 1, 1, bias=False)
            ])
        
        self.pitch_embed = Embedding(300, feat_dim, None)
        self.pitch_encoder = nn.Sequential(*[
                nn.Conv1d(feat_dim, feat_dim , 3, 1, 1, bias=False),
                nn.BatchNorm1d(feat_dim),
                nn.GELU(),
                nn.Conv1d(feat_dim, feat_dim, 3, 1, 1, bias=False)
            ])

        cond_dim = feat_dim + feat_dim + feat_dim

        if hparams.get('use_mouth_amp_embed', False):
            self.mouth_amp_embed = nn.Parameter(torch.randn(feat_dim))
            cond_dim += feat_dim

        if hparams.get('use_eye_amp_embed', False):
            self.eye_amp_embed = nn.Parameter(torch.randn(feat_dim))
            cond_dim += feat_dim

        self.cond_proj = nn.Linear(cond_dim, feat_dim, bias=True)

        self.cond_drop = cond_drop
        if self.cond_drop:
            self.dropout = nn.Dropout(0.5)

        self.in_dim, self.out_dim = in_out_dim, in_out_dim
        self.sqz_prior = sqz_prior
        self.use_prior_flow = use_prior_flow
        self.vae = FVAE(in_out_channels=in_out_dim, hidden_channels=256, latent_size=16, kernel_size=5,
            enc_n_layers=8, dec_n_layers=4, gin_channels=feat_dim, strides=[4,],
            use_prior_glow=self.use_prior_flow, glow_hidden=64, glow_kernel_size=3, glow_n_blocks=4,sqz_prior=sqz_prior)
        self.downsampler = LambdaLayer(lambda x: F.interpolate(x.transpose(1,2), scale_factor=0.5, mode='nearest').transpose(1,2))

    def num_params(self, model, print_out=True, model_name="model"):
        parameters = filter(lambda p: p.requires_grad, model.parameters())
        parameters = sum([np.prod(p.size()) for p in parameters]) / 1_000_000
        if print_out:
            print(f'| {model_name} Trainable Parameters: %.3fM' % parameters)
        return parameters
    
    @property
    def device(self):
        return self.vae.parameters().__next__().device

    def forward(self, batch, ret, train=True, return_latent=False, temperature=1.):
        infer = not train
        hparams = self.hparams
        mask = batch['y_mask'].to(self.device)
        mel = batch['audio'].to(self.device)
        f0 = batch['f0'].to(self.device) # [b,t]
        if 'blink' not in batch:
            batch['blink'] = torch.zeros([f0.shape[0], f0.shape[1], 1], dtype=torch.long, device=f0.device)
        blink = batch['blink'].to(self.device)
        blink_feat = self.blink_embed(blink.squeeze(2))

        blink_feat  = self.downsampler(blink_feat)
        mel = self.downsampler(mel)
        f0 = self.downsampler(f0.unsqueeze(-1)).squeeze(-1)
        f0_coarse = f0_to_coarse(f0)
        pitch_emb = self.pitch_embed(f0_coarse)
        cond_feat = self.mel_encoder(mel.transpose(1,2)).transpose(1,2)
        pitch_feat = self.pitch_encoder(pitch_emb.transpose(1,2)).transpose(1,2)

        cond_feats = [cond_feat, pitch_feat, blink_feat]
        if hparams.get('use_mouth_amp_embed', False):
            mouth_amp = batch.get('mouth_amp', torch.ones([f0.shape[0], 1], device=f0.device) * 0.4)
            mouth_amp_feat = mouth_amp.unsqueeze(1) * self.mouth_amp_embed.unsqueeze(0)
            mouth_amp_feat = mouth_amp_feat.repeat([1,cond_feat.shape[1],1])
            cond_feats.append(mouth_amp_feat)

        if hparams.get('use_eye_amp_embed', False):
            eye_amp = batch.get('eye_amp', torch.ones([f0.shape[0], 1], device=f0.device) * 0.4)
            eye_amp_feat = eye_amp.unsqueeze(1) * self.eye_amp_embed.unsqueeze(0)
            eye_amp_feat = eye_amp_feat.repeat([1,cond_feat.shape[1],1])
            cond_feats.append(eye_amp_feat)

        cond_feat = torch.cat(cond_feats, dim=-1)
        cond_feat = self.cond_proj(cond_feat)

        if self.cond_drop:
            cond_feat = self.dropout(cond_feat)
        
        if not infer:
            exp = batch['y'].to(self.device)
            x = exp
            x_recon, loss_kl, z_p, m_q, logs_q = self.vae(x=x, x_mask=mask, g=cond_feat, infer=False)
            x_recon = x_recon * mask.unsqueeze(-1)
            ret['pred'] = x_recon
            ret['mask'] = mask
            ret['loss_kl'] = loss_kl
            if return_latent:
                ret['m_q'] = m_q
                ret['z_p'] = z_p
            return x_recon, loss_kl, m_q, logs_q
        else:
            x_recon, z_p = self.vae(x=None, x_mask=mask, g=cond_feat, infer=True, temperature=temperature)
            x_recon = x_recon * mask.unsqueeze(-1)
            ret['pred'] = x_recon
            ret['mask'] = mask

            return x_recon


if __name__ == '__main__':
    model = FVAE(in_out_channels=64, hidden_channels=128, latent_size=32,kernel_size=3, enc_n_layers=6, dec_n_layers=2, 
        gin_channels=80, strides=[4], use_prior_glow=False, glow_hidden=128, glow_kernel_size=3, glow_n_blocks=3)
    x = torch.rand([8, 64, 1000])
    x_mask = torch.ones([8,1,1000])
    g = torch.rand([8, 80, 1000])
    train_out = model(x,x_mask,g,infer=False)
    x_recon, loss_kl, z_p, m_q, logs_q = train_out
    print(" ")
    infer_out = model(x,x_mask,g,infer=True)
    x_recon, z_p = infer_out
    print(" ")