import os
import tensorflow as tf
import librosa
import argparse
from tools import MiscFns
import io
import numpy as np
import logging
import glob
import shutil


def split_train_valid_and_test_files_fn(args):
    """
    generate non-overlapped training-test set partition

    After downloading and unzipping the MAPS dataset,
    1. define an environment variable called maps to point to the directory of the MAPS dataset,
    2. populate test_dirs with the actual directories of the close and the ambient setting generated by
       the Disklavier piano,
    3. and populate train_dirs with the actual directoreis of the other 7 settings generated by the synthesizer.
    test_dirs = ['ENSTDkCl_2/MUS', 'ENSTDkAm_2/MUS']
    train_dirs = ['AkPnBcht_2/MUS', 'AkPnBsdf_2/MUS', 'AkPnCGdD_2/MUS', 'AkPnStgb_2/MUS',
                  'SptkBGAm_2/MUS', 'SptkBGCl_2/MUS', 'StbgTGd2_2/MUS']
    """
    '''
    put all test files in a directory, ".../maps/test"
    put all training files in another directory, ".../maps/train"
    '''
    test_dirs = ['test']
    train_dirs = ['train']
    maps_dir = args.data_root

    test_files = []
    for directory in test_dirs:
        path = os.path.join(maps_dir, directory)
        path = os.path.join(path, '*.wav')
        wav_files = glob.glob(path)
        test_files += wav_files

    test_ids = set([MiscFns.filename_to_id(wav_file) for wav_file in test_files])
    assert len(test_ids) == 53

    training_files = []
    validation_files = []
    for directory in train_dirs:
        path = os.path.join(maps_dir, directory)
        path = os.path.join(path, '*.wav')
        wav_files = glob.glob(path)
        for wav_file in wav_files:
            me_id = MiscFns.filename_to_id(wav_file)
            if me_id not in test_ids:
                training_files.append(wav_file)
            else:
                validation_files.append(wav_file)

    assert len(training_files) == 139 and len(test_files) == 60 and len(validation_files) == 71

    return dict(training=training_files, test=test_files, validation=validation_files)


def _dataset_iter_fn(name, args):
    """dataset generator"""
    assert name in ('validation', 'training', 'test')
    file_names = split_train_valid_and_test_files_fn(args)
    logging.debug('{} - enter generator'.format(name))

    if name == 'test' and args.test_with_30_secs:
        _duration = 30
    else:
        _duration = None
    logging.debug('{} - generate spectrograms and labels'.format(name))
    dataset = []
    for file_idx, wav_file_name in enumerate(file_names[name]):
        print('file_idx={}'.format(file_idx))
        logging.debug('{}/{} - {}'.format(
            file_idx + 1, len(file_names),
            os.path.basename(wav_file_name))
        )
        samples, unused_sr = librosa.load(
            mono=False, path=wav_file_name, sr=16000, duration=_duration, dtype=np.float32)
        assert unused_sr == 16000
        assert samples.shape[0] == 2
        spectrogram = []
        for ch in range(2):
            sg = MiscFns.spectrogram_fn(
                samples=samples[ch],
                log_filter_bank_basis=MiscFns.log_filter_bank_fn(),
                spec_stride=512
            )

            spectrogram.append(sg)

        spectrogram = np.stack(spectrogram, axis=-1)
        assert spectrogram.shape[1:] == (229, 2)

        mid_file_name = wav_file_name[:-4] + '.mid'
        # spectrogram.shape[0]: the number of frames
        label = MiscFns.label_fn(mid_file_name=mid_file_name, num_frames=spectrogram.shape[0], spec_stride=512)
        dataset.append([spectrogram, label])
        logging.debug('number of frames - {}'.format(spectrogram.shape[0]))

    rec_start_end_for_shuffle = []
    for rec_idx, rec_dict in enumerate(dataset):
        num_frames = len(rec_dict[0])
        split_frames = list(range(0, num_frames + 1, 900))
        if split_frames[-1] != num_frames:
            split_frames.append(num_frames)
        start_end_frame_pairs = zip(split_frames[:-1], split_frames[1:])
        rec_start_end_idx_list = [[rec_idx] + list(start_end_pair) for start_end_pair in start_end_frame_pairs]
        rec_start_end_for_shuffle += rec_start_end_idx_list

    if name=='training':
        np.random.shuffle(rec_start_end_for_shuffle)

    new_dataset = []
    for rec_idx, start_frame, end_frame in rec_start_end_for_shuffle:
        rec_dict = dataset[rec_idx]
        new_spectrogram=rec_dict[0][start_frame:end_frame]
        new_label=rec_dict[1][start_frame:end_frame]
        new_dataset.append([new_spectrogram, new_label, rec_idx])

    return new_dataset


def int64_feature(value):
    return tf.train.Feature(int64_list=tf.train.Int64List(value=[value]))

def int64_list_feature(value):
    return tf.train.Feature(int64_list=tf.train.Int64List(value=value))

def bytes_feature(value):
    return tf.train.Feature(bytes_list=tf.train.BytesList(value=[value]))

def bytes_list_feature(value):
    return tf.train.Feature(bytes_list=tf.train.BytesList(value=value))

def float_list_feature(value):
    return tf.train.Feature(float_list=tf.train.FloatList(value=value))

def create_tf_example(spectrogram, label, me_id):
    shape = np.array(spectrogram.shape, np.int32)
    assert shape[1] == 229
    assert shape[2] == 2
    shape_label = np.array(label.shape, np.int32)
    assert shape_label[1] == 88
    assert shape[0] == shape_label[0]

    spectrogram = spectrogram.tobytes()

    label = label.tobytes()

    tf_example = tf.train.Example(
        features=tf.train.Features(feature={
            'spectrogram': bytes_feature(spectrogram),
            'label': bytes_feature(label),
            'h': int64_feature(shape[0]),
            'me_id': int64_feature(me_id),
            'shape_spec': int64_list_feature(shape),
            'shape_label': int64_list_feature(shape_label)
        })
    )
    return tf_example

def generate_tfrecord(annotation_list, output_path):
    num_valid_tf_example = 0
    writer = tf.python_io.TFRecordWriter(output_path)
    print('length of annotation_list={}'.format(len(annotation_list)))
    for i in range(len(annotation_list)):
        print('make tfrecord:{}'.format(i))
        spectrogram = annotation_list[i][0]    #float32
        label = annotation_list[i][1]    #bool
        label_new = label.astype(np.int32)
        me_id = annotation_list[i][2]    #int64
        tf_example = create_tf_example(spectrogram, label_new, me_id)

        if tf_example:
            writer.write(tf_example.SerializeToString())
            num_valid_tf_example += 1

            if num_valid_tf_example % 100 == 0:
                print('Create %d TF_Example.' % num_valid_tf_example)
    writer.close()
    print('Total create TF_Example: %d' % num_valid_tf_example)


def make_file(args):
    print('make testing data')
    test_files = _dataset_iter_fn(name='test', args=args)
    print('generate_test_tfrecord')
    test_output_path = os.path.join(args.test_root, 'test.tfrecord')
    generate_tfrecord(test_files, output_path=test_output_path)

    print('make validation data')
    valid_files = _dataset_iter_fn(name='validation', args=args)
    print('generate_valid_tfrecord')
    valid_output_path = os.path.join(args.valid_root, 'valid.tfrecord')
    generate_tfrecord(valid_files, output_path=valid_output_path)


    print('make training data')
    train_files = _dataset_iter_fn(name='training', args=args)
    print('generate_train_tfrecord')
    train_output_path = os.path.join(args.train_root, 'train.tfrecord')
    generate_tfrecord(train_files, output_path=train_output_path)

    print('Make data finished!')


def set_up(args):
    args.train_root = os.path.join(args.save_dir, 'train')
    args.valid_root = os.path.join(args.save_dir, 'valid')
    args.test_root = os.path.join(args.save_dir, 'test')

    # make dir.
    if os.path.isdir(args.save_dir):
        logging.info('delete dir: {}'.format(args.save_dir))
        shutil.rmtree(args.save_dir)
    logging.info('mkdir: {}'.format(args.save_dir))
    os.makedirs(args.save_dir)

    if os.path.isdir(args.train_root):
        logging.info('delete dir: {}'.format(args.train_root))
        shutil.rmtree(args.train_root)
    logging.info('mkdir: {}'.format(args.train_root))
    os.makedirs(args.train_root)

    if os.path.isdir(args.valid_root):
        logging.info('delete dir: {}'.format(args.valid_root))
        shutil.rmtree(args.valid_root)
    logging.info('mkdir: {}'.format(args.valid_root))
    os.makedirs(args.valid_root)

    if os.path.isdir(args.test_root):
        logging.info('delete dir: {}'.format(args.test_root))
        shutil.rmtree(args.test_root)
    logging.info('mkdir: {}'.format(args.test_root))
    os.makedirs(args.test_root)

    return args

def get_args():
    parser = argparse.ArgumentParser(formatter_class=argparse.ArgumentDefaultsHelpFormatter)
    parser.add_argument('--data_root', type=str, default='/titan_data1/zhangwen/maps')
    parser.add_argument('--save_dir', type=str, default='/titan_data1/zhangwen/maps-tf')
    parser.add_argument('--gpus', type=str, default='2')
    parser.add_argument('--test_with_30_secs', type=bool, default=False)
    args = parser.parse_args()
    return args


if __name__ == '__main__':
    args = get_args()
    args = set_up(args)
    logging.info(args)
    os.environ["CUDA_VISIBLE_DEVICES"] = args.gpus

    make_file(args)