import re
import logging
import torch
import torchaudio
import random
import speechbrain as sb
import torch as nn
from speechbrain.utils.fetching import fetch
from speechbrain.inference.interfaces import Pretrained
from speechbrain.inference.text import GraphemeToPhoneme


logger = logging.getLogger(__name__)


class TTSModel(Pretrained):
    """
    A ready-to-use wrapper for Transformer TTS (text -> mel_spec).
    Arguments
    ---------
    hparams
        Hyperparameters (from HyperPyYAML)"""

    HPARAMS_NEEDED = ["model", "blank_index", "padding_mask", "lookahead_mask", "mel_spec_feats", "label_encoder"]
    MODULES_NEEDED = ["modules"]

    def __init__(self, *args, **kwargs):
        super().__init__(*args, **kwargs)
        self.label_encoder = self.hparams.label_encoder
        self.label_encoder.update_from_iterable(self.hparams["lexicon"], sequence_input=False)
        self.g2p = GraphemeToPhoneme.from_hparams("speechbrain/soundchoice-g2p")


    def text_to_phoneme(self, text):
        """
        Generates phoneme sequences for the given text using a Grapheme-to-Phoneme (G2P) model.

        Args:
            text (str): The input text.

        Returns:
            list: List of phoneme sequences for the words in the text.
        """
        abbreviation_expansions = {
            "Mr.": "Mister",
            "Mrs.": "Misess",
            "Dr.": "Doctor",
            "No.": "Number",
            "St.": "Saint",
            "Co.": "Company",
            "Jr.": "Junior",
            "Maj.": "Major",
            "Gen.": "General",
            "Drs.": "Doctors",
            "Rev.": "Reverend",
            "Lt.": "Lieutenant",
            "Hon.": "Honorable",
            "Sgt.": "Sergeant",
            "Capt.": "Captain",
            "Esq.": "Esquire",
            "Ltd.": "Limited",
            "Col.": "Colonel",
            "Ft.": "Fort"
        }

        # Expand abbreviations
        for abbreviation, expansion in abbreviation_expansions.items():
            text = text.replace(abbreviation, expansion)

        phonemes = self.g2p(text)
        phonemes = self.label_encoder.encode_sequence(phonemes)
        phoneme_seq = torch.LongTensor(phonemes)

        return phoneme_seq, len(phoneme_seq)

    def encode_batch(self, texts):
        """Computes mel-spectrogram for a list of texts

        Texts must be sorted in decreasing order on their lengths

        Arguments
        ---------
        texts: List[str]
            texts to be encoded into spectrogram

        Returns
        -------
        tensors of output spectrograms, output lengths and alignments
        """
        with torch.no_grad():
            phoneme_seqs = [self.text_to_phoneme(text)[0] for text in texts]
            phoneme_seqs_padded, input_lengths = self.pad_sequences(phoneme_seqs)

            encoded_phoneme = self.mods.encoder_emb(phoneme_seqs_padded)
            encoder_emb = self.mods.enc_pre_net(encoded_phoneme)
            pos_emb_enc = self.mods.pos_emb_enc(encoder_emb)
            encoder_emb = encoder_emb + pos_emb_enc


            stop_generated = False
            decoder_input = torch.zeros(1, 80, 1, device=self.device)
            stop_tokens_logits = []
            max_generation_length = 1000
            sequence_length = 0

            result = []
            result.append(decoder_input)

            src_mask = torch.zeros(encoder_emb.size(1), encoder_emb.size(1), device=self.device)
            src_key_padding_mask = self.hparams.padding_mask(encoder_emb, self.hparams.blank_index)


            while not stop_generated and sequence_length < max_generation_length:
                encoded_mel = self.mods.dec_pre_net(decoder_input)
                pos_emb_dec = self.mods.pos_emb_dec(encoded_mel)
                decoder_emb = encoded_mel + pos_emb_dec

                decoder_output = self.mods.Seq2SeqTransformer(
                    encoder_emb, decoder_emb, src_mask=src_mask,
                    src_key_padding_mask=src_key_padding_mask)

                mel_output = self.mods.mel_lin(decoder_output)

                stop_token_logit = self.mods.stop_lin(decoder_output).squeeze(-1)

                post_mel_outputs = self.mods.postnet(mel_output.to(self.device))
                refined_mel_output = mel_output + post_mel_outputs.to(self.device)
                refined_mel_output = refined_mel_output.transpose(1, 2)

                stop_tokens_logits.append(stop_token_logit)
                stop_token_probs = torch.sigmoid(stop_token_logit)

                if torch.any(stop_token_probs[:, -1] >= self.hparams.stop_threshold):
                    stop_generated = True

                decoder_input = refined_mel_output
                result.append(decoder_input)
                sequence_length += 1

            results = torch.cat(result, dim=2)
            stop_tokens_logits = torch.cat(stop_tokens_logits, dim=1)

        return results

    def pad_sequences(self, sequences):
      """Pad sequences to the maximum length sequence in the batch.

      Arguments
      ---------
      sequences: List[torch.Tensor]
          The sequences to pad

      Returns
      -------
      Padded sequences and original lengths
      """
      max_length = max([len(seq) for seq in sequences])
      padded_seqs = torch.zeros(len(sequences), max_length, dtype=torch.long)
      lengths = []
      for i, seq in enumerate(sequences):
          length = len(seq)
          padded_seqs[i, :length] = seq
          lengths.append(length)
      return padded_seqs, torch.tensor(lengths)

    def encode_text(self, text):
        """Runs inference for a single text str"""
        return self.encode_batch(text)

    def forward(self, texts):
        "Encodes the input texts."
        return self.encode_batch(texts)