import cPickle
import functools
import os

import fuel
import numpy
from fuel.schemes import (
    ConstantScheme, ShuffledExampleScheme)
from fuel.streams import DataStream
from fuel.transformers import (
    SortMapping, Padding, ForceFloatX, Batch, Mapping, Unpack, Filter,
    FilterSources, Transformer)

from lvsr.datasets.h5py import H5PYAudioDataset
from lvsr.preprocessing import log_spectrogram


def switch_first_two_axes(batch):
    result = []
    for array in batch:
        if array.ndim == 2:
            result.append(array.transpose(1, 0))
        else:
            result.append(array.transpose(1, 0, 2))
    return tuple(result)


def _length(example):
    return len(example[0])


def apply_preprocessing(preprocessing, example):
    recording, label = example
    return (numpy.asarray(preprocessing(recording)), label)


class _AddEosLabelEnd(object):

    def __init__(self, eos_label):
        self.eos_label = eos_label

    def __call__(self, example):
        return (example[0], list(example[1]) + [self.eos_label]) + tuple(example[2:])


class _AddEosLabelBeginEnd(object):

    def __init__(self, eos_label):
        self.eos_label = eos_label

    def __call__(self, example):
        return (example[0], [self.eos_label] + list(example[1]) + [self.eos_label]) + tuple(example[2:])


class _LengthFilter(object):

    def __init__(self, max_length):
        self.max_length = max_length

    def __call__(self, example):
        if self.max_length:
            return len(example[0]) <= self.max_length
        return True


class ForceCContiguous(Transformer):
    """Force all floating point numpy arrays to be floatX."""
    def __init__(self, data_stream):
        super(ForceCContiguous, self).__init__(
            data_stream, axis_labels=data_stream.axis_labels)

    def get_data(self, request=None):
        if request is not None:
            raise ValueError
        data = next(self.child_epoch_iterator)
        result = []
        for piece in data:
            if isinstance(piece, numpy.ndarray):
                result.append(numpy.ascontiguousarray(piece))
            else:
                result.append(piece)
        return tuple(result)


class Data(object):
    """Dataset manager.

    This class is in charge of accessing different datasets
    and building preprocessing pipelines.

    Parameters
    ----------