import os
import torch
import se_extractor
from api import BaseSpeakerTTS, ToneColorConverter
ckpt_base = 'checkpoints/base_speakers/EN'
ckpt_converter = 'checkpoints/converter'
device = 'cuda:0'
output_dir = 'outputs'

base_speaker_tts = BaseSpeakerTTS(f'{ckpt_base}/config.json', device=device)
base_speaker_tts.load_ckpt(f'{ckpt_base}/checkpoint.pth')

tone_color_converter = ToneColorConverter(f'{ckpt_converter}/config.json', device=device)
tone_color_converter.load_ckpt(f'{ckpt_converter}/checkpoint.pth')

os.makedirs(output_dir, exist_ok=True)

source_se = torch.load(f'{ckpt_base}/en_default_se.pth').to(device)

reference_speaker = '/root/src/videly/voices/m-us-3.wav'
target_se, audio_name = se_extractor.get_se(reference_speaker, tone_color_converter, target_dir='processed', vad=True)
save_path = f'{output_dir}/output_en_default.wav'

# Run the base speaker tts
text = '''
Harry Potter stood silently at the edge of the Forbidden Forest, his wand gripped tightly in his hand. The moon cast a silvery glow over the ancient trees, creating a scene both eerie and beautiful. Harry's mind was racing with thoughts of the recent events at Hogwarts. The castle, usually a place of learning and friendship, had transformed into a battleground where the forces of good and evil clashed.

In the distance, the unmistakable silhouette of Hogwarts Castle stood against the starry sky, its spires reaching up like fingers grasping for hope. Harry felt a deep connection to the place, not just as a student, but as someone who had faced and overcome great trials within its walls. He thought of his friends Ron and Hermione, who had been with him through thick and thin, their loyalty never wavering.

As an owl hooted in the distance, Harry remembered the letters that had started his journey into the wizarding world. He had been just a boy then, unaware of his lineage and the future that awaited him. Now, he stood as a symbol of courage and resilience, a young wizard who had faced Lord Voldemort and lived to tell the tale.

The events of the Triwizard Tournament flashed through his mind. The excitement, the danger, the tragic loss of Cedric Diggory - it was a reminder of the thin line between life and death in their world. Harry felt a pang of sadness, knowing that the path ahead was fraught with peril. Yet, he also felt a strong sense of purpose. He was not just fighting for himself, but for the entire wizarding community, for a world free from the tyranny of Voldemort.
'''
src_path = f'{output_dir}/tmp.wav'
base_speaker_tts.tts(text, src_path, speaker='default', language='English', speed=1.0)

# Run the tone color converter
tone_color_converter.convert(
    audio_src_path=src_path, 
    src_se=source_se, 
    tgt_se=target_se, 
    output_path=save_path,
    message='')