import os

import pickle

import numpy as np

from scipy.io import wavfile

from scipy.fftpack import fft

from scipy.signal import butter, lfilter

from sklearn.preprocessing import normalize

from sklearn.cross_validation import train_test_split

from collections import namedtuple

from sklearn.cross_validation import check_random_state

class PCG:

    """

    PCG is a container for loading phonocardiogram (PCG) data for the [2016 physionet

    challenge](http://physionet.org/challenge/2016). Raw wav files are parsed into

    features, class labels are extracted from header files and data is split into

    training and testing groups.

    """

    def __init__(self, basepath, random_state=42):

        self.basepath = basepath

        self.class_name_to_id = {"normal": 0, "abnormal": 1}

        self.nclasses = len(self.class_name_to_id.keys())

        self.train = None

        self.test = None

        self.n_samples = 0

        self.X = None

        self.y = None

        self.random_state = random_state

    def initialize_wav_data(self):

        """

        Load the original wav files and extract features. Warning, this can take a

        while due to slow FFTs.

        Parameters

        ----------

        None

        Returns

        -------

        None

        """

        self.__load_wav_file()

        self.__split_train_test()

        # TODO: check if directory exists

        self.save("/tmp")

    def save(self, save_path):

        """

        Persist the PCG class to disk

        Parameters

        ----------

        save_path: str

            Location on disk to store the parsed PCG metadata

        Returns

        -------

        None

        """

        np.save(os.path.join(save_path, "X.npy"), self.X)

        np.save(os.path.join(save_path, "y.npy"), self.y)

        with open( os.path.join(save_path, "meta"), "w") as fout:

            pickle.dump((self.basepath, self.class_name_to_id, self.nclasses,

                         self.n_samples, self.random_state), fout)

    def load(self, load_path):

        """

        Load a previously stored PCG class.

        Parameters

        ----------

        load_path: str

            Location on disk to load parsed PCG data

        Returns

        -------

        None

        """

        self.X = np.load(os.path.join(load_path, "X.npy"))

        self.y = np.load(os.path.join(load_path, "y.npy"))

        with open(os.path.join(load_path, "meta"), "r") as fin:

            (self.basepath, self.class_name_to_id, self.nclasses,

             self.n_samples, self.random_state) = pickle.load(fin)

        self.__split_train_test()

    def __load_wav_file(self, doFFT=True):

        """

        Loads physio 2016 challenge dataset from self.basepath by crawling the path.

        For each discovered wav file:

        * Attempt to parse the header file for class label

        * Attempt to load the wav file

        * Calculate features from the wav file. if doFFT, features are

        the Fourier transform of the original signal. Else, features are

        the raw signal itself truncated to a fixed length

        Parameters

        ----------

        doFFT: bool

            True if features to be calculated are the FFT of the original signal

        Returns

        -------

        None

        """

        # First pass to calculate number of samples

        # ensure each wav file has an associated and parsable

        # Header file

        wav_file_names = []

        class_labels = []

        for root, dirs, files in os.walk(self.basepath):

            # Ignore validation for now!

            if "validation" in root:

                continue

            for file in files:

                if file.endswith('.wav'):

                    try:

                        base_file_name = file.rstrip(".wav")

                        label_file_name = os.path.join(root, base_file_name + ".hea")

                        class_label = self.__parse_class_label(label_file_name)

                        class_labels.append(self.class_name_to_id[class_label])

                        wav_file_names.append(os.path.join(root, file))

                        self.n_samples += 1

                    except InvalidHeaderFileException as e:

                        print e

        if doFFT:

            fft_embedding_size = 400

            highpass_embedding_size = 200

            X = np.zeros([self.n_samples, fft_embedding_size + highpass_embedding_size])

        else:

            # Truncating the length of each wav file to the

            # min file size (10611) (Note: this is bad

            # And causes loss of information!)

            embedding_size = 10611

            X = np.zeros([self.n_samples, embedding_size])

        for idx, wavfname in enumerate(wav_file_names):

            rate, wf = wavfile.read(wavfname)

            wf = normalize(wf.reshape(1, -1))

            if doFFT:

                # We only care about the magnitude of each frequency

                wf_fft = np.abs(fft(wf))

                wf_fft = wf_fft[:, :fft_embedding_size].reshape(-1)

                # Filter out high frequencies via Butter transform

                # The human heart maxes out around 150bpm = 2.5Hz

                # Let's filter out any frequency significantly above this

                nyquist = 0.5 * rate

                cutoff_freq = 4.0  # Hz

                w0, w1 = butter(5, cutoff_freq / nyquist, btype='low', analog=False)

                wf_low_pass = lfilter(w0, w1, wf)

                # FFT the filtered signal

                wf_low_pass_fft = np.abs(fft(wf_low_pass))

                wf_low_pass_fft = wf_low_pass_fft[:, :highpass_embedding_size].reshape(-1)

                features = np.concatenate((wf_fft, wf_low_pass_fft))

            else:

                features = wf[:embedding_size]

            X[idx, :] = features

            idx += 1

        self.X = X

        class_labels = np.array(class_labels)

        # Map from dense to one hot

        self.y = np.eye(self.nclasses)[class_labels]

    def __parse_class_label(self, label_file_name):

        """

        Parses physio bank header files, where the class label

        is located in the last line of the file. An example header

        file could contain:

        f0112 1 2000 60864

        f0112.wav 16+44 1 16 0 0 0 0 PCG

        # Normal

        Parameters

        ----------

        label_file_name: str

            Path to a specific header file

        Returns

        -------

        class_label: str

            One of `normal` or `abnormal`

        """

        with open(label_file_name, 'r') as fin:

            header = fin.readlines()

        comments = [line for line in header if line.startswith("#")]

        if not len(comments) == 1:

            raise InvalidHeaderFileException("Invalid label file %s" % label_file_name)

        class_label = str(comments[0]).lstrip("#").rstrip("\r").strip().lower()

        if not class_label in self.class_name_to_id.keys():

            raise InvalidHeaderFileException("Invalid class label %s" % class_label)

        return class_label

    def __split_train_test(self):

        """

        Splits internal features (self.X) and class labels (self.y) into

        balanced training and test sets using sklearn's helper function.

        Notes:

         * if self.random_state is None, splits will be randomly seeded

           otherwise, self.random_state defines the random seed to deterministicly

           split training and test data

         * For now, class balancing is done by subsampling the overrepresented class.

          Ideally this would be pushed down to the cost function in TensorFlow.

        Returns

        -------

        None

        """

        mlData = namedtuple('ml_data', 'X y')

        num_pos, num_neg = np.sum(self.y, axis=0)

        # Remove samples to rebalance classes

        # TODO: push this down into the cost function

        undersample_rate = num_neg / num_pos

        over_represented_idxs = self.y[:, 1] == 0

        under_represented_idxs = self.y[:, 1] == 1

        random_indexes_to_remove = np.random.rand(self.y.shape[0]) < undersample_rate

        sample_idxs = (over_represented_idxs & random_indexes_to_remove |

                       under_represented_idxs)

        X_balanced = self.X[sample_idxs, :]

        y_balanced = self.y[sample_idxs, :]

        X_train, X_test, y_train, y_test = train_test_split(X_balanced, y_balanced, test_size=0.25,

                                                            random_state=self.random_state)

        self.train = mlData(X=X_train, y=y_train)

        self.test = mlData(X=X_test, y=y_test)

    def get_mini_batch(self, batch_size):

        """

        Helper function for sampling mini-batches from the training

        set. Note, random_state needs to be set to None or the same

        mini batch will be sampled eternally!

        Parameters

        ----------

        batch_size: int

            Number of elements to return in the mini batch

        Returns

        -------

        X: np.ndarray

            A feature matrix subsampled from self.train

        y: np.ndarray

            A one-hot matrix of class labels subsampled from self.train

        """

        random_state = check_random_state(None)  # self.random_state)

        n_training_samples = self.train.X.shape[0]

        minibatch_indices = random_state.randint(0, n_training_samples - 1, batch_size)

        return self.train.X[minibatch_indices, :], self.train.y[minibatch_indices, :]

class InvalidHeaderFileException(Exception):

    def __init__(self, *args, **kwargs):

        super(args, kwargs)