import librosa
import numpy as np
import torch
import torchaudio
from cached_path import cached_path
import random
import nltk
from models import build_model
from text_utils import TextCleaner
from nltk.tokenize import word_tokenize
import phonemizer
from Modules.diffusion.sampler import DiffusionSampler, ADPM2Sampler, KarrasSchedule
from utils import recursive_munch
from Utils.PLBERT.util import load_plbert

nltk.download("punkt")
np.random.seed(0)
random.seed(0)
torch.manual_seed(0)
torch.backends.cudnn.benchmark = False
torch.backends.cudnn.deterministic = True

global_phonemizer = phonemizer.backend.EspeakBackend(
    language="en-us", preserve_punctuation=True, with_stress=True
)


textcleaner = TextCleaner()


to_mel = torchaudio.transforms.MelSpectrogram(
    n_mels=80, n_fft=2048, win_length=1200, hop_length=300
)
mean, std = -4, 4


def length_to_mask(lengths):
    mask = (
        torch.arange(lengths.max())
        .unsqueeze(0)
        .expand(lengths.shape[0], -1)
        .type_as(lengths)
    )
    mask = torch.gt(mask + 1, lengths.unsqueeze(1))
    return mask


def preprocess(wave):
    wave_tensor = torch.from_numpy(wave).float()
    mel_tensor = to_mel(wave_tensor)
    mel_tensor = (torch.log(1e-5 + mel_tensor.unsqueeze(0)) - mean) / std
    return mel_tensor


def compute_style(path):
    wave, sr = librosa.load(path, sr=24000)
    audio, index = librosa.effects.trim(wave, top_db=30)
    if sr != 24000:
        audio = librosa.resample(audio, sr, 24000)
    mel_tensor = preprocess(audio).to(device)

    with torch.no_grad():
        ref_s = model.style_encoder(mel_tensor.unsqueeze(1))
        ref_p = model.predictor_encoder(mel_tensor.unsqueeze(1))

    return torch.cat([ref_s, ref_p], dim=1)


device = "cpu"
if torch.cuda.is_available():
    device = "cuda"
elif torch.backends.mps.is_available():
    print("MPS would be available but cannot be used rn")
    # device = "mps"

# config = yaml.safe_load(open("Models/LibriTTS/config.yml"))
config = {
    "ASR_config": "Utils/ASR/config.yml",
    "ASR_path": "Utils/ASR/epoch_00080.pth",
    "F0_path": "Utils/JDC/bst.t7",
    "PLBERT_dir": "Utils/PLBERT/",
    "batch_size": 8,
    "data_params": {
        "OOD_data": "Data/OOD_texts.txt",
        "min_length": 50,
        "root_path": "",
        "train_data": "Data/train_list.txt",
        "val_data": "Data/val_list.txt",
    },
    "device": "cuda",
    "epochs_1st": 40,
    "epochs_2nd": 25,
    "first_stage_path": "first_stage.pth",
    "load_only_params": False,
    "log_dir": "Models/LibriTTS",
    "log_interval": 10,
    "loss_params": {
        "TMA_epoch": 4,
        "diff_epoch": 0,
        "joint_epoch": 0,
        "lambda_F0": 1.0,
        "lambda_ce": 20.0,
        "lambda_diff": 1.0,
        "lambda_dur": 1.0,
        "lambda_gen": 1.0,
        "lambda_mel": 5.0,
        "lambda_mono": 1.0,
        "lambda_norm": 1.0,
        "lambda_s2s": 1.0,
        "lambda_slm": 1.0,
        "lambda_sty": 1.0,
    },
    "max_len": 300,
    "model_params": {
        "decoder": {
            "resblock_dilation_sizes": [[1, 3, 5], [1, 3, 5], [1, 3, 5]],
            "resblock_kernel_sizes": [3, 7, 11],
            "type": "hifigan",
            "upsample_initial_channel": 512,
            "upsample_kernel_sizes": [20, 10, 6, 4],
            "upsample_rates": [10, 5, 3, 2],
        },
        "diffusion": {
            "dist": {
                "estimate_sigma_data": True,
                "mean": -3.0,
                "sigma_data": 0.19926648961191362,
                "std": 1.0,
            },
            "embedding_mask_proba": 0.1,
            "transformer": {
                "head_features": 64,
                "multiplier": 2,
                "num_heads": 8,
                "num_layers": 3,
            },
        },
        "dim_in": 64,
        "dropout": 0,
        "hidden_dim": 512,
        "max_conv_dim": 512,
        "max_dur": 50,
        "multispeaker": True,
        "n_layer": 3,
        "n_mels": 80,
        "n_token": 178,
        "slm": {
            "hidden": 768,
            "initial_channel": 64,
            "model": "microsoft/wavlm-base-plus",
            "nlayers": 13,
            "sr": 16000,
        },
        "style_dim": 128,
    },
    "optimizer_params": {"bert_lr": 1e-05, "ft_lr": 1e-05, "lr": 0.0001},
    "preprocess_params": {
        "spect_params": {"hop_length": 300, "n_fft": 2048, "win_length": 1200},
        "sr": 24000,
    },
    "pretrained_model": "Models/LibriTTS/epoch_2nd_00002.pth",
    "save_freq": 1,
    "second_stage_load_pretrained": True,
    "slmadv_params": {
        "batch_percentage": 0.5,
        "iter": 20,
        "max_len": 500,
        "min_len": 400,
        "scale": 0.01,
        "sig": 1.5,
        "thresh": 5,
    },
}


BERT_path = config.get("PLBERT_dir", False)
plbert = load_plbert(BERT_path)


model_params = recursive_munch(config["model_params"])
model = build_model(model_params, plbert)
_ = [model[key].eval() for key in model]
_ = [model[key].to(device) for key in model]

# for key in model:
#     print(f"Compiling {key}")
#     model[key] = torch.compile(model[key])
#     print(f"Compiled {key}")


params_whole = torch.load(
    str(cached_path("https://base-weights.weights.gg/epochs_2nd_00020.pth")),
    map_location="cpu",
)
params = params_whole["net"]

for key in model:
    if key in params:
        print("%s loaded" % key)
        try:
            model[key].load_state_dict(params[key])
        except:
            from collections import OrderedDict

            state_dict = params[key]
            new_state_dict = OrderedDict()
            for k, v in state_dict.items():
                name = k[7:]  # remove `module.`
                new_state_dict[name] = v
            # load params
            model[key].load_state_dict(new_state_dict, strict=False)
#             except:
#                 _load(params[key], model[key])
_ = [model[key].eval() for key in model]


sampler = DiffusionSampler(
    model.diffusion.diffusion,
    sampler=ADPM2Sampler(),
    sigma_schedule=KarrasSchedule(
        sigma_min=0.0001, sigma_max=3.0, rho=9.0
    ),  # empirical parameters
    clamp=False,
)


def inference(
    text,
    ref_s,
    alpha=0.3,
    beta=0.7,
    diffusion_steps=5,
    embedding_scale=1,
    use_gruut=False,
):
    text = text.strip()
    ps = global_phonemizer.phonemize([text])
    ps = word_tokenize(ps[0])
    ps = " ".join(ps)
    tokens = textcleaner(ps)
    tokens.insert(0, 0)
    tokens = torch.LongTensor(tokens).to(device).unsqueeze(0)

    with torch.no_grad():
        input_lengths = torch.LongTensor([tokens.shape[-1]]).to(device)
        text_mask = length_to_mask(input_lengths).to(device)

        t_en = model.text_encoder(tokens, input_lengths, text_mask)
        bert_dur = model.bert(tokens, attention_mask=(~text_mask).int())
        d_en = model.bert_encoder(bert_dur).transpose(-1, -2)

        s_pred = sampler(
            noise=torch.randn((1, 256)).unsqueeze(1).to(device),
            embedding=bert_dur,
            embedding_scale=embedding_scale,
            features=ref_s,  # reference from the same speaker as the embedding
            num_steps=diffusion_steps,
        ).squeeze(1)

        s = s_pred[:, 128:]
        ref = s_pred[:, :128]

        ref = alpha * ref + (1 - alpha) * ref_s[:, :128]
        s = beta * s + (1 - beta) * ref_s[:, 128:]

        d = model.predictor.text_encoder(d_en, s, input_lengths, text_mask)

        x, _ = model.predictor.lstm(d)
        duration = model.predictor.duration_proj(x)

        duration = torch.sigmoid(duration).sum(axis=-1)
        pred_dur = torch.round(duration.squeeze()).clamp(min=1)

        pred_aln_trg = torch.zeros(input_lengths, int(pred_dur.sum().data))
        c_frame = 0
        for i in range(pred_aln_trg.size(0)):
            pred_aln_trg[i, c_frame : c_frame + int(pred_dur[i].data)] = 1
            c_frame += int(pred_dur[i].data)

        # encode prosody
        en = d.transpose(-1, -2) @ pred_aln_trg.unsqueeze(0).to(device)
        asr_new = torch.zeros_like(en)
        asr_new[:, :, 0] = en[:, :, 0]
        asr_new[:, :, 1:] = en[:, :, 0:-1]
        en = asr_new

        F0_pred, N_pred = model.predictor.F0Ntrain(en, s)

        asr = t_en @ pred_aln_trg.unsqueeze(0).to(device)
        asr_new = torch.zeros_like(asr)
        asr_new[:, :, 0] = asr[:, :, 0]
        asr_new[:, :, 1:] = asr[:, :, 0:-1]
        asr = asr_new

        out = model.decoder(asr, F0_pred, N_pred, ref.squeeze().unsqueeze(0))

    return (
        out.squeeze().cpu().numpy()[..., :-50]
    )  # weird pulse at the end of the model, need to be fixed later