import torch
import torchaudio
from tqdm import tqdm
from underthesea import sent_tokenize
import os

from TTS.tts.configs.xtts_config import XttsConfig
from TTS.tts.models.xtts import Xtts
from TTS.tts.layers.xtts.tokenizer import VoiceBpeTokenizer  # Import the tokenizer

# Device configuration
device = "cuda:0" if torch.cuda.is_available() else "cpu"

# Model paths
checkpoint_dir = "/export/home/vivian/svarah/XTTSv2-Finetuning-for-New-Languages/checkpoints/GPT_XTTS_FT-March-12-2025_11+28PM-8e59ec3"
# xtts_checkpoint = os.path.join(checkpoint_dir, "checkpoint_18000.pth")
xtts_checkpoint = os.path.join(checkpoint_dir, "best_model.pth")
xtts_config = os.path.join(checkpoint_dir, "config.json")
xtts_vocab = "/export/home/vivian/svarah/XTTSv2-Finetuning-for-New-Languages/checkpoints/XTTS_v2.0_original_model_files/vocab.json"  # Path to vocab file
speaker_file_path = os.path.join(checkpoint_dir, "speakers_xtts.pth")  # Path to speaker file

# Load model
config = XttsConfig()
config.load_json(xtts_config)

# Initialize the tokenizer
tokenizer = VoiceBpeTokenizer(xtts_vocab)  # Manually initialize the tokenizer

# Initialize the model
XTTS_MODEL = Xtts.init_from_config(config)

# Load checkpoint
XTTS_MODEL.load_checkpoint(
    config,
    checkpoint_path=xtts_checkpoint,
    checkpoint_dir=checkpoint_dir,  # Explicitly provide checkpoint_dir
    vocab_path=xtts_vocab,
    speaker_file_path=speaker_file_path,  # Explicitly provide speaker file path
    use_deepspeed=False,
)
XTTS_MODEL.to(device)
# Manually set the tokenizer

print("Model loaded successfully!")

# Inference
tts_text = "दोस्तों इंडिया ने आईसीसी चैंपियंस ट्रॉफी twentytwentyfive जीत ली है! और उसके बाद एसएस राजामौली, रोहित शर्मा की बायोपिक बनाने वाले हैं, जिसमें उनका रोल उनके ही जैसे नयन-नक्श और कद-काठी रखने वाले जूनियर एनटीआर करेंगे।फिल्म में विराट कोहली का रोल रामचरण निभाएंगे, वहीं एमएस धोनी का रोल प्रभास को मिला है, क्योंकि उनके नाम में भी सात लेटर ही आते हैं। "# थाला फॉर अ रीजन..सभी स्टार्स अपने रोल की तैयारी में जीजान से जुट गए हैं। जहां रामचरण लंदन जा रहे हैं, वहीं जूनियर एनटीआर अपना वजन बढ़ाने के साथ-साथ खुद को भूलने की बीमारी कराने के लिए अपने सिर पर चीजें गिरवा रहे हैं। जबकि प्रभास, फिल्म में लीड रोल ना होने के बावजूद, सारा क्रेडिट खुद लेना सीख रहे हैं।फिल्म में युजवेंद्र चहल भी होंगे, जिनका रोल डॉली चायवाला निभाने वाले हैं।और हाँ, यह वीडियो audiopod dot ai की मदद से बनाया गया है, क्योंकि हमें लगा कि इतनी क्रिएटिव और फनी स्क्रिप्ट को सिर्फ टेक्स्ट में छोड़ना जस्टिस नहीं होगा"
# tts_text="सुबह की सुनहरी किरणों ने मेरे चेहरे पर मुस्कान ला दी। आज का दिन बहुत खूबसूरत है, आसमान में बादल नाच रहे हैं और हवा में खुशियों की लहर है। मैंने अपने परिवार के साथ मीठे पकवान बनाए और हम सब ने मिलकर संगीत के साथ उत्सव मनाया।"
speaker_audio_file = "/export/home/vivian/svarah/Karan (happy)-20250311T073926Z-001/Karan (happy)/karan happy emotions new1.90.wav"
# speaker_audio_file="/export/home/vivian/svarah/Karan Narration-20250311T073420Z-001/Karan Narration/Karan narration emotions.02.wav"
# speaker_audio_file="/export/home/vivian/svarah/XTTSv2-Finetuning-for-New-Languages/checkpoints/GPT_XTTS_FT-March-12-2025_11+28PM-8e59ec3/pratap_hindi - denoised.wav"
# speaker_audio_file="/export/home/vivian/svarah/XTTSv2-Finetuning-for-New-Languages/checkpoints/GPT_XTTS_FT-March-12-2025_11+28PM-8e59ec3/rakesh_hindi_recording.wav"
# speaker_audio_file='/export/home/vivian/svarah/kathbath_wav/844424930591664-939-f.wav'
# speaker_audio_file='/export/home/vivian/svarah/XTTSv2-Finetuning-for-New-Languages/checkpoints/GPT_XTTS_FT-March-12-2025_11+28PM-8e59ec3/shahrukh_voice.flac'
# speaker_audio_file='/export/home/vivian/svarah/XTTSv2-Finetuning-for-New-Languages/checkpoints/GPT_XTTS_FT-March-12-2025_11+28PM-8e59ec3/download (2).wav'
lang = "hi"

gpt_cond_latent, speaker_embedding = XTTS_MODEL.get_conditioning_latents(
    audio_path=speaker_audio_file,
    gpt_cond_len=XTTS_MODEL.config.gpt_cond_len,
    max_ref_length=XTTS_MODEL.config.max_ref_len,
    sound_norm_refs=XTTS_MODEL.config.sound_norm_refs,
)

tts_texts = sent_tokenize(tts_text)
num=0
wav_chunks = []
pause_duration = 0.4  # Duration of the pause in seconds
sample_rate = 24000  # Sample rate of the audio
pause_samples = int(pause_duration * sample_rate)  # Convert pause duration to samples

for text in tqdm(tts_texts):
    num=num+1
    wav_chunk = XTTS_MODEL.inference(
        text=text,
        language=lang,
        gpt_cond_latent=gpt_cond_latent,
        speaker_embedding=speaker_embedding,
        # temperature=0.6,
        # length_penalty=1.0,
        # repetition_penalty=2.0,
        # top_k=60,
        # top_p=0.8, 
        temperature=0.1,
        length_penalty=1.0,
        repetition_penalty=2.0,
        top_k=10,
        top_p=0.3,
    )
    wav_chunks.append(torch.tensor(wav_chunk["wav"]))
    # Add a pause (silent segment) after each chunk
    wav_chunks.append(torch.zeros(pause_samples))

# Concatenate all chunks including pauses
out_wav = torch.cat(wav_chunks, dim=0).unsqueeze(0).cpu()
print(num)
# Play audio (for Jupyter Notebook)
output_audio_path = "/export/home/vivian/svarah/XTTSv2-Finetuning-for-New-Languages/checkpoints/output_audio.wav"
torchaudio.save(output_audio_path, out_wav, sample_rate=24000)

print(f"Audio saved to {output_audio_path}")