# Copyright (c) 2022, 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 process_speech_commands_data.py \
    --data_root=<absolute path to where the data should be stored> \
    --data_version=<either 1 or 2, indicating version of the dataset> \
    --class_split=<either "all" or "sub", indicates whether all 30/35 classes should be used, or the 10+2 split should be used> \
    --num_processes=<number of processes to use for data preprocessing> \
    --rebalance \
    --log
"""

import argparse
import glob
import json
import logging
import os
import re
import tarfile
import urllib.request
from collections import defaultdict
from functools import partial
from multiprocessing import Pool
from typing import Dict, List, Set, Tuple

import librosa
import numpy as np
import soundfile
from tqdm import tqdm

URL_v1 = 'http://download.tensorflow.org/data/speech_commands_v0.01.tar.gz'
URL_v2 = 'http://download.tensorflow.org/data/speech_commands_v0.02.tar.gz'


def __maybe_download_file(destination: str, source: str) -> str:
    """
    Downloads source to destination if it doesn't exist.
    If exists, skips download

    Args:
        destination: local filepath
        source: url of resource

    Returns:
        Local filepath of the downloaded file
    """
    if not os.path.exists(destination):
        logging.info(f'{destination} does not exist. Downloading ...')
        urllib.request.urlretrieve(source, filename=destination + '.tmp')
        os.rename(destination + '.tmp', destination)
        logging.info(f'Downloaded {destination}.')
    else:
        logging.info(f'Destination {destination} exists. Skipping.')
    return destination


def __extract_all_files(filepath: str, data_dir: str):
    if not os.path.exists(data_dir):
        extract_file(filepath, data_dir)
    else:
        logging.info(f'Skipping extracting. Data already there {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 __get_mp_chunksize(dataset_size: int, num_processes: int) -> int:
    """
    Returns the number of chunks to split the dataset into for multiprocessing.

    Args:
        dataset_size: size of the dataset
        num_processes: number of processes to use for multiprocessing

        Returns:
            Number of chunks to split the dataset into for multiprocessing
    """
    chunksize = dataset_size // num_processes
    return chunksize if chunksize > 0 else 1


def __construct_filepaths(
    all_files: List[str],
    valset_uids: Set[str],
    testset_uids: Set[str],
    class_split: str,
    class_subset: List[str],
    pattern: str,
) -> Tuple[Dict[str, int], Dict[str, List[tuple]], List[tuple], List[tuple], List[tuple], List[tuple], List[tuple]]:
    """
    Prepares the filepaths for the dataset.

    Args:
        all_files: list of all files in the dataset
        valset_uids: set of uids of files in the validation set
        testset_uids: set of uids of files in the test set
        class_split: whether to use all classes as distinct labels, or to use
            10 classes subset and rest of the classes as noise or background
        class_subset: list of classes to consider if `class_split` is set to `sub`
        pattern: regex pattern to match the file names in the dataset
    """

    label_count = defaultdict(int)
    label_filepaths = defaultdict(list)
    unknown_val_filepaths = []
    unknown_test_filepaths = []

    train, val, test = [], [], []
    for entry in all_files:
        r = re.match(pattern, entry)
        if r:
            label, uid = r.group(2), r.group(3)

            if label == '_background_noise_' or label == 'silence':
                continue

        if class_split == 'sub' and label not in class_subset:
            label = 'unknown'

            if uid in valset_uids:
                unknown_val_filepaths.append((label, entry))
            elif uid in testset_uids:
                unknown_test_filepaths.append((label, entry))

        if uid not in valset_uids and uid not in testset_uids:
            label_count[label] += 1
            label_filepaths[label].append((label, entry))

        if label == 'unknown':
            continue

        if uid in valset_uids:
            val.append((label, entry))
        elif uid in testset_uids:
            test.append((label, entry))
        else:
            train.append((label, entry))

    return {
        'label_count': label_count,
        'label_filepaths': label_filepaths,
        'unknown_val_filepaths': unknown_val_filepaths,
        'unknown_test_filepaths': unknown_test_filepaths,
        'train': train,
        'val': val,
        'test': test,
    }


def __construct_silence_set(
    rng: np.random.RandomState, sampling_rate: int, silence_stride: int, data_folder: str, background_noise: str
) -> List[str]:
    """
    Creates silence files given a background noise.

    Args:
        rng: Random state for random number generator
        sampling_rate: sampling rate of the audio
        silence_stride: stride for creating silence files
        data_folder: folder containing the silence directory
        background_noise: filepath of the background noise

    Returns:
        List of filepaths of silence files
    """
    silence_files = []
    if '.wav' in background_noise:
        y, sr = librosa.load(background_noise, sr=sampling_rate)

        for i in range(0, len(y) - sampling_rate, silence_stride):
            file_path = f'silence/{os.path.basename(background_noise)[:-4]}_{i}.wav'
            y_slice = y[i : i + sampling_rate] * rng.uniform(0.0, 1.0)
            out_file_path = os.path.join(data_folder, file_path)
            soundfile.write(out_file_path, y_slice, sr)

            silence_files.append(('silence', out_file_path))

    return silence_files


def __rebalance_files(max_count: int, label_filepath: str) -> Tuple[str, List[str], int]:
    """
    Rebalance the number of samples for a class.

    Args:
        max_count: maximum number of samples for a class
        label_filepath: list of filepaths for a class

    Returns:
        Rebalanced list of filepaths along with the label name and the number of samples
    """
    command, samples = label_filepath
    filepaths = [sample[1] for sample in samples]

    rng = np.random.RandomState(0)
    filepaths = np.asarray(filepaths)
    num_samples = len(filepaths)

    if num_samples < max_count:
        difference = max_count - num_samples
        duplication_ids = rng.choice(num_samples, difference, replace=True)
        filepaths = np.append(filepaths, filepaths[duplication_ids], axis=0)

    return command, filepaths, num_samples


def __prepare_metadata(skip_duration, sample: Tuple[str, str]) -> dict:
    """
    Creates the manifest entry for a file.

    Args:
        skip_duration: Whether to skip the computation of duration
        sample: Tuple of label and filepath

    Returns:
        Manifest entry of the file
    """
    label, audio_path = sample
    return json.dumps(
        {
            'audio_filepath': audio_path,
            'duration': 0.0 if skip_duration else librosa.core.get_duration(filename=audio_path),
            'command': label,
        }
    )


def __process_data(
    data_folder: str,
    dst_folder: str,
    num_processes: int = 1,
    rebalance: bool = False,
    class_split: str = 'all',
    skip_duration: bool = False,
):
    """
    Processes the data and generates the manifests.

    Args:
        data_folder: source with wav files and validation / test lists
        dst_folder: where manifest files will be stored
        num_processes: number of processes
        rebalance: rebalance the classes to have same number of samples
        class_split: whether to use all classes as distinct labels, or to use
            10 classes subset and rest of the classes as noise or background
        skip_duration: Bool whether to skip duration computation. Use this only for
            colab notebooks where knowing duration is not necessary for demonstration
    """

    os.makedirs(dst_folder, exist_ok=True)

    # Used for 10 classes + silence + unknown class setup - Only used when class_split is 'sub'
    class_subset = ['yes', 'no', 'up', 'down', 'left', 'right', 'on', 'off', 'stop', 'go']

    pattern = re.compile(r'(.+\/)?(\w+)\/([^_]+)_.+wav')
    all_files = glob.glob(os.path.join(data_folder, '*/*wav'))

    # Get files in the validation set
    valset_uids = set()
    with open(os.path.join(data_folder, 'validation_list.txt')) as fin:
        for line in fin:
            r = re.match(pattern, line)
            if r:
                valset_uids.add(r.group(3))

    # Get files in the test set
    testset_uids = set()
    with open(os.path.join(data_folder, 'testing_list.txt')) as fin:
        for line in fin:
            r = re.match(pattern, line)
            if r:
                testset_uids.add(r.group(3))

    logging.info('Validation and test set lists extracted')

    filepath_info = __construct_filepaths(all_files, valset_uids, testset_uids, class_split, class_subset, pattern)
    label_count = filepath_info['label_count']
    label_filepaths = filepath_info['label_filepaths']
    unknown_val_filepaths = filepath_info['unknown_val_filepaths']
    unknown_test_filepaths = filepath_info['unknown_test_filepaths']
    train = filepath_info['train']
    val = filepath_info['val']
    test = filepath_info['test']

    logging.info('Prepared filepaths for dataset')

    pool = Pool(num_processes)

    # Add silence and unknown class label samples
    if class_split == 'sub':
        logging.info('Perforiming 10+2 class subsplit')

        silence_path = os.path.join(data_folder, 'silence')
        os.makedirs(silence_path, exist_ok=True)

        silence_stride = 1000  # 0.0625 second stride
        sampling_rate = 16000
        folder = os.path.join(data_folder, '_background_noise_')

        silence_files = []
        rng = np.random.RandomState(0)

        background_noise_files = [os.path.join(folder, x) for x in os.listdir(folder)]
        silence_set_fn = partial(__construct_silence_set, rng, sampling_rate, silence_stride, data_folder)
        for silence_flist in tqdm(
            pool.imap(
                silence_set_fn, background_noise_files, __get_mp_chunksize(len(background_noise_files), num_processes)
            ),
            total=len(background_noise_files),
            desc='Constructing silence set',
        ):
            silence_files.extend(silence_flist)

        rng = np.random.RandomState(0)
        rng.shuffle(silence_files)
        logging.info(f'Constructed silence set of {len(silence_files)}')

        # Create the splits
        rng = np.random.RandomState(0)
        silence_split = 0.1
        unknown_split = 0.1

        # train split
        num_total_samples = sum([label_count[cls] for cls in class_subset])
        num_silence_samples = int(np.ceil(silence_split * num_total_samples))

        # initialize sample
        label_count['silence'] = 0
        label_filepaths['silence'] = []

        for silence_id in range(num_silence_samples):
            label_count['silence'] += 1
            label_filepaths['silence'].append(silence_files[silence_id])

        train.extend(label_filepaths['silence'])

        # Update train unknown set
        unknown_train_samples = label_filepaths['unknown']

        rng.shuffle(unknown_train_samples)
        unknown_size = int(np.ceil(unknown_split * num_total_samples))

        label_count['unknown'] = unknown_size
        label_filepaths['unknown'] = unknown_train_samples[:unknown_size]

        train.extend(label_filepaths['unknown'])

        logging.info('Train set prepared')

        # val set silence
        num_val_samples = len(val)
        num_silence_samples = int(np.ceil(silence_split * num_val_samples))

        val_idx = label_count['silence'] + 1
        for silence_id in range(num_silence_samples):
            val.append(silence_files[val_idx + silence_id])

        # Update val unknown set
        rng.shuffle(unknown_val_filepaths)
        unknown_size = int(np.ceil(unknown_split * num_val_samples))

        val.extend(unknown_val_filepaths[:unknown_size])

        logging.info('Validation set prepared')

        # test set silence
        num_test_samples = len(test)
        num_silence_samples = int(np.ceil(silence_split * num_test_samples))

        test_idx = val_idx + num_silence_samples + 1
        for silence_id in range(num_silence_samples):
            test.append(silence_files[test_idx + silence_id])

        # Update test unknown set
        rng.shuffle(unknown_test_filepaths)
        unknown_size = int(np.ceil(unknown_split * num_test_samples))

        test.extend(unknown_test_filepaths[:unknown_size])

        logging.info('Test set prepared')

    max_command = None
    max_count = -1
    for command, count in label_count.items():
        if command == 'unknown':
            continue

        if count > max_count:
            max_count = count
            max_command = command

    if rebalance:
        logging.info(f'Command with maximum number of samples = {max_command} with {max_count} samples')
        logging.info(f'Rebalancing dataset by duplicating classes with less than {max_count} samples...')

        rebalance_fn = partial(__rebalance_files, max_count)
        for command, filepaths, num_samples in tqdm(
            pool.imap(rebalance_fn, label_filepaths.items(), __get_mp_chunksize(len(label_filepaths), num_processes)),
            total=len(label_filepaths),
            desc='Rebalancing dataset',
        ):
            if num_samples < max_count:
                logging.info(f'Extended class label {command} from {num_samples} samples to {len(filepaths)} samples')
                label_filepaths[command] = [(command, filepath) for filepath in filepaths]

        del train
        train = []
        for label, samples in label_filepaths.items():
            train.extend(samples)

    manifests = [
        ('train_manifest.json', train),
        ('validation_manifest.json', val),
        ('test_manifest.json', test),
    ]

    metadata_fn = partial(__prepare_metadata, skip_duration)
    for manifest_filename, dataset in manifests:
        num_files = len(dataset)

        logging.info(f'Preparing manifest : {manifest_filename} with #{num_files} files')

        manifest = [
            metadata
            for metadata in tqdm(
                pool.imap(metadata_fn, dataset, __get_mp_chunksize(len(dataset), num_processes)),
                total=num_files,
                desc=f'Preparing {manifest_filename}',
            )
        ]

        with open(os.path.join(dst_folder, manifest_filename), 'w') as fout:
            for metadata in manifest:
                fout.write(metadata + '\n')

        logging.info(f'Finished construction of manifest. Path: {os.path.join(dst_folder, manifest_filename)}')

    pool.close()

    if skip_duration:
        logging.info(
            f'\n<<NOTE>> Duration computation was skipped for demonstration purposes on Colaboratory.\n'
            f'In order to replicate paper results and properly perform data augmentation, \n'
            f'please recompute the manifest file without the `--skip_duration` flag !\n'
        )


def main():
    parser = argparse.ArgumentParser(description='Google Speech Commands Data download and preprocessing')
    parser.add_argument('--data_root', required=True, help='Root directory for storing data')
    parser.add_argument(
        '--data_version',
        required=True,
        default=1,
        type=int,
        choices=[1, 2],
        help='Version of the speech commands dataset to download',
    )
    parser.add_argument(
        '--class_split', default='all', choices=['all', 'sub'], help='Whether to consider all classes or only a subset'
    )
    parser.add_argument('--num_processes', default=1, type=int, help='Number of processes')
    parser.add_argument('--rebalance', action='store_true', help='Rebalance the number of samples in each class')
    parser.add_argument('--skip_duration', action='store_true', help='Skip computing duration of audio files')
    parser.add_argument('--log', action='store_true', help='Generate logs')
    args = parser.parse_args()

    if args.log:
        logging.basicConfig(level=logging.DEBUG)

    data_root = args.data_root
    data_set = f'google_speech_recognition_v{args.data_version}'
    data_folder = os.path.join(data_root, data_set)

    logging.info(f'Working on: {data_set}')

    URL = URL_v1 if args.data_version == 1 else URL_v2

    # Download and extract
    if not os.path.exists(data_folder):
        file_path = os.path.join(data_root, data_set + '.tar.bz2')
        logging.info(f'Getting {data_set}')
        __maybe_download_file(file_path, URL)
        logging.info(f'Extracting {data_set}')
        __extract_all_files(file_path, data_folder)

    logging.info(f'Processing {data_set}')
    __process_data(
        data_folder,
        data_folder,
        num_processes=args.num_processes,
        rebalance=args.rebalance,
        class_split=args.class_split,
        skip_duration=args.skip_duration,
    )
    logging.info('Done!')


if __name__ == '__main__':
    main()