import os

from trainer import Trainer, TrainerArgs

from TTS.tts.configs.shared_configs import BaseDatasetConfig, CharactersConfig
from TTS.tts.configs.vits_config import VitsConfig
from TTS.tts.datasets import load_tts_samples
from TTS.tts.models.vits import Vits, VitsAudioConfig
from TTS.tts.utils.text.tokenizer import TTSTokenizer
from TTS.utils.audio import AudioProcessor


if __name__ == "__main__":
    output_path = os.path.dirname(os.path.abspath(__file__))
    dataset_config = BaseDatasetConfig(
        formatter="ljspeech", meta_file_train="metadata.csv", path=os.path.join(output_path, "/home/fatima-w/workpj/texttospeech/yal_asheh1_book")
    )

    audio_config = VitsAudioConfig(
        sample_rate=22050, 
        win_length=1024, 
        hop_length=256, 
        num_mels=80, 
        mel_fmin=0, 
        mel_fmax=None
    )

    character_config = CharactersConfig(
        characters="ىءۀابتثجحخدذرزسشصضطظعغفقلمنهويِپچژکگیآأؤإئًَُّْ'ٔ٪۰۱۲۳۴۵۶۷۸۹ةabcdefghijklmnopqrstuvwxyzABCDEFGHIJKLMNOPQRSTUVWXYZكأي",
        # characters="ءۀابتثجحخدذرزسشصضطظعغفقلمنهويِپچژکگیآأؤإئًَُّ۰۱۲۳۴۵۶۷۸۹ة",
        punctuations='!(),-.–:;«»?— ̠،؛؟‌<>"',
        phonemes=None,
        pad="<PAD>",
        eos="<EOS>",
        bos="<BOS>",
        blank="<BLNK>",
        characters_class=None,
    )

    config = VitsConfig(
        audio=audio_config,
        run_name="vits_farsi",
        batch_size=20,
        eval_batch_size=10,
        batch_group_size=5,
        num_loader_workers=4,
        num_eval_loader_workers=2,
        run_eval=True,
        test_delay_epochs=-1,
        epochs=1000,
        text_cleaner="basic_cleaners",
        use_phonemes=False,
        phoneme_language="fa",
        phoneme_cache_path=os.path.join(output_path, "phoneme_cache"),
        compute_input_seq_cache=True,
        print_step=25,
        print_eval=True,
        mixed_precision=True,  # Keep mixed precision enabled
        output_path=output_path,
        datasets=[dataset_config],
        cudnn_benchmark=False,
        characters=character_config,
        test_sentences=[
        ["سلطان محمود در زمستانی سخت به طلخک گفت که: با این جامه ی یک لا در این سرما چه می کنی "],
        ["مردی نزد بقالی آمد و گفت پیاز هم ده تا دهان بدان خو شبوی سازم."],
        ["از مال خود پاره ای گوشت بستان و زیره بایی معطّر بساز"],
        ["یک بار هم از جهنم بگویید."],
        ["یکی اسبی به عاریت خواست"]
    ],
    )

   
    # INITIALIZE THE AUDIO PROCESSOR
    ap = AudioProcessor.init_from_config(config)

    # INITIALIZE THE TOKENIZER
    tokenizer, config = TTSTokenizer.init_from_config(config)

    # LOAD DATA SAMPLES
    train_samples, eval_samples = load_tts_samples(
        dataset_config,
        eval_split=True,
        eval_split_max_size=config.eval_split_max_size,
        eval_split_size=config.eval_split_size,
    )

    # INIT MODEL
    model = Vits(config, ap, tokenizer, speaker_manager=None)

    checkpoint_path = os.path.join('/home/fatima-w/workpj/coqui/vits_farsi-April-10-2025_10+51PM-0000000', "best_model.pth")

    # INIT THE TRAINER AND START TRAINING
    trainer = Trainer(
       TrainerArgs(restore_path=checkpoint_path),
        # TrainerArgs(),
        config,
        output_path,
        model=model,
        train_samples=train_samples,
        eval_samples=eval_samples,
    )
    trainer.fit()