from typing import List, Optional

import torch
import torch.nn as nn
from torch import Tensor

from mtts.models.decoder import Decoder
from mtts.models.encoder import FS2TransformerEncoder
from mtts.models.fs2_variance import VarianceAdaptor
from mtts.models.postnet import PostNet
from mtts.utils.logging import get_logger

ENCODERS = [
    FS2TransformerEncoder,
]

logger = get_logger(__file__)


def __read_vocab(file):
    with open(file) as f:
        lines = f.read().split('\n')
    lines = [line for line in lines if len(line) > 0]
    return lines


def _get_layer(emb_config: dict):  # -> Optional[List[nn.Module],List[float]]:
    logger.info(f'building embedding with config: {emb_config}')
    if emb_config['enable']:
        if emb_config['vocab'] is None:
            vocab_size = emb_config['vocab_size']
        else:
            vocab = __read_vocab(emb_config['vocab'])
            vocab_size = len(vocab)
        layer = nn.Embedding(vocab_size, emb_config['dim'], padding_idx=0)
        return layer, emb_config['weight']
    else:
        return None, None


def _build_embedding_layers(config):
    layers = nn.ModuleList()
    weights = []
    for c in (
            config['pinyin_embedding'],
            config['hanzi_embedding'],
            config['speaker_embedding'],
    ):

        layer, weight = _get_layer(c)
        if layer is not None:
            layers.append(layer)
            weights.append(weight)

    return layers, weights


def get_mask_from_lengths(lengths, max_len=None):
    batch_size = lengths.shape[0]
    if max_len is None:
        max_len = torch.max(lengths).item()
    ids = torch.arange(0, max_len).unsqueeze(0).expand(batch_size, -1).to(lengths.device)
    mask = (ids >= lengths.unsqueeze(1).expand(-1, max_len))

    return mask


class FastSpeech2(nn.Module):
    """ FastSpeech2 """
    def __init__(self, config):
        super(FastSpeech2, self).__init__()

        emb_layers, emb_weights = _build_embedding_layers(config)

        EncoderClass = eval(config['encoder']['encoder_type'])
        assert EncoderClass in ENCODERS
        encoder_conf = config['encoder']['conf']
        encoder_conf.update({'emb_layers': emb_layers})
        encoder_conf.update({'embeding_weights': emb_weights})
        self.encoder = EncoderClass(**encoder_conf)

        dur_config = config['duration_predictor']
        self.variance_adaptor = VarianceAdaptor(**dur_config)

        decoder_config = config['decoder']
        self.decoder = Decoder(**decoder_config)

        n_mels = config['fbank']['n_mels']
        mel_linear_input_dim = decoder_config['hidden_dim']
        self.mel_linear = nn.Linear(mel_linear_input_dim, n_mels)
        self.postnet = PostNet(**config['postnet'], )

    def forward(self,
                input_seqs: List[Tensor],
                seq_len: Tensor,
                mel_len: Optional[Tensor] = None,
                d_target: Optional[Tensor] = None,
                max_src_len=None,
                max_mel_len=None,
                d_control=1.0,
                p_control=1.0,
                e_control=1.0):
        src_mask = get_mask_from_lengths(seq_len, max_src_len)

        if mel_len is not None:
            mel_mask = get_mask_from_lengths(mel_len, max_mel_len)
        else:
            mel_mask = None

        encoder_output = self.encoder(input_seqs, src_mask)

        if d_target is not None:
            variance_adaptor_output, d_prediction, _, _ = self.variance_adaptor(encoder_output, src_mask, mel_mask,
                                                                                d_target, max_mel_len, d_control)
        else:
            variance_adaptor_output, d_prediction, mel_len, mel_mask = self.variance_adaptor(
                encoder_output, src_mask, mel_mask, d_target, max_mel_len, d_control)

        decoder_output = self.decoder(variance_adaptor_output, mel_mask)

        mel_pred = self.mel_linear(decoder_output)
        postnet_input = torch.unsqueeze(mel_pred, 1)
        postnet_output = self.postnet(postnet_input) + mel_pred

        return mel_pred, postnet_output, d_prediction, src_mask, mel_mask, mel_len


if __name__ == "__main__":
    # Test
    import yaml
    with open('../../examples/aishell3/config.yaml') as f:
        config = yaml.safe_load(f)
    model = FastSpeech2(config)
    print(model)
    print(sum(param.numel() for param in model.parameters()))