import os
import pathlib
import sklearn.model_selection
import tarfile
import torch
import torchaudio
import urllib.request

here = pathlib.Path(__file__).resolve().parent


def _split_data(tensor, stratify):
    # 0.7/0.15/0.15 train/val/test split
    (train_tensor, testval_tensor,
     train_stratify, testval_stratify) = sklearn.model_selection.train_test_split(tensor, stratify,
                                                                                  train_size=0.7,
                                                                                  random_state=0,
                                                                                  shuffle=True,
                                                                                  stratify=stratify)

    val_tensor, test_tensor = sklearn.model_selection.train_test_split(testval_tensor,
                                                                       train_size=0.5,
                                                                       random_state=1,
                                                                       shuffle=True,
                                                                       stratify=testval_stratify)
    return train_tensor, val_tensor, test_tensor


def _save_data(dir, **tensors):
    for tensor_name, tensor_value in tensors.items():
        torch.save(tensor_value, str(dir / tensor_name) + '.pt')


def download():
    base_base_loc = str(here / '../experiments/data')
    if not os.path.exists(base_base_loc):
        raise RuntimeError("data directory does not exist. Please create a directory called 'data' in the 'experiments'"
                           " directory. (We're going to put a lot of data there, so we don't make it automatically - "
                           "thus giving you the opportunity to make it a symlink rather than a normal directory, so "
                           "that the data can be stored elsewhere if you wish.)")
    base_loc = base_base_loc + '/SpeechCommands'
    loc = base_loc + '/speech_commands.tar.gz'
    if os.path.exists(loc):
        return
    if not os.path.exists(base_loc):
        os.mkdir(base_loc)
    urllib.request.urlretrieve('http://download.tensorflow.org/data/speech_commands_v0.02.tar.gz',
                               loc)

    with tarfile.open(loc, 'r') as f:
        f.extractall(base_loc)


def _process_data():
    base_loc = here / '..' / 'experiments' / 'data' / 'SpeechCommands'
    X = torch.empty(34975, 16000, 1)
    y = torch.empty(34975, dtype=torch.long)

    batch_index = 0
    y_index = 0
    for foldername in ('yes', 'no', 'up', 'down', 'left', 'right', 'on', 'off', 'stop', 'go'):
        loc = base_loc / foldername
        for filename in os.listdir(loc):
            audio, _ = torchaudio.load_wav(loc / filename, channels_first=False,
                                           normalization=False)  # for forward compatbility if they fix it
            audio = audio / 2 ** 15  # Normalization argument doesn't seem to work so we do it manually.

            # A few samples are shorter than the full length; for simplicity we discard them.
            if len(audio) != 16000:
                continue

            X[batch_index] = audio
            y[batch_index] = y_index
            batch_index += 1
        y_index += 1
    assert batch_index == 34975, "batch_index is {}".format(batch_index)

    audio_X = X

    # X is of shape (batch=34975, length=16000, channels=1)
    X = torchaudio.transforms.MFCC(log_mels=True)(X.squeeze(-1)).transpose(1, 2).detach()
    # X is of shape (batch=34975, length=81, channels=40). For some crazy reason it requires a gradient, so detach.

    train_X, _, _ = _split_data(X, y)
    out = []
    means = []
    stds = []
    for Xi, train_Xi in zip(X.unbind(dim=-1), train_X.unbind(dim=-1)):
        mean = train_Xi.mean()
        std = train_Xi.std()
        means.append(mean)
        stds.append(std)
        out.append((Xi - mean) / (std + 1e-5))
    X = torch.stack(out, dim=-1)

    train_audio_X, val_audio_X, test_audio_X = _split_data(audio_X, y)
    train_X, val_X, test_X = _split_data(X, y)
    train_y, val_y, test_y = _split_data(y, y)

    return train_X, val_X, test_X, train_y, val_y, test_y, torch.stack(means), torch.stack(stds), train_audio_X, \
           val_audio_X, test_audio_X


def main():
    download()
    (train_X, val_X, test_X, train_y, val_y, test_y, means, stds, train_audio_X, val_audio_X,
     test_audio_X) = _process_data()
    loc = here / '..' / 'experiments' / 'data' / 'speech_commands_data'
    if not os.path.exists(loc):
        os.mkdir(loc)
    _save_data(loc, train_X=train_X, val_X=val_X, test_X=test_X, train_y=train_y, val_y=val_y, test_y=test_y,
               means=means, stds=stds, train_audio_X=train_audio_X, val_audio_X=val_audio_X, test_audio_X=test_audio_X)


if __name__ == '__main__':
    main()

[((43, 7, 43, 26), 'os.path.exists', 'os.path.exists', ({(43, 22, 43, 25): 'loc'}, {}), '(loc)', False, 'import os\n'), ((56, 8, 56, 36), 'torch.empty', 'torch.empty', ({(56, 20, 56, 25): '34975', (56, 27, 56, 32): '16000', (56, 34, 56, 35): '1'}, {}), '(34975, 16000, 1)', False, 'import torch\n'), ((57, 8, 57, 44), 'torch.empty', 'torch.empty', (), '', False, 'import torch\n'), ((94, 8, 94, 32), 'torch.stack', 'torch.stack', (), '', False, 'import torch\n'), ((36, 11, 36, 40), 'os.path.exists', 'os.path.exists', ({(36, 26, 36, 39): 'base_base_loc'}, {}), '(base_base_loc)', False, 'import os\n'), ((45, 11, 45, 35), 'os.path.exists', 'os.path.exists', ({(45, 26, 45, 34): 'base_loc'}, {}), '(base_loc)', False, 'import os\n'), ((46, 8, 46, 26), 'os.mkdir', 'os.mkdir', ({(46, 17, 46, 25): 'base_loc'}, {}), '(base_loc)', False, 'import os\n'), ((50, 9, 50, 31), 'tarfile.open', 'tarfile.open', ({(50, 22, 50, 25): 'loc', (50, 27, 50, 30): '"""r"""'}, {}), "(loc, 'r')", False, 'import tarfile\n'), ((63, 24, 63, 39), 'os.listdir', 'os.listdir', ({(63, 35, 63, 38): 'loc'}, {}), '(loc)', False, 'import os\n'), ((100, 59, 100, 77), 'torch.stack', 'torch.stack', ({(100, 71, 100, 76): 'means'}, {}), '(means)', False, 'import torch\n'), ((100, 79, 100, 96), 'torch.stack', 'torch.stack', ({(100, 91, 100, 95): 'stds'}, {}), '(stds)', False, 'import torch\n'), ((109, 11, 109, 30), 'os.path.exists', 'os.path.exists', ({(109, 26, 109, 29): 'loc'}, {}), '(loc)', False, 'import os\n'), ((110, 8, 110, 21), 'os.mkdir', 'os.mkdir', ({(110, 17, 110, 20): 'loc'}, {}), '(loc)', False, 'import os\n'), ((9, 7, 9, 29), 'pathlib.Path', 'pathlib.Path', ({(9, 20, 9, 28): '__file__'}, {}), '(__file__)', False, 'import pathlib\n'), ((64, 23, 65, 63), 'torchaudio.load_wav', 'torchaudio.load_wav', (), '', False, 'import torchaudio\n'), ((81, 8, 81, 49), 'torchaudio.transforms.MFCC', 'torchaudio.transforms.MFCC', (), '', False, 'import torchaudio\n')]