import argparse
import torch
import random
import numpy as np
import time
import random
import yaml
import numpy as np
import torch.nn.functional as F
import torchaudio
import librosa
import os
import nltk
import re

from nltk.tokenize import word_tokenize
from Utils.JDC import model
from models import *
from utils import *
from text_utils_gal import TextCleanerGal
from munch import Munch
from Utils.ASR.AuxiliaryASR.phonemize import run_cotovia_with_phrase, clean_output
from Utils.PLBERT.util import load_plbert
from Modules.diffusion.sampler import DiffusionSampler, ADPM2Sampler, KarrasSchedule
from speechmos import dnsmos

import torch.nn.functional as F


def split_text_into_sentences(text):
    # Separa por ., : o ? y conserva el delimitador en cada segmento
    sentences = re.findall(r'[^.:?!]+[.:?!]|[^.:?!]+$', text, flags=re.S)
    return [s.strip() for s in sentences if s.strip()]


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):
    to_mel = torchaudio.transforms.MelSpectrogram(
        n_mels=80, n_fft=2048, win_length=1200, hop_length=300)
    mean, std = -4, 4
    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, model, device):
    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)


def LFinference(text, s_prev, ref_s, textcleaner, device, sampler, model, alpha=0.3, beta=0.9, t=0.7, diffusion_steps=5, embedding_scale=1):
    text = text.strip()

    ps = clean_output(run_cotovia_with_phrase(text))

    print(f"Phonemized text: {ps}")

    tokens = textcleaner(ps)
    blank_index = textcleaner([" "], mode="phoneme")[0]
    tokens.insert(0, blank_index)  # add a blank at the beginning (silence)
    tokens = torch.LongTensor(tokens).to(device).unsqueeze(0)
    print(f"Token IDs: {tokens}")

    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, num_steps=diffusion_steps,
                         embedding_scale=embedding_scale, features=ref_s).squeeze(1)

        if s_prev is not None:
            # convex combination of previous and current style
            s_pred = t * s_prev + (1 - t) * s_pred

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

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

        s_pred = torch.cat([ref, s], dim=-1)

        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))

        F0_pred, N_pred = model.predictor.F0Ntrain(en, s)
        asr = (t_en @ pred_aln_trg.unsqueeze(0).to(device))

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

    return out.squeeze().cpu().numpy()[..., :-50], s_pred


def main(args=None):
    torch.manual_seed(0)
    torch.backends.cudnn.benchmark = False
    torch.backends.cudnn.deterministic = True

    random.seed(0)
    np.random.seed(0)

    config = yaml.safe_load(open(args.config))

    if args.device != "cpu":
        os.environ["CUDA_VISIBLE_DEVICES"] = args.device

    device = 'cuda' if torch.cuda.is_available() and args.device != "cpu" else 'cpu'
    textcleaner = TextCleanerGal()
    model_config = yaml.safe_load(open(config['model_config']))

    nltk.download('punkt_tab')

    # load pretrained ASR model
    ASR_config = model_config.get('ASR_config', False)
    ASR_path = model_config.get('ASR_path', False)
    text_aligner = load_ASR_models(ASR_path, ASR_config)

    # load pretrained F0 model
    F0_path = model_config.get('F0_path', False)
    pitch_extractor = load_F0_models(F0_path)

    # load BERT model
    BERT_path = model_config.get('PLBERT_dir', False)
    plbert = load_plbert(BERT_path)

    model = build_model(recursive_munch(
        model_config['model_params']), text_aligner, pitch_extractor, plbert)
    _ = [model[key].eval() for key in model]
    _ = [model[key].to(device) 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
    )

    params_whole = torch.load(config['checkpoint_path'], 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)

    _ = [model[key].eval() for key in model]

    if args.text is None and args.file is None:
        print("Por favor, proporciona o argumento --text ou --file .")
        return

    if args.text is None and args.file is not None:
        with open(args.file, 'r') as f:
            args.text = f.read()

    sentences = split_text_into_sentences(args.text)
    wavs = []
    s_prev = None
    # Obtener parámetros de argumentos o config
    alfa = args.alpha if hasattr(
        args, 'alpha') and args.alpha is not None else config.get('alpha', 0.3)
    beta = args.beta if hasattr(
        args, 'beta') and args.beta is not None else config.get('beta', 0.7)
    t = args.t if hasattr(
        args, 't') and args.t is not None else config.get('t', 0.7)
    diffusion_steps = args.diffusion_steps if hasattr(
        args, 'diffusion_steps') and args.diffusion_steps is not None else config.get('diffusion_steps', 5)
    embedding_scale = args.embedding_scale if hasattr(
        args, 'embedding_scale') and args.embedding_scale is not None else config.get('embedding_scale', 1)
    interrogative_reference = compute_style(
        config['interrogative_reference'], model, device)
    exclamative_reference = compute_style(
        config['exclamative_reference'], model, device)
    normal_reference = compute_style(config['normal_reference'], model, device)

    for text in sentences:
        if text.strip() == "":
            continue

        if "?" in text:
            s_ref = interrogative_reference
        elif "!" in text:
            s_ref = exclamative_reference
        else:
            s_ref = normal_reference

        wav, s_prev = LFinference(
            text, s_prev, s_ref, textcleaner, device, sampler, model,
            diffusion_steps=diffusion_steps, embedding_scale=embedding_scale, alpha=alfa, beta=beta, t=t)

        wavs.append(wav)

    wav = np.concatenate(wavs)
    os.makedirs(args.output_dir, exist_ok=True)
    torchaudio.save(os.path.join(args.output_dir, args.output_file+"_"+str(diffusion_steps)+"_"+str(embedding_scale)+".wav"),
                    torch.from_numpy(wav).unsqueeze(0), config['sample_rate'])

    if args.evaluate:
        wav_resampled = librosa.resample(wav, orig_sr=config['sample_rate'],
                                         target_sr=16000, res_type='kaiser_best', fix=True)
        mos_dict = dnsmos.run(wav_resampled, sr=16000)
        print(f"DNSMOS scores: {mos_dict}")


if __name__ == "__main__":
    parser = argparse.ArgumentParser(
        description="Inference script for Galician Style TTS")
    parser.add_argument("--config", type=str, required=True,
                        help="Path to the inference config file")
    parser.add_argument("--text", type=str,
                        help="Text to synthesize")
    parser.add_argument("--file", type=str,
                        help="Path to a text file containing sentences to synthesize, one per line")
    parser.add_argument("--device", type=str, default="0",
                        help="Device to run the inference on")
    parser.add_argument("--output_dir", type=str, default="outputs",
                        help="Directory to save the output audio")
    parser.add_argument("--output_file", type=str, default="output",
                        help="Name of the output audio file (without extension)")
    parser.add_argument("--evaluate", action="store_true")
    parser.add_argument("--alpha", type=float, default=None,
                        help="Valor de alpha")
    parser.add_argument("--beta", type=float, default=None,
                        help="Valor de beta")
    parser.add_argument("--t", type=float, default=None,
                        help="Valor de t")
    parser.add_argument("--diffusion_steps", type=int, default=None,
                        help="Valor de diffusion_steps")
    parser.add_argument("--embedding_scale", type=float, default=None,
                        help="Valor de embedding_scale")

    args = parser.parse_args()
    main(args)