# Copyright (c) 2021, NVIDIA CORPORATION.  All rights reserved.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#     http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.

# USAGE: python get_aishell_data.py --data_root=<where to put data>

import argparse
import json
import logging
import os
import subprocess
import tarfile
import urllib.request

from tqdm import tqdm

parser = argparse.ArgumentParser(description="Aishell Data download")
parser.add_argument("--data_root", required=True, default=None, type=str)
args = parser.parse_args()

URL = {"data_aishell": "http://www.openslr.org/resources/33/data_aishell.tgz"}


def __retrieve_with_progress(source: str, filename: str):
    """
    Downloads source to destination
    Displays progress bar
    Args:
        source: url of resource
        destination: local filepath
    Returns:
    """
    with open(filename, "wb") as f:
        response = urllib.request.urlopen(source)
        total = response.length

        if total is None:
            f.write(response.content)
        else:
            with tqdm(total=total, unit="B", unit_scale=True, unit_divisor=1024) as pbar:
                for data in response:
                    f.write(data)
                    pbar.update(len(data))


def __maybe_download_file(destination: str, source: str):
    """
    Downloads source to destination if it doesn't exist.
    If exists, skips download
    Args:
        destination: local filepath
        source: url of resource

    Returns:

    """
    source = URL[source]
    if not os.path.exists(destination):
        logging.info("{0} does not exist. Downloading ...".format(destination))
        __retrieve_with_progress(source, filename=destination + ".tmp")
        os.rename(destination + ".tmp", destination)
        logging.info("Downloaded {0}.".format(destination))
    else:
        logging.info("Destination {0} exists. Skipping.".format(destination))
    return destination


def __extract_all_files(filepath: str, data_root: str, data_dir: str):
    if not os.path.exists(data_dir):
        extract_file(filepath, data_root)
        audio_dir = os.path.join(data_dir, "wav")
        for subfolder, _, filelist in os.walk(audio_dir):
            for ftar in filelist:
                extract_file(os.path.join(subfolder, ftar), subfolder)
    else:
        logging.info("Skipping extracting. Data already there %s" % data_dir)


def extract_file(filepath: str, data_dir: str):
    try:
        tar = tarfile.open(filepath)
        tar.extractall(data_dir)
        tar.close()
    except Exception:
        logging.info("Not extracting. Maybe already there?")


def __process_data(data_folder: str, dst_folder: str):
    """
    To generate manifest
    Args:
        data_folder: source with wav files
        dst_folder: where manifest files will be stored
    Returns:

    """

    if not os.path.exists(dst_folder):
        os.makedirs(dst_folder)

    transcript_file = os.path.join(data_folder, "transcript", "aishell_transcript_v0.8.txt")
    transcript_dict = {}
    with open(transcript_file, "r", encoding="utf-8") as f:
        for line in f:
            line = line.strip()
            audio_id, text = line.split(" ", 1)
            # remove white space
            text = text.replace(" ", "")
            transcript_dict[audio_id] = text

    data_types = ["train", "dev", "test"]
    vocab_count = {}
    for dt in data_types:
        json_lines = []
        audio_dir = os.path.join(data_folder, "wav", dt)
        for sub_folder, _, file_list in os.walk(audio_dir):
            for fname in file_list:
                audio_path = os.path.join(sub_folder, fname)
                audio_id = fname.strip(".wav")
                if audio_id not in transcript_dict:
                    continue
                text = transcript_dict[audio_id]
                for li in text:
                    vocab_count[li] = vocab_count.get(li, 0) + 1
                duration = subprocess.check_output("soxi -D {0}".format(audio_path), shell=True)
                duration = float(duration)
                json_lines.append(
                    json.dumps(
                        {"audio_filepath": os.path.abspath(audio_path), "duration": duration, "text": text,},
                        ensure_ascii=False,
                    )
                )

        manifest_path = os.path.join(dst_folder, dt + ".json")
        with open(manifest_path, "w", encoding="utf-8") as fout:
            for line in json_lines:
                fout.write(line + "\n")

    vocab = sorted(vocab_count.items(), key=lambda k: k[1], reverse=True)
    vocab_file = os.path.join(dst_folder, "vocab.txt")
    with open(vocab_file, "w", encoding="utf-8") as f:
        for v, c in vocab:
            f.write(v + "\n")


def main():
    data_root = args.data_root
    data_set = "data_aishell"
    logging.info("\n\nWorking on: {0}".format(data_set))
    file_path = os.path.join(data_root, data_set + ".tgz")
    logging.info("Getting {0}".format(data_set))
    __maybe_download_file(file_path, data_set)
    logging.info("Extracting {0}".format(data_set))
    data_folder = os.path.join(data_root, data_set)
    __extract_all_files(file_path, data_root, data_folder)
    logging.info("Processing {0}".format(data_set))
    __process_data(data_folder, data_folder)
    logging.info("Done!")


if __name__ == "__main__":
    main()