<file_sep>import numpy as np
from numpy.lib.stride_tricks import as_strided
from audio_utils import AudioSegment
from scipy.io.wavfile import read

EPS = 1e-5


def specgram_real(data, NFFT=256, Fs=2, noverlap=128):
    """
    Computes the spectrogram for a real signal.

    The parameters follow the naming convention of
    matplotlib.mlab.specgram

    Params:
    data     : 1D numpy array
    NFFT     : number of elements in the window
    Fs       : sample rate
    noverlap : number of elements to overlap each window

    Returns:
    x        : 2D numpy array, frequency x time
    freq     : 1D array of frequency of each row in x

    Note this is a truncating computation e.g. if NFFT=10,
    noverlap=5 and the signal has 23 elements, then the
    last 3 elements will be truncated.
    """
    assert not np.iscomplexobj(data), "Must not pass in complex numbers"

    stride = NFFT - noverlap
    window = np.hanning(NFFT)[:, None]
    window_norm = np.sum(window**2)

    # The scaling below follows the convention of
    # matplotlib.mlab.specgram which is the same as
    # matlabs specgram.  Seems somewhat arbitrary thus
    # TODO, awni, check on the validity of this scaling.
    scale = window_norm * Fs

    trunc = (len(data) - NFFT) % stride
    x = data[:len(data) - trunc]

    # "stride trick" reshape to include overlap
    nshape = (NFFT, (len(x) - NFFT) // stride + 1)
    nstrides = (x.strides[0], x.strides[0] * stride)
    x = as_strided(x, shape=nshape, strides=nstrides)

    # broadcast window, compute fft over columns and square mod
    x = np.fft.rfft(x * window, axis=0)
    x = np.absolute(x)**2

    # scale, 2.0 for everything except dc and nfft/2
    x[1:-1, :] *= (2.0 / scale)
    x[(0, -1), :] /= scale

    freqs = float(Fs) / NFFT * np.arange(x.shape[0])

    return (x, freqs)


def get_feat_dim(step, window, max_freq, **kwargs):
    # TODO sanjeev Figure this out!
    feat_dim_map = {(10, 20, 8000): 81, (10, 20, 16000): 161}
    key = (step, window, max_freq)
    assert key in feat_dim_map
    return feat_dim_map[key]


def compute_volume_normalized_feature(audio_file, feature_info):
    audio = AudioSegment.from_wav(audio_file)
    audio = audio.normalize()
    assert audio.normalized, "Audio segment was not normalized"
    return compute_raw_feature(audio.wav_data, feature_info)


def compute_raw_feature(audio_file, feature_info):
    """
    Compute time x frequency feature matrix *without* any normalization,
    striding, or splicing. We *do* convert to log-space.
    """
    # spec is frequency x time
    spec, _ = compute_specgram(
        audio_file,
        step=feature_info['step'],
        window=feature_info['window'],
        max_freq=feature_info['max_freq'])

    # transpose to time x frequency
    return np.log(spec + EPS).astype(np.float32).T


def compute_specgram(audio_file, step=10, window=20, max_freq=8000, shift=0):
    """
    audio_file : audio wav file
    step : step size in ms between windows
    window : window size in ms
    max_freq : keep frequencies below value
    shift: time in ms to be clipped form the audio data,
           use negative values to clip form the end

    Uses default window for specgram (hamming).
    Return value is Frequency x Time matrix.
    """

    Fs, data = read(audio_file)
    return compute_specgram_raw(Fs, data, step, window, max_freq, shift)


def compute_specgram_raw(Fs, data, step, window, max_freq,
                         shift, spec_fn=specgram_real):
    """
    Same as above when you've already read audio_file
    """

    shift = (Fs/1000) * shift
    if shift < 0:
        data[-shift:] = data[:shift]
    elif shift > 0:
        data[:-shift] = data[shift:]

    # If we have two channels, take just the first
    if len(data.shape) == 2:
        data = data[:, 0].squeeze()

    step_sec = step/1000.
    window_sec = window/1000.

    if step_sec > window_sec:
        raise ValueError("Step size must not be greater than window size")

    noverlap = int((window_sec - step_sec) * Fs)
    NFFT = int(window_sec * Fs)
    Pxx, freqs = spec_fn(data, NFFT=NFFT, Fs=Fs, noverlap=noverlap)[0:2]
    try:
        ind = np.where(freqs == max_freq/2)[0][0] + 1
    except IndexError:
        return None, None
    return Pxx[:ind, :], freqs[:ind]
<file_sep>''' Collection of classes to help with finding and adding noise to audio
'''

import os
import wave
import pydub
import random
import string
import subprocess
import numpy as np
import glob
from scipy import stats
from cStringIO import StringIO
# from utils import parmap
import scikits.audiolab as ab
import linecache
import re
import json

MSEC = 1000 # pydub frames are in ms, so multiply with MSEC when indexing
FRAME_RATE = 16000
DEVNULL = open(os.devnull, 'w')
ROOT_DIR = '/scratch/sanjeev/random-youtube-1/'
TMP_DIR = '/scratch/sanjeev/tmp/'

class AudioSegment(pydub.AudioSegment):
    ''' Helper class to manipulate `AudioSegment`s
        Some functions may use sox
    '''
    sox_cmd = 'sox -t wav - -t wav -'
    effects = ['echo', 'chorus', 'speed', 'pitch', 'flanger', 'reverb', 'tempo', 'overdrive']

    @property
    def power_dB(self):
        ''' power of the segment in decibels
        '''
        if self.rms == 0:
            return 0
        return pydub.utils.ratio_to_db(self.rms)

    def normalize(self, required_power_in_dB=70):
        ''' Performs a best effort power normalization
            If the power could not be set to the required_power_in_dB
            it returns the altered audio without failing!
        '''
        delta = required_power_in_dB - self.power_dB
        self.normalize_gap = delta
        self = self + delta
        self = self.set_frame_rate(FRAME_RATE)
        self.__class__ = AudioSegment
        self.normalize_gap = abs(self.power_dB - required_power_in_dB) < 5
        return self

    @property
    def wav_data(self):
        assert self.frame_rate == FRAME_RATE, "Frame rate set incorrectly."
        assert hasattr(self, 'normalize_gap'), "Normalize audio before accessing wav_data."
        wav_file = StringIO()
        self.export(wav_file, format='wav')
        wav_file.seek(0)
        return wav_file

    @classmethod
    def get_random_effect(cls):
        return cls.effects[random.randint(0, len(cls.effects)-1)]

    def add_effect(self, effect, params=[]):
        ''' Uses sox to add one of these effects
            Read AudioSegment.effects to see the list of effects available
        '''
        # Write the current file to console
        assert self.duration_seconds < 100, '{0:.2f}'.format(self.duration_seconds)
        original_file = StringIO()
        self.export(original_file, format='wav')
        original_file.seek(0)
        wav_data = original_file.read()

        # Add the effect using sox
        try:
            effect_cmd = AudioSegment._get_effect_param(effect, params)
            sox_cmd = '{0} {1}'.format(self.sox_cmd, effect_cmd)
            process = subprocess.Popen(sox_cmd, stdout=subprocess.PIPE,
                                       stdin=subprocess.PIPE, stderr=DEVNULL,
                                       shell=True)
            effected_file = StringIO(process.communicate(wav_data)[0])
            effected_file.seek(0)

            # Return the updated one
            return AudioSegment.from_wav(effected_file)
        except Exception as msg:
            print msg, effect
            return self

    @classmethod
    def _get_effect_param(cls, effect, params=[]):
        # Make sure params is a list
        if not hasattr(params, '__getitem__'): params = [params]


        if effect == 'echo':
            delay = random.randint(50, 350)
            decay = random.randint(1, 7) * 0.1
            return 'echo 1 1 {0} {1:.2}'.format(delay, decay)

        elif effect == 'chorus':
            opts = 'chorus 1 1 '
            num_people = random.randint(1, 5)
            for i in range(num_people):
                delay = random.randint(21, 90)
                decay = random.randint(1, 7) * 0.1
                speed = random.random() + 0.5
                depth = random.random() + 1.5
                opts += ' {0} {1:.2} {2:.2} {3:.2} -s' \
                        .format(delay, decay, speed, depth)
            return opts

        elif effect == 'flanger':
            delay = random.randint(0, 30)
            depth = random.randint(0, 10)
            regen = random.randint(-95, 95)
            width = random.randint(60, 100)
            speed = random.random() + 0.1
            phase = random.randint(0, 100)
            return 'flanger {0} {1} {2} {3} {4:.2} sin {5} lin' \
                    .format(delay, depth, regen, width, speed, phase)

        elif effect == 'reverb':
            return 'reverb '

        elif effect == 'tempo':
            amount = params[0] if len(params) > 0 else random.random() + 0.5
            return 'tempo -s {0:.2}'.format(amount)

        elif effect == 'speed':
            amount = params[0] if len(params) > 0 else random.random() + 0.5
            return 'speed {0:.2}'.format(amount)

        elif effect == 'pitch':
            amount = params[0] if len(params) > 0 else random.randint(-400, 400)
            return 'pitch {0}'.format(amount)

        elif effect == 'overdrive':
            amount = params[0] if len(params) > 0 else 20*abs(random.random())
            return 'overdrive {0:.2}'.format(amount)


        return ''

    def sample_clip(self, required_duration):
        ''' Randomly sample a clip of `required_duration` seconds from self
        '''
        self_duration = int(self.duration_seconds * MSEC) # in ms
        required_duration = int(required_duration * MSEC) # in ms
        assert self_duration >= 20 + required_duration, 'Cannot sample!'
        start_time = random.randint(10, self_duration-required_duration-10)
        clip = self[start_time:start_time+required_duration]
        clip.__class__ = AudioSegment
        return clip

    def add_noise(self, noise, snr_dB, final_duration=None, effect=None):
        ''' Adds the given noise signal such that the snr_dB is maintained
        '''
        signal_duration = int(self.duration_seconds * MSEC) # in ms
        noise_duration = int(noise.duration_seconds * MSEC) # in ms
        assert noise_duration >= signal_duration, 'Noise smaller than signal'
        if effect:
            noise = noise.add_effect(effect)
        if final_duration:
            noise = noise.sample_clip(final_duration or self.duration_seconds)
        delta = self.power_dB - noise.power_dB - snr_dB
        noise_clip = noise + delta
        merged = noise_clip * self
        merged.__class__ = AudioSegment
        return merged

    def make_segments(self, prefix, duration=60, outdir=TMP_DIR):
        ''' Saves consecutive `duration` seconds long clips from self into
            outdir with names of form `prefix.{starting_second}.wav`
        '''
        split_files = []
        duration = int(duration * MSEC)
        for start_pos in range(0, len(self), duration):
            seg_x = self[start_pos:start_pos + duration]
            if seg_x.duration_seconds * MSEC >= duration:
                basename = '{0}.{1}.wav'.format(prefix, start_pos)
                fname = os.path.join(outdir, basename)
                seg_x.set_channels(1)
                seg_x.set_frame_rate(FRAME_RATE)
                seg_x.export(fname, format='wav')
                split_files.append((basename, duration))
        return split_files

class Youtube(object):
    ''' Pulls and converts audio from youtube '''

    url_alpha = list(string.ascii_letters + string.digits)
    base_url = 'https://www.youtube.com/watch?v=%s'

    ydl_cmd = ['youtube-dl', '-x', '--cache-dir', TMP_DIR, '--output']

    avconv_cmd = 'avconv -y -f {2} -i {0} -vn -b 64k -ar {3} -ac 1 -f flac {1}'
    ydl_filename_format = '%(id)s.%(ext)s'

    def __init__(self, yt_id, deloop=False):
        ''' Initializing this class downloads a youtube file in wav format
            into `ROOT_DIR`
        '''
        self.yt_id = yt_id
        url = self.base_url % self.yt_id
        audio_fname = os.path.join(TMP_DIR, self.ydl_filename_format)
        ydl_params = self.ydl_cmd + [audio_fname, url]
        tmp_file = glob.glob(audio_fname % {'id': yt_id, 'ext' : '*'})
        final_fname = os.path.join(ROOT_DIR, self.ydl_filename_format)
        sound_file = final_fname % {'id': yt_id, 'ext': 'wav'}
        if os.path.exists(sound_file):
            print sound_file, 'exists'
            return

        if len(tmp_file) == 0:
            status = subprocess.call(ydl_params, stdin=subprocess.PIPE,
                                  stdout=subprocess.PIPE, stderr=DEVNULL)
            if status != 0:
                print 'downloading ', yt_id, ' failed'
                return

        tmp_file = glob.glob(audio_fname % {'id': yt_id, 'ext' : '*'})[0]
        inpformat = tmp_file[tmp_file.rindex('.')+1:]

        if deloop:
            # The headers are wrong in long files, use flac
            # for all processing on large files
            sound_file = self._flac_encode(tmp_file)
            self.audio = self.get_loopless_segment(sound_file)
        else:
            self.audio = AudioSegment.from_file(tmp_file, format=inpformat)

        print 'saving to', sound_file
        self.audio.set_channels(1)
        self.audio.set_frame_rate(FRAME_RATE)
        # TODO Best if the bit rate could also be set to 64k,
        # usually gives about 2x savings in disk space consumed
        self.audio.export(sound_file, format='wav')

    def _random_yt_id(self):
        ''' Get a random youtube video id '''
        return random.sample(self.url_alpha, 11)

    def _flac_encode(self, inpfile):
        ''' Youtube sounds are downloaded as m4a files. These cannot be loaded
            by scikit.audiolab. Wav files do not support files > 2G,
            so we temporarily convert to a flac format with smaller sample and
            bit rate
        '''
        inpformat = inpfile[inpfile.rindex('.')+1:]
        flacfile = inpfile.replace(inpformat, 'flac')
        conv_cmd = self.avconv_cmd.format(inpfile, flacfile, inpformat, FRAME_RATE)
        subprocess.call(conv_cmd.split(' '), stdin=subprocess.PIPE,
                        stdout=subprocess.PIPE, stderr=DEVNULL)
        os.remove(inpfile)
        return flacfile

    def get_loopless_segment(self, inpfile):
        inpformat = inpfile[inpfile.rindex('.')+1:]
        corrfname = inpfile.replace(inpformat, 'corr.npy')
        if not os.path.exists(corrfname):
            print corrfname, 'not found'
            source = ab.Sndfile(inpfile, 'r')
            source_data = source.read_frames(source.nframes)
            num_frames = source_data.shape[0]
            clip_size = int(0.1 * FRAME_RATE)
            max_start_pos = FRAME_RATE * 10
            start_pos = random.randint(0, num_frames-clip_size-10)
            start_pos = min(max_start_pos, start_pos)
            clip_data = source_data[start_pos:start_pos+clip_size]
            corr = norm_correlation(clip_data, source_data)
            # Why are there duplicates!
            with open(corrfname, 'w') as save_file:
                np.save(save_file, corr)
        else:
            corr = np.load(corrfname)

        reps = np.sort(np.unique(np.where(corr > 0.98)[0]))
        gaps = np.array([])

        if reps.shape[0] > 2:
            gaps = reps[1:] - reps[:-1]
            gaps = gaps[np.unique(np.where(gaps > 10.)[0])]

        audio = AudioSegment.from_file(inpfile, inpformat)
        os.remove(inpfile)
        file_duration = audio.duration_seconds / 60.
        if gaps.shape[0] > 0:
            loop_duration = stats.mode(gaps)[0][0]/FRAME_RATE
            audio = audio[0:loop_duration * MSEC]
            print '{0} {1:.2f} {2:.2f}'.format(inpfile, file_duration,
                                               audio.duration_seconds / 60.)
        return audio

def norm_correlation(t, s):
    ''' Fast normalized correlation function for matching template `t` against
        the source `s`
        Refer http://stackoverflow.com/questions/23705107
    '''
    import pandas as pd
    n = len(t)
    nt = (t-np.mean(t))/(np.std(t)*n)
    sum_nt = nt.sum()
    a = pd.rolling_mean(s, n)[n-1:-1]
    b = pd.rolling_std(s, n)[n-1:-1]
    b *= np.sqrt((n-1.0) / n)
    c = np.convolve(nt[::-1], s, mode="valid")[:-1]
    result = (c - sum_nt * a) / b
    return result

def get_pcm_audio(pcmdata, nchannel=1, samplerate=8000, nbit=16):
    ''' Converts pcm data to an AudioSegment '''
    wavbuffer = StringIO()
    wavfile = wave.open(wavbuffer, 'wb')
    wavfile.setparams((nchannel, nbit/8, samplerate, 0, 'NONE', 'NONE'))
    wavfile.writeframes(pcmdata)
    wavfile.close()

    wavbuffer.seek(0)
    return AudioSegment.from_wav(wavbuffer)


current_seq_filename = "" #Need filename to know when to clear cache
def getline(seq_filename, line_num):
    '''Gets line of text.  Clears linecache if file changes'''
    global current_seq_filename

    if seq_filename != current_seq_filename:
        linecache.clearcache()
        current_seq_filename = seq_filename

    return linecache.getline(seq_filename, line_num)

def get_raw_audio(fname):
    ''' Loads audio segment from file. Load method depends on file extension '''

    audio_fmt = fname[1+fname.rindex('.'):]
    if audio_fmt == 'pcm':
        with open(fname, 'rb') as pcmfile:
            pcmdata = pcmfile.read()
        audio = get_pcm_audio(pcmdata)

    elif re.match("""seq_\d+""", audio_fmt):
        pos = fname.rindex('_')
        seq_filename, line_num = fname[:pos], int(fname[1+pos:])

        data_line = getline(seq_filename, line_num)
        audio_data = json.loads(data_line)["audio_data"].decode('base64')

        try:
            # Assume this is a wav file
            audio = AudioSegment(data=audio_data)
        except Exception as e:
            # if not, its pcm. We only have these two.
            audio = get_pcm_audio(audio_data)
    else:
        audio = AudioSegment.from_file(fname, format=audio_fmt)

    return audio

def get_wav_duration(wavfile):
    return AudioSegment.from_wav(wavfile).duration_seconds

def download_sounds(ytid):
    try:
        ytid = ytid.strip()
        if ytid:
            Youtube(ytid)
    except Exception as error:
        print ytid, error

# if __name__ == '__main__':
#     with open('ytids', 'r') as yt:
#         ids = yt.readlines()
#     parmap(download_sounds, ids)<file_sep>from __future__ import print_function
import os
import struct
import random
import math
import numpy
import theano.tensor as T
from opendeep.data import Dataset, TRAIN, VALID, TEST
from opendeep.log import config_root_logger
from opendeep.models.single_layer.lstm import LSTM
from opendeep.optimization import RMSProp
from opendeep.monitor import Monitor, Plot
from opendeep import dataset_shared

basedir = 'data/ecg/apnea-ecg'
label_ext = ".apn"
data_ext = ".dat"
extra_ext = ".qrs"


def find_train_files(directory, find_ext):
    for root, dirs, files in os.walk(directory):
        for basename in files:
            name, ext = os.path.splitext(basename)
            if ext == find_ext and name.split(os.sep)[-1][0] in ['a', 'b', 'c'] and len(name.split(os.sep)[-1]) == 3:
                filename = os.path.join(root, basename)
                yield filename

class ECG(Dataset):
    def __init__(self, source_dir, train_split=0.8, valid_split=0.15):
        self.train = None, None
        self.train_shape = None
        self.valid = None, None
        self.valid_shape = None
        self.test = None, None
        self.test_shape = None


    def getSubset(self, subset):
        if subset is TRAIN:
            return self.train
        elif subset is VALID:
            return self.valid
        elif subset is TEST:
            return self.test
        else:
            return None, None

    def getDataShape(self, subset):
        if subset is TRAIN:
            return self.train_shape
        elif subset is VALID:
            return self.valid_shape
        elif subset is TEST:
            return self.test_shape
        else:
            return None


def main():
    pass



if __name__ == '__main__':
    # http://www.physionet.org/physiobank/database/apnea-ecg/

    # if we want logging
    config_root_logger()

    i=0
    for f in find_train_files(basedir, label_ext):
        if i==0:
            data = numpy.fromfile(f, dtype=numpy.bool)
            print(data.shape)
        else:
            pass
        i+=1

    i=0
    for f in find_train_files(basedir, data_ext):
        if i == 0:
            data = numpy.fromfile(f, dtype=numpy.float16)
            print(data.shape)
        else:
            pass
        i += 1

    i = 0
    for f in find_train_files(basedir, extra_ext):
        if i == 0:
            data = numpy.fromfile(f, dtype=numpy.bool)
            print(data.shape)
        else:
            pass
        i += 1<file_sep>import os
import numpy
from feature_helpers import compute_specgram

audio_ext = ['.wav']
basedir = 'data/sounds/'

def main(steps, windows, max_freqs):
    # create the datasets!
    # go through everything in the data/sounds directory
    for f in find_audio_files(basedir):
        name, ext = os.path.splitext(f)
        for max_freq in max_freqs:
            for window in windows:
                for step in steps:
                    processed_name = name + '_processed_%d_%d_%d.npy'%(step, window, max_freq)
                    # if the current processed version doesn't exist
                    if not os.path.exists(processed_name):
                        # process the file!
                        processed, _freqs = compute_specgram(audio_file=f, step=step, window=window, max_freq=max_freq)
                        if processed is not None:
                            print f, '--->', processed_name
                            processed = processed.transpose()
                            print processed.shape
                            numpy.save(processed_name, processed)

def find_audio_files(directory):
    for root, dirs, files in os.walk(directory):
        for basename in files:
            _, ext = os.path.splitext(basename)
            if ext in audio_ext:
                filename = os.path.join(root, basename)
                yield filename

if __name__ == '__main__':
    freqs = [2000, 4000, 8000]
    windows = [10, 20]
    steps = [10]
    main(steps, windows, freqs)<file_sep>from __future__ import print_function
import os
import random
import math
import numpy
import theano.tensor as T
from opendeep.data import Dataset, TRAIN, VALID, TEST
from opendeep.log import config_root_logger
from opendeep.models import Prototype, SoftmaxLayer
from opendeep.models.single_layer.lstm import LSTM
from opendeep.optimization import RMSProp
from opendeep.monitor import Monitor, Plot
from opendeep.utils.misc import numpy_one_hot
from opendeep import dataset_shared

basedir = 'data/sounds/'

sizes = {'10_10_2000': 11,
         '10_20_2000': 21,
         '10_10_4000': 21,
         '10_20_4000': 41,
         '10_10_8000': 41,
         '10_20_8000': 81}

classes = ['normal', 'murmer', 'extrahls', 'artifact', 'extrastole']

def find_processed_files(directory, size_key):
    for root, dirs, files in os.walk(directory):
        for basename in files:
            name, ext = os.path.splitext(basename)
            if ext == ".npy" and "processed_%s" % size_key in name:
                filename = os.path.join(root, basename)
                yield filename

def get_label(filename):
    for i, c in enumerate(classes):
        if c in filename:
            return i
    return None

class HeartSound(Dataset):
    def __init__(self, source_dir, size_key, one_hot=False, train_split=0.9, valid_split=0.1):
        print("Getting dataset %s" % size_key)
        # grab the datasets from the preprocessed files
        datasets = [(numpy.load(f), get_label(f)) for f in find_processed_files(source_dir, size_key)]
        # make sure they are all the correct dimensionality
        datasets = [(data.shape, data, label) for data, label in datasets
                    if data.shape[1] == sizes[size_key] and label is not None]

        print("Found %d examples" % len(datasets))

        # shuffle!
        random.seed(1)
        random.shuffle(datasets)

        # shapes
        shapes = [shape for shape, _, _ in datasets]
        # data
        dataset = [data for _, data, _ in datasets]
        # labels
        labels = numpy.asarray([label for _, _, label in datasets], dtype='int8')
        # make the labels into one-hot vectors
        if one_hot:
            labels = numpy_one_hot(labels, n_classes=5)

        train_len = int(math.floor(train_split * len(dataset)))
        valid_len = int(math.floor(valid_split * len(dataset)))
        valid_stop = train_len + valid_len
        print("# train: %d examples" % train_len)

        train_datasets = dataset[:train_len]
        train_labels = labels[:train_len]

        if valid_len > 0:
            valid_datasets = dataset[train_len:valid_stop]
            valid_labels = labels[train_len:valid_stop]
        else:
            valid_datasets = []

        if train_len + valid_len < len(dataset):
            test_datasets = dataset[valid_stop:]
            test_labels = labels[valid_stop:]
        else:
            test_datasets = []

        # median_train_len = int(numpy.median(numpy.asarray(shapes[:train_len]), axis=0)[0])
        min_train_len = int(numpy.min(numpy.asarray(shapes[:train_len]), axis=0)[0])
        max_train_len = int(numpy.max(numpy.asarray(shapes[:train_len]), axis=0)[0])
        # train = numpy.array([data[:min_train_len] for data in train_datasets], dtype='float32')
        train = numpy.array(
            [numpy.pad(data, [(0, max_train_len-data.shape[0]), (0, 0)], mode='constant') for data in train_datasets],
            dtype='float32'
        )
        self.train_shape = train.shape
        self.train = (dataset_shared(train, borrow=True), dataset_shared(train_labels, borrow=True))

        if len(valid_datasets) > 0:
            min_valid_len = int(numpy.min(numpy.asarray(shapes[train_len:valid_stop]), axis=0)[0])
            max_valid_len = int(numpy.max(numpy.asarray(shapes[train_len:valid_stop]), axis=0)[0])
            # valid = numpy.array([data[:min_valid_len] for data in valid_datasets], dtype='float32')
            valid = numpy.array(
                [numpy.pad(data, [(0, max_valid_len-data.shape[0]), (0, 0)], mode='constant') for data in
                 valid_datasets],
                dtype='float32'
            )
            self.valid_shape = valid.shape
            self.valid = (dataset_shared(valid, borrow=True), dataset_shared(valid_labels, borrow=True))
        else:
            self.valid = None, None
            self.valid_shape = None

        if len(test_datasets) > 0:
            min_test_len = int(numpy.min(numpy.asarray(shapes[valid_stop:]), axis=0)[0])
            max_test_len = int(numpy.max(numpy.asarray(shapes[valid_stop:]), axis=0)[0])
            # valid = numpy.array([data[:min_test_len] for data in test_datasets], dtype='float32')
            test = numpy.array(
                [numpy.pad(data, [(0, max_test_len - data.shape[0]), (0, 0)], mode='constant') for data in
                 test_datasets],
                dtype='float32'
            )
            self.test_shape = test.shape
            self.test = (dataset_shared(test, borrow=True), dataset_shared(test_labels, borrow=True))
        else:
            self.test = None, None
            self.test_shape = None


        print("Train shape: %s" % str(self.train_shape))
        print("Valid shape: %s" % str(self.valid_shape))
        print("Test shape: %s" % str(self.test_shape))
        print("Dataset %s initialized!" % size_key)


    def getSubset(self, subset):
        if subset is TRAIN:
            return self.train
        elif subset is VALID:
            return self.valid
        elif subset is TEST:
            return self.test
        else:
            return None, None

    def getDataShape(self, subset):
        if subset is TRAIN:
            return self.train_shape
        elif subset is VALID:
            return self.valid_shape
        elif subset is TEST:
            return self.test_shape
        else:
            return None


def main(size_key):

    out_vector = True

    # grab the data for this step size, window, and max frequency
    heartbeats = HeartSound(basedir, size_key, one_hot=out_vector)

    # define our model! we are using lstm with mean-pooling and softmax as classification
    hidden_size = int(max(math.floor(sizes[size_key]*.4), 1))
    n_classes = 5

    lstm_layer = LSTM(input_size=sizes[size_key],
                      hidden_size=hidden_size,
                      output_size=1,  # don't care about output size
                      hidden_activation='tanh',
                      inner_hidden_activation='sigmoid',
                      weights_init='uniform',
                      weights_interval='montreal',
                      r_weights_init='orthogonal',
                      clip_recurrent_grads=5.,
                      noise='dropout',
                      noise_level=0.6,
                      direction='bidirectional')

    # mean of the hiddens across timesteps (reduces ndim by 1)
    mean_pooling = T.mean(lstm_layer.get_hiddens(), axis=0)

    # last hiddens as an alternative to mean pooling over time
    last_hiddens = lstm_layer.get_hiddens()[-1]

    # now the classification layer
    softmax_layer = SoftmaxLayer(inputs_hook=(hidden_size, mean_pooling),
                                 output_size=n_classes,
                                 out_as_probs=out_vector)

    # make it into a prototype!
    model = Prototype(layers=[lstm_layer, softmax_layer], outdir='outputs/prototype%s' % size_key)

    # optimizer
    optimizer = RMSProp(dataset=heartbeats,
                        model=model,
                        n_epoch=500,
                        batch_size=10,
                        save_frequency=10,
                        learning_rate=1e-4, #1e-6
                        grad_clip=None,
                        hard_clip=False
                        )

    # monitors
    errors = Monitor(name='error', expression=model.get_monitors()['softmax_error'], train=True, valid=True, test=True)

    # plot the monitor
    plot = Plot('heartbeat %s' % size_key, monitor_channels=[errors], open_browser=True)

    optimizer.train(plot=plot)



if __name__ == '__main__':
    # http://www.peterjbentley.com/heartchallenge/

    # if we want logging
    config_root_logger()

    for step in [10]:
        for freq in [8000]:
            for window in [20, 10]:
                size_key = "%d_%d_%d" % (step, window, freq)
                main(size_key)
