import spaces
import accelerate
import gradio as gr
import torch
import safetensors
from huggingface_hub import hf_hub_download
import soundfile as sf
import os

import numpy as np
import librosa
from models.codec.kmeans.repcodec_model import RepCodec
from models.tts.maskgct.maskgct_s2a import MaskGCT_S2A
from models.tts.maskgct.maskgct_t2s import MaskGCT_T2S
from models.codec.amphion_codec.codec import CodecEncoder, CodecDecoder
from transformers import Wav2Vec2BertModel
from utils.util import load_config
from models.tts.maskgct.g2p.g2p_generation import g2p, chn_eng_g2p

from transformers import SeamlessM4TFeatureExtractor

import whisper
import langid

processor = SeamlessM4TFeatureExtractor.from_pretrained("facebook/w2v-bert-2.0")

device = torch.device("cuda" if torch.cuda.is_available() else "cpu")

whisper_model = whisper.load_model("turbo")

def detect_speech_language(speech_file):
    # load audio and pad/trim it to fit 30 seconds
    audio = whisper.load_audio(speech_file)
    audio = whisper.pad_or_trim(audio)

    # make log-Mel spectrogram and move to the same device as the model
    mel = whisper.log_mel_spectrogram(audio, n_mels=128).to(whisper_model.device)

    # detect the spoken language
    _, probs = whisper_model.detect_language(mel)
    return max(probs, key=probs.get)


def detect_text_language(text):
    return langid.classify(text)[0]

@torch.no_grad()
def get_prompt_text(speech_16k, language):
    full_prompt_text = ""
    shot_prompt_text = ""
    short_prompt_end_ts = 0.0

    asr_result = whisper_model.transcribe(speech_16k, language=language)
    full_prompt_text = asr_result["text"] # whisper asr result
    #text = asr_result["segments"][0]["text"] # whisperx asr result
    shot_prompt_text = ""
    short_prompt_end_ts = 0.0
    for segment in asr_result["segments"]:
        shot_prompt_text = shot_prompt_text +  segment['text']
        short_prompt_end_ts = segment['end']
        if short_prompt_end_ts >= 4:
            break
    return full_prompt_text, shot_prompt_text, short_prompt_end_ts


def g2p_(text, language):
    if language in ["zh", "en"]:
        return chn_eng_g2p(text)
    else:
        return g2p(text, sentence=None, language=language)


def build_t2s_model(cfg, device):
    t2s_model = MaskGCT_T2S(cfg=cfg)
    t2s_model.eval()
    t2s_model.to(device)
    return t2s_model


def build_s2a_model(cfg, device):
    soundstorm_model = MaskGCT_S2A(cfg=cfg)
    soundstorm_model.eval()
    soundstorm_model.to(device)
    return soundstorm_model


def build_semantic_model(device):
    semantic_model = Wav2Vec2BertModel.from_pretrained("facebook/w2v-bert-2.0")
    semantic_model.eval()
    semantic_model.to(device)
    stat_mean_var = torch.load("./models/tts/maskgct/ckpt/wav2vec2bert_stats.pt")
    semantic_mean = stat_mean_var["mean"]
    semantic_std = torch.sqrt(stat_mean_var["var"])
    semantic_mean = semantic_mean.to(device)
    semantic_std = semantic_std.to(device)
    return semantic_model, semantic_mean, semantic_std


def build_semantic_codec(cfg, device):
    semantic_codec = RepCodec(cfg=cfg)
    semantic_codec.eval()
    semantic_codec.to(device)
    return semantic_codec


def build_acoustic_codec(cfg, device):
    codec_encoder = CodecEncoder(cfg=cfg.encoder)
    codec_decoder = CodecDecoder(cfg=cfg.decoder)
    codec_encoder.eval()
    codec_decoder.eval()
    codec_encoder.to(device)
    codec_decoder.to(device)
    return codec_encoder, codec_decoder


@torch.no_grad()
def extract_features(speech, processor):
    inputs = processor(speech, sampling_rate=16000, return_tensors="pt")
    input_features = inputs["input_features"][0]
    attention_mask = inputs["attention_mask"][0]
    return input_features, attention_mask


@torch.no_grad()
def extract_semantic_code(semantic_mean, semantic_std, input_features, attention_mask):
    vq_emb = semantic_model(
        input_features=input_features,
        attention_mask=attention_mask,
        output_hidden_states=True,
    )
    feat = vq_emb.hidden_states[17]  # (B, T, C)
    feat = (feat - semantic_mean.to(feat)) / semantic_std.to(feat)

    semantic_code, rec_feat = semantic_codec.quantize(feat)  # (B, T)
    return semantic_code, rec_feat


@torch.no_grad()
def extract_acoustic_code(speech):
    vq_emb = codec_encoder(speech.unsqueeze(1))
    _, vq, _, _, _ = codec_decoder.quantizer(vq_emb)
    acoustic_code = vq.permute(1, 2, 0)
    return acoustic_code


@torch.no_grad()
def text2semantic(
    device,
    prompt_speech,
    prompt_text,
    prompt_language,
    target_text,
    target_language,
    target_len=None,
    n_timesteps=50,
    cfg=2.5,
    rescale_cfg=0.75,
):

    prompt_phone_id = g2p_(prompt_text, prompt_language)[1]

    target_phone_id = g2p_(target_text, target_language)[1]

    if target_len < 0:
        target_len = int(
            (len(prompt_speech) * len(target_phone_id) / len(prompt_phone_id))
            / 16000
            * 50
        )
    else:
        target_len = int(target_len * 50)

    prompt_phone_id = torch.tensor(prompt_phone_id, dtype=torch.long).to(device)
    target_phone_id = torch.tensor(target_phone_id, dtype=torch.long).to(device)

    phone_id = torch.cat([prompt_phone_id, target_phone_id])

    input_fetures, attention_mask = extract_features(prompt_speech, processor)
    input_fetures = input_fetures.unsqueeze(0).to(device)
    attention_mask = attention_mask.unsqueeze(0).to(device)
    semantic_code, _ = extract_semantic_code(
        semantic_mean, semantic_std, input_fetures, attention_mask
    )

    predict_semantic = t2s_model.reverse_diffusion(
        semantic_code[:, :],
        target_len,
        phone_id.unsqueeze(0),
        n_timesteps=n_timesteps,
        cfg=cfg,
        rescale_cfg=rescale_cfg,
    )

    combine_semantic_code = torch.cat([semantic_code[:, :], predict_semantic], dim=-1)
    prompt_semantic_code = semantic_code

    return combine_semantic_code, prompt_semantic_code


@torch.no_grad()
def semantic2acoustic(
    device,
    combine_semantic_code,
    acoustic_code,
    n_timesteps=[25, 10, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1],
    cfg=2.5,
    rescale_cfg=0.75,
):

    semantic_code = combine_semantic_code

    cond = s2a_model_1layer.cond_emb(semantic_code)
    prompt = acoustic_code[:, :, :]
    predict_1layer = s2a_model_1layer.reverse_diffusion(
        cond=cond,
        prompt=prompt,
        temp=1.5,
        filter_thres=0.98,
        n_timesteps=n_timesteps[:1],
        cfg=cfg,
        rescale_cfg=rescale_cfg,
    )

    cond = s2a_model_full.cond_emb(semantic_code)
    prompt = acoustic_code[:, :, :]
    predict_full = s2a_model_full.reverse_diffusion(
        cond=cond,
        prompt=prompt,
        temp=1.5,
        filter_thres=0.98,
        n_timesteps=n_timesteps,
        cfg=cfg,
        rescale_cfg=rescale_cfg,
        gt_code=predict_1layer,
    )

    vq_emb = codec_decoder.vq2emb(predict_full.permute(2, 0, 1), n_quantizers=12)
    recovered_audio = codec_decoder(vq_emb)
    prompt_vq_emb = codec_decoder.vq2emb(prompt.permute(2, 0, 1), n_quantizers=12)
    recovered_prompt_audio = codec_decoder(prompt_vq_emb)
    recovered_prompt_audio = recovered_prompt_audio[0][0].cpu().numpy()
    recovered_audio = recovered_audio[0][0].cpu().numpy()
    combine_audio = np.concatenate([recovered_prompt_audio, recovered_audio])

    return combine_audio, recovered_audio


# Load the model and checkpoints
def load_models():
    cfg_path = "./models/tts/maskgct/config/maskgct.json"

    cfg = load_config(cfg_path)
    semantic_model, semantic_mean, semantic_std = build_semantic_model(device)
    semantic_codec = build_semantic_codec(cfg.model.semantic_codec, device)
    codec_encoder, codec_decoder = build_acoustic_codec(
        cfg.model.acoustic_codec, device
    )
    t2s_model = build_t2s_model(cfg.model.t2s_model, device)
    s2a_model_1layer = build_s2a_model(cfg.model.s2a_model.s2a_1layer, device)
    s2a_model_full = build_s2a_model(cfg.model.s2a_model.s2a_full, device)

    # Download checkpoints
    semantic_code_ckpt = hf_hub_download(
        "amphion/MaskGCT", filename="semantic_codec/model.safetensors"
    )
    # codec_encoder_ckpt = hf_hub_download(
    #     "amphion/MaskGCT", filename="acoustic_codec/model.safetensors"
    # )
    # codec_decoder_ckpt = hf_hub_download(
    #     "amphion/MaskGCT", filename="acoustic_codec/model_1.safetensors"
    # )
    t2s_model_ckpt = hf_hub_download(
        "amphion/MaskGCT", filename="t2s_model/model.safetensors"
    )
    s2a_1layer_ckpt = hf_hub_download(
        "amphion/MaskGCT", filename="s2a_model/s2a_model_1layer/model.safetensors"
    )
    s2a_full_ckpt = hf_hub_download(
        "amphion/MaskGCT", filename="s2a_model/s2a_model_full/model.safetensors"
    )

    safetensors.torch.load_model(semantic_codec, semantic_code_ckpt)
    # safetensors.torch.load_model(codec_encoder, codec_encoder_ckpt)
    # safetensors.torch.load_model(codec_decoder, codec_decoder_ckpt)
    accelerate.load_checkpoint_and_dispatch(codec_encoder, "./acoustic_codec/model.safetensors")
    accelerate.load_checkpoint_and_dispatch(codec_decoder, "./acoustic_codec/model_1.safetensors")
    safetensors.torch.load_model(t2s_model, t2s_model_ckpt)
    safetensors.torch.load_model(s2a_model_1layer, s2a_1layer_ckpt)
    safetensors.torch.load_model(s2a_model_full, s2a_full_ckpt)

    return (
        semantic_model,
        semantic_mean,
        semantic_std,
        semantic_codec,
        codec_encoder,
        codec_decoder,
        t2s_model,
        s2a_model_1layer,
        s2a_model_full,
    )


@torch.no_grad()
def maskgct_inference(
    prompt_speech_path,
    target_text,
    target_len=None,
    n_timesteps=25,
    cfg=2.5,
    rescale_cfg=0.75,
    n_timesteps_s2a=[25, 10, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1],
    cfg_s2a=2.5,
    rescale_cfg_s2a=0.75,
    device=torch.device("cuda:0"),
):
    speech_16k = librosa.load(prompt_speech_path, sr=16000)[0]
    speech = librosa.load(prompt_speech_path, sr=24000)[0]

    prompt_language = detect_speech_language(prompt_speech_path)
    full_prompt_text, short_prompt_text, shot_prompt_end_ts = get_prompt_text(prompt_speech_path,
                                                                              prompt_language)
    # use the first 4+ seconds wav as the prompt in case the prompt wav is too long
    speech = speech[0: int(shot_prompt_end_ts * 24000)]
    speech_16k = speech_16k[0: int(shot_prompt_end_ts*16000)]
    target_language = detect_text_language(target_text)
    combine_semantic_code, _ = text2semantic(
        device,
        speech_16k,
        short_prompt_text,
        prompt_language,
        target_text,
        target_language,
        target_len,
        n_timesteps,
        cfg,
        rescale_cfg,
    )
    acoustic_code = extract_acoustic_code(torch.tensor(speech).unsqueeze(0).to(device))
    _, recovered_audio = semantic2acoustic(
        device,
        combine_semantic_code,
        acoustic_code,
        n_timesteps=n_timesteps_s2a,
        cfg=cfg_s2a,
        rescale_cfg=rescale_cfg_s2a,
    )

    return recovered_audio


@spaces.GPU
def inference(
    prompt_wav,
    target_text,
    target_len,
    n_timesteps,
):
    save_path = "./output/output.wav"
    os.makedirs("./output", exist_ok=True)
    recovered_audio = maskgct_inference(
        prompt_wav,
        target_text,
        target_len=target_len,
        n_timesteps=int(n_timesteps),
        device=device,
    )
    sf.write(save_path, recovered_audio, 24000)
    return save_path

# Load models once
(
    semantic_model,
    semantic_mean,
    semantic_std,
    semantic_codec,
    codec_encoder,
    codec_decoder,
    t2s_model,
    s2a_model_1layer,
    s2a_model_full,
) = load_models()

# Language list
language_list = ["en", "zh", "ja", "ko", "fr", "de"]

# Gradio interface
iface = gr.Interface(
    fn=inference,
    inputs=[
        gr.Audio(label="Upload Prompt Wav", type="filepath"),
        gr.Textbox(label="Target Text"),
        gr.Number(
            label="Target Duration (in seconds), if the target duration is less than 0, the system will estimate a duration.", value=-1
        ),  # Removed 'optional=True'
        gr.Slider(
            label="Number of Timesteps", minimum=15, maximum=100, value=25, step=1
        ),
    ],
    outputs=gr.Audio(label="Generated Audio"),
    title="MaskGCT TTS Demo",
    description="""
    ## MaskGCT: Zero-Shot Text-to-Speech with Masked Generative Codec Transformer
    
    [![arXiv](https://img.shields.io/badge/arXiv-Paper-COLOR.svg)](https://arxiv.org/abs/2409.00750)
    
    [![hf](https://img.shields.io/badge/%F0%9F%A4%97%20HuggingFace-model-yellow)](https://huggingface.co/amphion/maskgct)
    
    [![hf](https://img.shields.io/badge/%F0%9F%A4%97%20HuggingFace-demo-pink)](https://huggingface.co/spaces/amphion/maskgct)
    
    [![readme](https://img.shields.io/badge/README-Key%20Features-blue)](https://github.com/open-mmlab/Amphion/tree/main/models/tts/maskgct)
    """
)

# Launch the interface
iface.launch(allowed_paths=["./output"])