from __future__ import print_function

import cPickle as pickle

import glob

import numpy as np

import os

import subprocess

WFDB = "/deep/group/med/tools/wfdb-10.5.24/build/bin/"

DATA = "/deep/group/med/mitdb"

def extract_wave(idx):

    """

    Reads .dat file and returns in numpy array. Assumes 2 channels.  The

    returned array is n x 3 where n is the number of samples. The first column

    is the sample number and the second two are the first and second channel

    respectively.

    """

    rdsamp = os.path.join(WFDB, 'rdsamp')

    output = subprocess.check_output([rdsamp, '-r', idx], cwd=DATA)

    data = np.fromstring(output, dtype=np.int32, sep=' ')

    return data.reshape((-1, 3))

def extract_annotation(idx):

    """

    The annotation file column names are:

        Time, Sample #, Type, Sub, Chan, Num, Aux

    The Aux is optional, it could be left empty. Type is the beat type and Aux

    is the transition label.

    """

    rdann = os.path.join(WFDB, 'rdann')

    output = subprocess.check_output([rdann, '-r', idx, '-a', 'atr'], cwd=DATA)

    labels = (line.split() for line in output.strip().split("\n"))

    labels = [(l[0], int(l[1]), l[2], l[6] if len(l) == 7 else None)

                for l in labels]

    return labels

def extract(idx):

    """

    Extracts data and annotations from .dat and .atr files.

    Returns a numpy array for the data and a list of tuples for the labels.

    """

    data = extract_wave(idx)

    labels = extract_annotation(idx)

    return data, labels

def save(example, idx):

    """

    Saves data with numpy.save (load with numpy.load) and pickles labels. The

    files are saved in the same place as the raw data.

    """

    np.save(os.path.join(DATA, idx), example[0])

    with open(os.path.join(DATA, "{}.pkl".format(idx)), 'w') as fid:

        pickle.dump(example[1], fid)

if __name__ == "__main__":

    files = glob.glob(os.path.join(DATA, "*.dat"))

    idxs = [os.path.basename(f).split(".")[0] for f in files]

    for idx in idxs:

        example = extract(idx)

        save(example, idx)

        print("Example {}".format(idx))