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-February-25-2025_12+32AM-8e59ec3"
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 = "கிராமத்தில் வசிக்கும் குமரன் ஒரு நேர்மையான விவசாயி. ஒரு நாள், அவன் வயலில் வேலைசெய்துக்கொண்டிருக்கும்போது, ஒரு பசி பித்தியான கழுகு பிணியாக விழுந்தது. அதன் இறக்கைகள் உடைந்திருந்ததால் அது பறக்க முடியவில்லை. இரக்கமடைந்த குமரன், அதைப் பிடித்து வீட்டிற்கு கொண்டு சென்று சிகிச்சை அளித்தான். சில வாரங்கள் கழித்து, கழுகு மீண்டும் வலிமை பெற்றது. விடைபெறுவதற்காக அது குமரனை நேசமுடன் பார்த்து ஒருமுறை பறக்க முயன்றது. அந்தக் கணமே, வானத்தில் வீசும் காற்றில் தனது சக்தியை உணர்ந்தது. புன்னகைத்த குமரன், 'நம்மை மீண்டும் எழுப்பும் வரை நேர்மையும் உதவியும் தேவை,' என்று மனதில் நினைத்தான்."
speaker_audio_file = "/export/home/vivian/svarah/XTTSv2-Finetuning-for-New-Languages/checkpoints/GPT_XTTS_FT-February-25-2025_12+32AM-8e59ec3/download (1).mp3"

lang = "ta"

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.1,
        length_penalty=1.0,
        repetition_penalty=10.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}")