import random
import torch 
import torch.nn as nn 

class AccousticTransformer(nn.Module):
    def __init__(self, max_seq_len, phonemes_vocab, emb_dim, n_head, max_time_frames):
        super(AccousticTransformer, self).__init__()
        self.embedding = nn.Embedding(phonemes_vocab, emb_dim) # shape: [batch, vocab_size, emb_dim]

        self.pos_embedding = nn.Embedding(max_seq_len, emb_dim) # shape: [bath, max_seq_len=256, emb_dim]

        self.mel_pos_embedding = nn.Embedding(max_time_frames, emb_dim) # same emb_dim is to be used as before

        self.encoder_layer = nn.TransformerEncoderLayer(
            d_model=emb_dim,
            nhead=n_head, 
            dim_feedforward=2048, 
            dropout=0.1,
            batch_first=True
        )

        self.encoder = nn.TransformerEncoder(
            self.encoder_layer, 
            num_layers=6
        )

        self.decoder_layer = nn.TransformerDecoderLayer(
            d_model=emb_dim,
            nhead=n_head,
            dim_feedforward=2048,
            dropout=0.1,
            batch_first=True
        )
        self.decoder = nn.TransformerDecoder(
            self.decoder_layer,
            num_layers=6
        )
        
        self.duration_predictor = DurationPredictor(emb_dim=emb_dim)
        self.length_regulator = LengthRegulator()

        self.output_layer = nn.Linear(emb_dim, 80)  # 80 is freq_bins

    def forward(self, X, gt_durations = None):
    
        positions = torch.arange(X.shape[1], device=X.device)
        X = self.embedding(X) + self.pos_embedding(positions)  # [batch, seq_len, emb_dim]
        encoder_output = self.encoder(X)  # [batch, seq_len, emb_dim]

        predicted_durations = self.duration_predictor(encoder_output)  # shape: [batch, seq_len]

        if gt_durations is not None:
            durations = gt_durations                        # training
        else:
            durations = predicted_durations.round().long()  # inference

        expanded = self.length_regulator(encoder_output, durations)

        mel_positions = torch.arange(expanded.shape[1], device=expanded.device)
        expanded = expanded + self.mel_pos_embedding(mel_positions)  # [batch, T, emb_dim]

        decoder_output = self.decoder(
                tgt=expanded,
                memory=encoder_output
            )  # shape: [batch, T, emb_dim]

        predicted_mel = self.output_layer(decoder_output) # shape: [batch, T, 80]
        predicted_mel = predicted_mel.transpose(1,2)      # shape: [batch, 80, T]

        return predicted_mel, predicted_durations


class DurationPredictor(nn.Module):
    def __init__(self, emb_dim):
        super(DurationPredictor, self).__init__()
        self.conv1 = nn.Conv1d(
            in_channels= emb_dim,
            out_channels= emb_dim, 
            kernel_size= 3, 
            padding= 1
        )

        self.conv2 = nn.Conv1d(
            in_channels= emb_dim,
            out_channels= emb_dim, 
            kernel_size= 3, 
            padding= 1
        )

        self.norm1 = nn.LayerNorm(emb_dim)  # this is the output of the conv1 as that layer 
                                            # does makes the change in the emb_dim 

        self.norm2 = nn.LayerNorm(emb_dim)

        self.linear = nn.Linear(emb_dim, 1)  # this comes int the play as when the 2 conv layers does acts on emb_dim
                                             # then this emb_dim till this TIME does have the rich representation
                                             # as it catures the what phoneme it is , what it's neighours are and 
                                             # what is the context from the encoder

                                             # so hence in_features : emb_dim, out_features : 1
                                            

        self.output_activation = nn.ReLU()   # to ensure the OUTPUT value can;t be negative as 
                                      # these are the frame counts so it cant be negative

        self.activation = nn.GELU()

    def forward(self, x):
        # shape x: [batch, seq_len, emb_dim]
        x = x.transpose(1,2)    # shape: [batch, emb_dim, seq_len]
        x = self.conv1(x)       # convol is being acted, as the filter have to be moved 
                                # along the last dimension. 
        x = x.transpose(1,2)    # shape: [batch, seq_len, emb_dim]
                                # DONE to apply Norm on it ans Norm is being across the 
                                # last DIMENSION, hence it is being acted upon emb_dim
        x = self.norm1(x)       
        x = self.activation(x)  # ACtivation happened After NORM, as on well-scaled values 

        # shape: [batch, seq_len, emb_dim]
        x = x.transpose(1,2)    # shape: [batch, emb_dim, seq_len]
        x = self.conv2(x)
        x = x.transpose(1,2)    # shape: [batch, seq_len, emb_dim]
        x = self.norm2(x)
        x = self.activation(x)

        # shape: [batch, seq_len, emb_dim]
        x = self.linear(x)    #shape: [batch, seq_len, 1] 
        x = self.output_activation(x)
        x = x.squeeze(-1)     #shape: [batch, seq_len] 
        # as this does RETURNS the ove value per-phoneme NOT one value
        # so thats why the shape is : [batch, seq_len]
        return x 

class LengthRegulator(nn.Module):
    def __init__(self):
        super(LengthRegulator, self).__init__()

    def forward(self, encoder_output, durations):
        # shape encoder_output: [batch, seq_len, emb_dim]
        # shape durations: [batch, seq_len]

        batch_size = encoder_output.shape[0]
        outputs = []

        for i in range(batch_size):

            expanded = torch.repeat_interleave(
                input=encoder_output[i],
                repeats=durations[i].long(),  # as NOW the INT val is being needed 
                                              # instead of the real values which the 
                                              # DuraionPredictor was giving
                dim=0
            )

            outputs.append(expanded)

        # as now the phonemes are being expanded so inside the batch their will 
        # be the uneven sequences of length NOW. so we need to do the padding 
        # of them for batching. 

        output = torch.nn.utils.rnn.pad_sequence(
            sequences=outputs, 
            batch_first = True,   # to tell that the INPUT TENSOR does have the .shape[0]
                                  # as batch
            padding_value = 0.0                                  
        )

        # output shape: [batch, T_max, emb_dim]
        return output