import os
import sys
import time
import torch
import configparser
import urllib.request

from .tacotron2.model import Tacotron2
from .waveglow.model import WaveGlow
from .text import text_to_sequence
from .waveglow.denoiser import Denoiser

from speechbrain.pretrained import Tacotron2 as SpeechBrain
from speechbrain.pretrained import HIFIGAN
from scipy.io.wavfile import write

config = configparser.ConfigParser()
config.read(os.path.join(os.path.dirname(__file__), "config.cfg"))

force_reload = bool(int(config["HPARAMS"]["FORCE_RELOAD"]))
rate = int(config["HPARAMS"]["RATE"])
denoicer_strength = float(config["HPARAMS"]["DENOICER_STRENGTH"])


def _download_checkpoint(checkpoint, force_reload):
    """
    This method download the tacotron checkpoint weights from the checkpoint url

    Parameters:
    -----------
    checkpoint: String, url to the checkpoint.
    force_reload: Boolean.

    Return:
    -----------
    ckpt_file: String, Path to downloaded checkpoint file

    """
    model_dir = os.path.join(torch.hub._get_torch_home(), "checkpoints")
    if not os.path.exists(model_dir):
        os.makedirs(model_dir)
    ckpt_file = os.path.join(model_dir, os.path.basename(checkpoint))
    if not os.path.exists(ckpt_file) or force_reload:
        sys.stderr.write("Downloading checkpoint from {}\n".format(checkpoint))
        urllib.request.urlretrieve(checkpoint, ckpt_file)
    return ckpt_file


def _checkpoint_from_distributed(state_dict):
    """
    Checks whether checkpoint was generated by DistributedDataParallel. DDP
    wraps model in additional "module.", it needs to be unwrapped for single
    GPU inference.

    Parameters:
    -----------
    state_dict: Dictionary

    Return:
    -----------
    ret: Boolean

    """
    ret = False
    for key, _ in state_dict.items():
        if key.find("module.") != -1:
            ret = True
            break
    return ret


def _unwrap_distributed(state_dict):
    """
    Unwraps model from DistributedDataParallel.
    DDP wraps model in additional "module.", it needs to be removed for single
    GPU inference.

    Parameters:
    -----------

    state_dict: Dictionary, model's state dict

    Return:
    -----------

    new_state_dict = Dictionary, modified model's state dict
    """
    new_state_dict = {}
    for key, value in state_dict.items():
        new_key = key.replace("module.1.", "")
        new_key = new_key.replace("module.", "")
        new_state_dict[new_key] = value
    return new_state_dict


def _load_tacotron(checkpoint, force_reload):
    """
    Downloads tacotron2 checkpoints and loading state dictionery

    Parameters:
    -----------
    checkpoint: String, url to the tacotron2 checkpoint.
    force_reload: Boolean. setting this value to true will ignore the downloaded checkpoints from the cache.
    Return:
    -----------
    tacotron2: Model Tacotron
    """
    ckpt_file = _download_checkpoint(checkpoint, force_reload)
    ckpt = torch.load(ckpt_file, map_location=torch.device("cpu"))
    state_dict = ckpt["state_dict"]
    if _checkpoint_from_distributed(state_dict):
        state_dict = _unwrap_distributed(state_dict)
    config = ckpt["config"]
    tacotron2 = Tacotron2(**config)
    tacotron2.load_state_dict(state_dict)

    return tacotron2


def _load_waveglow(checkpoint, force_reload):
    """
    Downloads waveglow checkpoints and loading state dictionery

    Parameters:
    -----------
    checkpoint: String, url to the waveglow checkpoint.
    force_reload: Boolean. setting this value to true will ignore the downloaded checkpoints from the cache.
    Return:
    -----------
    waveglow: Model, waveglow
    """

    ckpt_file = _download_checkpoint(checkpoint, force_reload)
    ckpt = torch.load(ckpt_file, map_location=torch.device("cpu"))
    state_dict = ckpt["state_dict"]
    if _checkpoint_from_distributed(state_dict):
        state_dict = _unwrap_distributed(state_dict)
    config = ckpt["config"]
    waveglow = WaveGlow(**config)
    waveglow.load_state_dict(state_dict)

    return waveglow


def _pretrained_models_file_path(model_name):
    """
    Creating a space in cache to save the downloaded model checkpoint.
    Parameters:
    -----------
    model_name: String, name of the pretrained model.

    Return:
    -----------
    model_ckpt_file: Str, name of the checkpoint directory.
    """
    model_dir = os.path.join(torch.hub._get_torch_home(), "pretrained_models")
    if not os.path.exists(model_dir):
        os.makedirs(model_dir)
    model_ckpt_file = os.path.join(model_dir, os.path.basename(model_name))
    return model_ckpt_file


def load_seq2seq_model(model_name, force_reload):
    """
    Loading the seq2seq model.
    """
    if model_name == "tacotron2_v1":
        return _load_tacotron(config["SEQ2SEQ_MODEL"]["TACOTRON2_V1"], force_reload)
    elif model_name == "tacotron2_v2":
        return SpeechBrain.from_hparams(
            source=config["SEQ2SEQ_MODEL"]["TACOTRON2_v2"],
            savedir=_pretrained_models_file_path("tacotron2_v1"),
        )
    else:
        return "Incorrect model name"


def load_vocorder_model(model_name, force_reload):
    """
    Loading the vocorder model.
    """
    if model_name == "waveglow":
        waveglow = _load_waveglow(config["VOCORDER"]["WAVEGLOW"], force_reload)
        return waveglow.remove_weightnorm(waveglow)
    elif model_name == "hifigan":
        return HIFIGAN.from_hparams(
            source="speechbrain/tts-hifigan-libritts-22050Hz",
            savedir=_pretrained_models_file_path("hifigan"),
        )
    else:
        return "The vocorder model is incorrect"


def _pad_sequences(batch):
    """
    Right zero-pad all one-hot text sequences to max input length.
    """
    input_lengths, ids_sorted_decreasing = torch.sort(
        torch.LongTensor([len(x) for x in batch]), dim=0, descending=True
    )
    max_input_len = input_lengths[0]

    text_padded = torch.LongTensor(len(batch), max_input_len)
    text_padded.zero_()
    for i in range(len(ids_sorted_decreasing)):
        text = batch[ids_sorted_decreasing[i]]
        text_padded[i, : text.size(0)] = text

    return text_padded, input_lengths


def prepare_input_sequence(texts, cpu_run=True):
    """
    Preprocess input text sequence with text cleaners.
    """
    d = []
    for i, text in enumerate(texts):
        d.append(torch.IntTensor(text_to_sequence(text, ["english_cleaners"])[:]))

    text_padded, input_lengths = _pad_sequences(d)
    if not cpu_run:
        text_padded = text_padded.cuda().long()
        input_lengths = input_lengths.cuda().long()
    else:
        text_padded = text_padded.long()
        input_lengths = input_lengths.long()

    return text_padded, input_lengths


def save_file(input_numpy, path="", rate=rate):
    """
    Saving the audio file the given path

    Parameters:
    -----------
    input_numpy: Numpy array, 1d numpy with floats contains audio data.
    path: String, path to the directory to save the file.
    rate: int, The rate of the audio.

    Return:
    -----------
    success: Boolean, True if successfuly saved.
    file_path: String, path to the saved file.
    """
    success = False
    ret_dict = {
        "save_status": success,
    }

    if os.path.exists(path):
        try:
            file_name = "audio_file_" + str(time.time()) + ".wav"
            file_path = os.path.join(path, file_name)
            write(file_path, rate, input_numpy)
            success = True
            ret_dict = {
                "save_status": success,
                "file_path": file_path,
            }
        except Exception as ex:
            print(ex)
            success = False
            file_path = None
    else:
        print("Set up directory path properly to save the audio file.")

    return ret_dict


def prediction(input_text, seq2seqmodel, vocorder):
    """
    Inferencing

    Parameters:
    -----------
    input_text: String, input text for preprocessing.

    Return:
    -----------
    audio_numpy: Numpy array, 1d numpy with floats contains audio data.
    """
    print(
        "Synthesizing speeches with "
        + seq2seqmodel.__class__.__name__
        + " and "
        + vocorder.__class__.__name__
        + "."
    )

    sequences, lengths = prepare_input_sequence([input_text], cpu_run=True)

    with torch.no_grad():
        mel, _, _ = seq2seqmodel.infer(sequences, lengths)

        if (
            vocorder.__class__.__name__ == "WaveGlow"
            or vocorder.__class__.__name__ == "WAVEGLOW"
        ):
            audio = vocorder.infer(mel)
            denoiser = Denoiser(vocorder)
            audio = denoiser(audio, strength=denoicer_strength).squeeze(1)
            audio_numpy = audio[0].data.cpu().numpy()

        elif vocorder.__class__.__name__ == "HIFIGAN":
            audio = vocorder.decode_batch(mel)
            audio_numpy = audio[0].data.cpu().numpy()[0]

        else:
            print("No valid vocorder")
    return audio_numpy
