import pandas as pd
import numpy as np
from spectro_gen import generate_aiff_spectrogram
from sklearn.model_selection import train_test_split
import itertools
import gc


def get_siamese_pairs(groupF1, groupF2):
    group1 = [r for (i, r) in groupF1.iterrows()]
    group2 = [r for (i, r) in groupF2.iterrows()]
    f = [(g1, g2) for g2 in group2 for g1 in group1]
    t = [i for i in itertools.combinations(group1, 2)
         ] + [i for i in itertools.combinations(group2, 2)]
    return (t, f)


def append_zeros(spgr, max_samples):
    return np.lib.pad(spgr, [(0, max_samples - spgr.shape[0]), (0, 0)],
                      'median')


def create_pair(l, r, max_samples):
    l_sample = append_zeros(l, max_samples)
    r_sample = append_zeros(r, max_samples)
    return np.asarray([l_sample, r_sample])


def create_test_pair(l, r, max_samples):
    l_sample = append_zeros(l, max_samples)
    r_sample = append_zeros(r, max_samples)
    return np.asarray([[l_sample, r_sample]])


def create_X(sp, max_samples):
    return create_pair(sp[0]['spectrogram'], sp[1]['spectrogram'], max_samples)


def get_word_pairs_data(word, max_samples):
    audio_samples = pd.read_csv(
        './outputs/audio.csv',
        names=['word', 'voice', 'rate', 'variant', 'file'])
    audio_samples = audio_samples.loc[audio_samples['word'] ==
                                      word].reset_index(drop=True)
    audio_samples.loc[:, 'spectrogram'] = audio_samples.loc[:, 'file'].apply(
        lambda x: 'outputs/audio/' + x).apply(generate_aiff_spectrogram)
    # max_samples = audio_samples['spectrogram'].apply(
    #     lambda x: x.shape[0]).max()
    same_data, diff_data = [], []
    for (w, g) in audio_samples.groupby(audio_samples['word']):
        sample_norm = g.loc[audio_samples['variant'] == 'normal']
        sample_phon = g.loc[audio_samples['variant'] == 'phoneme']
        same, diff = get_siamese_pairs(sample_norm, sample_phon)
        same_data.extend([create_X(s, max_samples) for s in same[:10]])
        diff_data.extend([create_X(d, max_samples) for d in diff[:10]])
    Y = np.hstack([np.ones(len(same_data)), np.zeros(len(diff_data))])
    X = np.asarray(same_data + diff_data)
    # tr_pairs, te_pairs, tr_y, te_y = train_test_split(X, Y, test_size=0.1)
    return (X, Y)


def create_spectrogram_data(audio_group='audio'):
    audio_samples = pd.read_csv(
        './outputs/' + audio_group + '.csv',
        names=['word', 'voice', 'rate', 'variant', 'file'])
    # audio_samples = audio_samples.loc[audio_samples['word'] ==
    #                                   'sunflowers'].reset_index(drop=True)
    audio_samples.loc[:, 'spectrogram'] = audio_samples.loc[:, 'file'].apply(
        lambda x: 'outputs/' + audio_group + '/' + x).apply(
            generate_aiff_spectrogram)
    audio_samples.to_pickle('outputs/spectrogram.pkl')


def create_speech_pairs_data(audio_group='audio'):
    audio_samples = pd.read_pickle('outputs/spectrogram.pkl')
    max_samples = audio_samples['spectrogram'].apply(
        lambda x: x.shape[0]).max()
    # sample_size = audio_samples['spectrogram'][0].shape[1]

    print('generating siamese speech pairs')
    same_data, diff_data = [], []
    for (w, g) in audio_samples.groupby(audio_samples['word']):
        sample_norm = g.loc[audio_samples['variant'] == 'normal']
        sample_phon = g.loc[audio_samples['variant'] == 'phoneme']
        same, diff = get_siamese_pairs(sample_norm, sample_phon)
        same_data.extend([create_X(s, max_samples) for s in same[:10]])
        diff_data.extend([create_X(d, max_samples) for d in diff[:10]])
    print('creating all speech pairs')
    Y = np.hstack([np.ones(len(same_data)), np.zeros(len(diff_data))])
    print('casting as array speech pairs')
    X = np.asarray(same_data + diff_data)
    print('pickling X/Y')
    np.save('outputs/X.npy', X)
    np.save('outputs/Y.npy', Y)
    del same_data
    del diff_data
    gc.collect()
    print('train/test splitting speech pairs')
    tr_pairs, te_pairs, tr_y, te_y = train_test_split(X, Y, test_size=0.1)
    print('pickling train/test')
    np.save('outputs/tr_pairs.npy', tr_pairs)
    np.save('outputs/te_pairs.npy', te_pairs)
    np.save('outputs/tr_y.npy', tr_y)
    np.save('outputs/te_y.npy', te_y)


def speech_model_data():
    tr_pairs = np.load('outputs/tr_pairs.npy') / 255.0
    te_pairs = np.load('outputs/te_pairs.npy') / 255.0
    tr_pairs[tr_pairs < 0] = 0
    te_pairs[te_pairs < 0] = 0
    tr_y = np.load('outputs/tr_y.npy')
    te_y = np.load('outputs/te_y.npy')
    return tr_pairs, te_pairs, tr_y, te_y


if __name__ == '__main__':
    # sunflower_pairs_data()
    # create_spectrogram_data()
    create_speech_pairs_data()
    # print(speech_model_data())