#!/usr/bin/env python3

# Copyright 2021 Carnegie Mellon University (Peter Wu)
#  Apache 2.0  (http://www.apache.org/licenses/LICENSE-2.0)


import argparse
import os
import random


if __name__ == "__main__":
    parser = argparse.ArgumentParser()
    parser.add_argument("-d", help="downloads directory", type=str, default="downloads")
    args = parser.parse_args()

    tsv_path = "%s/line_index.tsv" % args.d

    with open(tsv_path, "r") as inf:
        tsv_lines = inf.readlines()
    tsv_lines = [line.strip() for line in tsv_lines]

    spk2utt = {}
    utt2text = {}
    for line in tsv_lines:
        l_list = line.split("\t")
        fid = l_list[0]
        spk = l_list[0].split("_")[1]
        text = l_list[1]
        path = "%s/%s.wav" % (args.d, fid)
        if os.path.exists(path):
            utt2text[fid] = text
            if spk in spk2utt:
                spk2utt[spk].append(fid)
            else:
                spk2utt[spk] = [fid]

    spks = sorted(list(spk2utt.keys()))
    num_fids = 0
    num_test_spks = 0
    for spk in spks:
        num_test_spks += 1
        fids = sorted(list(set(spk2utt[spk])))
        num_fids += len(fids)
        if num_fids >= 2000:
            break

    num_test_spks = 2
    test_spks = spks[:num_test_spks]
    train_dev_spks = spks[num_test_spks:]
    random.Random(0).shuffle(train_dev_spks)
    num_train = int(len(train_dev_spks) * 0.9)
    train_spks = train_dev_spks[:num_train]
    dev_spks = train_dev_spks[num_train:]

    spks_by_phase = {"train": train_spks, "dev": dev_spks, "test": test_spks}
    flac_dir = "%s" % args.d
    sr = 16000
    for phase in spks_by_phase:
        spks = spks_by_phase[phase]
        text_strs = []
        wav_scp_strs = []
        spk2utt_strs = []
        num_fids = 0
        for spk in spks:
            fids = sorted(list(set(spk2utt[spk])))
            num_fids += len(fids)
            if phase == "test" and num_fids > 2000:
                curr_num_fids = num_fids - 2000
                random.Random(1).shuffle(fids)
                fids = fids[:curr_num_fids]
            utts = [spk + "-" + f for f in fids]
            utts_str = " ".join(utts)
            spk2utt_strs.append("%s %s" % (spk, utts_str))
            for fid, utt in zip(fids, utts):
                cmd = "ffmpeg -i %s/%s.wav -f wav -ar %d -ab 16 -ac 1 - |" % (
                    flac_dir,
                    fid,
                    sr,
                )
                text_strs.append("%s %s" % (utt, utt2text[fid]))
                wav_scp_strs.append("%s %s" % (utt, cmd))
        phase_dir = "data/marathi_%s" % phase
        if not os.path.exists(phase_dir):
            os.makedirs(phase_dir)

        text_strs = sorted(text_strs)
        wav_scp_strs = sorted(wav_scp_strs)
        spk2utt_strs = sorted(spk2utt_strs)

        with open(os.path.join(phase_dir, "text"), "w+") as ouf:
            for s in text_strs:
                ouf.write("%s\n" % s)
        with open(os.path.join(phase_dir, "wav.scp"), "w+") as ouf:
            for s in wav_scp_strs:
                ouf.write("%s\n" % s)
        with open(os.path.join(phase_dir, "spk2utt"), "w+") as ouf:
            for s in spk2utt_strs:
                ouf.write("%s\n" % s)

[(13, 'argparse.ArgumentParser', 'argparse.ArgumentParser', False, 'import argparse\n'), (31, 'os.path.exists', 'os.path.exists', False, 'import os\n'), (51, 'random.Random', 'random.Random', False, 'import random\n'), (84, 'os.path.exists', 'os.path.exists', False, 'import os\n'), (85, 'os.makedirs', 'os.makedirs', False, 'import os\n'), (91, 'os.path.join', 'os.path.join', False, 'import os\n'), (94, 'os.path.join', 'os.path.join', False, 'import os\n'), (97, 'os.path.join', 'os.path.join', False, 'import os\n'), (70, 'random.Random', 'random.Random', False, 'import random\n')]