import argparse
import os
import torch
import torchaudio
import text
import utils.make_html as html

from utils import progbar, read_lines_from_file

# default:
# python inference.py --list data/infer_text.txt --out_dir samples/results --model fastpitch --checkpoint pretrained/fastpitch_ar_adv.pth --batch_size 2 --denoise 0

# Examples:
# python inference.py --list data/infer_text.txt --out_dir samples/res_tc2_adv0 --model tacotron2 --checkpoint pretrained/tacotron2_ar_adv.pth --batch_size 2
# python inference.py --list data/infer_text.txt --out_dir samples/res_tc2_adv1 --model tacotron2 --checkpoint pretrained/tacotron2_ar_adv.pth --batch_size 2 --denoise 0.005
# python inference.py --list data/infer_text.txt --out_dir samples/res_fp_adv0 --model fastpitch --checkpoint pretrained/fastpitch_ar_adv.pth --batch_size 2
# python inference.py --list data/infer_text.txt --out_dir samples/res_fp_adv1 --model fastpitch --checkpoint pretrained/fastpitch_ar_adv.pth --batch_size 2 --denoise 0.005
# python inference.py --list data/infer_text.txt --out_dir samples/res_fp_adv2 --model fastpitch --checkpoint pretrained/fastpitch_ar_adv.pth --batch_size 2 --denoise 0.005 --vocoder_sd pretrained/hifigan-asc-v1/g_02500000 --vocoder_config pretrained/hifigan-asc-v1/config.json


def infer(args):
    use_cuda_if_available = not args.cpu
    device = torch.device(
        'cuda' if torch.cuda.is_available() and use_cuda_if_available else 'cpu')

    if args.model == 'fastpitch':
        from models.fastpitch import FastPitch2Wave
        model = FastPitch2Wave(args.checkpoint, 
                               vocoder_sd=args.vocoder_sd,
                               vocoder_config=args.vocoder_config)
    elif args.model == 'tacotron2':
        from models.tacotron2 import Tacotron2Wave
        model = Tacotron2Wave(args.checkpoint, 
                              vocoder_sd=args.vocoder_sd,
                              vocoder_config=args.vocoder_config)
    else:
        raise "model type not supported"

    model = model.to(device)
    model.eval()

    if not os.path.exists(f"{args.out_dir}/wavs"):
        os.makedirs(f"{args.out_dir}/wavs")

    static_lines = read_lines_from_file(args.list)
    static_batches = [static_lines[k:k+args.batch_size]
                      for k in range(0, len(static_lines), args.batch_size)]

    idx = 0
    with open(os.path.join(args.out_dir, 'index.html'), 'w', encoding='utf-8') as f:
        f.write(html.make_html_start())

        for batch in progbar(static_batches):
            # infer batch
            wav_list = model.tts(batch,
                                 batch_size=args.batch_size,
                                 denoise=args.denoise,
                                 speed=args.speed)

            # save wavs and add entries to html file
            for (text_line, wav) in zip(batch, wav_list):
                torchaudio.save(f'{args.out_dir}/wavs/static{idx}.wav',
                                wav.unsqueeze(0),
                                22_050)

                text_buckw = text.arabic_to_buckwalter(text_line)
                text_arabic = text.buckwalter_to_arabic(text_buckw)
                t_phon = text.buckwalter_to_phonemes(text_buckw)
                t_phon = text.simplify_phonemes(
                    t_phon.replace(' ', '').replace('+', ' '))

                f.write(html.make_sample_entry2(
                    f'wavs/static{idx}.wav',
                    text_arabic,
                    f"{idx}) {t_phon}"))

                idx += 1

        f.write(html.make_volume_script(0.5))
        f.write(html.make_html_end())

    print(f"Saved files to: {args.out_dir}")


def main():
    parser = argparse.ArgumentParser()
    parser.add_argument(
        '--list', type=str, default='./data/infer_text.txt')
    parser.add_argument(
        '--model', type=str, default='fastpitch')
    parser.add_argument(
        '--checkpoint', type=str, default='pretrained/fastpitch_ar_adv.pth')
    parser.add_argument('--vocoder_sd', type=str, default=None)
    parser.add_argument('--vocoder_config', type=str, default=None) 
    parser.add_argument('--out_dir', type=str, default='samples/results')
    parser.add_argument('--speed', type=float, default=1.0)
    parser.add_argument('--denoise', type=float, default=0)
    parser.add_argument('--batch_size', type=int, default=2)
    parser.add_argument('--cpu', action='store_true')
    args = parser.parse_args()

    infer(args)


if __name__ == '__main__':
    main()