import argparse
import json
import os
from pathlib import Path

import torch
from transformers import Seq2SeqTrainingArguments, Seq2SeqTrainer

from deep_voice_cloning.cloning.model import CloningModel
from deep_voice_cloning.transcriber.model import TranscriberModel
from deep_voice_cloning.data.collator import TTSDataCollatorWithPadding
from deep_voice_cloning.data.dataset import get_cloning_dataset


if __name__ == "__main__":
    parser = argparse.ArgumentParser()
    parser.add_argument("--lang", type=str, default=None, help="Language of speech samples")
    parser.add_argument("--audio_path", type=str, default=None, help="Path to training audio file")
    parser.add_argument("--output_dir", type=str, default=None, help="Path to output directory for trained model")
    args = parser.parse_args()

    with open(os.path.join(os.path.dirname(__file__), "training_config.json")) as f:
        training_config = json.load(f)

    if args.lang is not None:
        training_config['lang'] = args.lang
    if args.audio_path is not None:
        training_config['audio_path'] = Path(args.audio_path)
    if args.output_dir is not None:
        training_config['output_dir'] = Path(args.output_dir)

    transcriber_model = TranscriberModel(lang=training_config['lang'])
    cloning_model = CloningModel(lang=training_config['lang'])

    dataset = get_cloning_dataset(training_config['audio_path'], transcriber_model, cloning_model)
    data_collator = TTSDataCollatorWithPadding(processor=cloning_model.processor, model=cloning_model.model)

    training_args = Seq2SeqTrainingArguments(
        output_dir=training_config["output_dir"],
        per_device_train_batch_size=training_config['batch_size'],
        gradient_accumulation_steps=2,
        overwrite_output_dir=True,
        learning_rate=training_config['learning_rate'],
        warmup_steps=training_config['warmup_steps'],
        max_steps=training_config['max_steps'],
        gradient_checkpointing=True,
        fp16=transcriber_model.device == torch.device("cuda"),
        evaluation_strategy="steps",
        per_device_eval_batch_size=8,
        save_strategy="no",
        eval_steps=100,
        logging_steps=20,
        load_best_model_at_end=False,
        greater_is_better=False,
        label_names=["labels"],
    )

    trainer = Seq2SeqTrainer(
        args=training_args,
        model=cloning_model.model,
        train_dataset=dataset,
        eval_dataset=dataset,
        data_collator=data_collator,
        tokenizer=cloning_model.processor.tokenizer,
    )

    trainer.train()
    cloning_model.save_pretrained(Path(training_config["output_dir"]) /
                                  Path(cloning_model.config['model_path'].replace('/', '_')
                                       + '_' + Path(training_config['audio_path']).stem)
                                  )