import mxnet as mx
import numpy as np
import sys
from io_func.feat_io import DataReadStream

# The interface of a data iter that works for bucketing
#
# DataIter
#   - default_bucket_key: the bucket key for the default symbol.
#
# DataBatch
#   - provide_data: same as DataIter, but specific to this batch
#   - provide_label: same as DataIter, but specific to this batch
#   - bucket_key: the key for the bucket that should be used for this batch


def read_content(path):
    with open(path) as input:
        content = input.read()
        content = content.replace('\n', ' <eos> ').replace('. ', ' <eos> ')
        return content


class SimpleBatch(object):
    def __init__(self, data_names, data, label_names, label, bucket_key,
                 utt_id=None, utt_len=0, effective_sample_count=None):
        self.data = data
        self.label = label
        self.data_names = data_names
        self.label_names = label_names
        self.bucket_key = bucket_key
        self.utt_id = utt_id
        self.utt_len = utt_len
        self.effective_sample_count = effective_sample_count

        self.pad = 0
        self.index = None  # TODO: what is index?

    @property
    def provide_data(self):
        return [(n, x.shape) for n, x in zip(self.data_names, self.data)]

    @property
    def provide_label(self):
        if len(self.label_names):
            return [(n, x.shape) for n, x in zip(self.label_names, self.label)]
        else:
            return None

class SimpleIter(mx.io.DataIter):
    """DataIter used in Calculate Statistics (in progress).

    Parameters
    ----------
    pad_zeros : bool
        Default `False`. Control the behavior of padding when we run
        out of the whole dataset. When true, we will pad with all-zeros.
        When false, will pad with a random sentence in the dataset.
        Usually, for training we would like to use `False`, but
        for testing use `True` so that the evaluation metric can
        choose to ignore the padding by detecting the zero-labels.
    """
    def __init__(self, train_sets, batch_size,
            init_states, delay=5, feat_dim=40, label_dim=1955,
            label_mean_sets=None, data_name='data',
            label_name='softmax_label', has_label=True, load_label_mean=True):

        self.train_sets = train_sets
        self.label_mean_sets = label_mean_sets
        self.train_sets.initialize_read()

        self.data_name = data_name
        if has_label:
            self.label_name = label_name

        features = []
        labels = []
        utt_lens = []
        utt_ids = []
        buckets = []
        self.has_label = has_label

        if label_mean_sets is not None:
            self.label_mean_sets.initialize_read()
            (feats, tgts, utt_id) = self.label_mean_sets.load_next_seq()

            self.label_mean = feats/np.sum(feats)
            for i, v in enumerate(feats):
                if v <= 1.0:
                    self.label_mean[i] = 1

        sys.stderr.write("Loading data...\n")
        buckets_map = {}
        n = 0
        while True:
            (feats, tgts, utt_id) = self.train_sets.load_next_seq()
            if utt_id is None:
                break
            if tgts is None and self.has_label:
                continue
            if feats.shape[0] == 0:
                continue
            features.append(feats)
            utt_lens.append(feats.shape[0])
            utt_ids.append(utt_id)
            if self.has_label:
                labels.append(tgts+1)
            if feats.shape[0] not in buckets:
                buckets_map[feats.shape[0]] = feats.shape[0]

        for k, v in buckets_map.iteritems():
            buckets.append(k)

        buckets.sort()
        i_max_bucket = len(buckets)-1
        max_bucket = buckets[i_max_bucket]
        self.buckets = buckets
        self.data = [[] for k in buckets]
        self.utt_id = [[] for k in buckets]
        self.utt_lens = [[] for k in buckets]
        self.feat_dim = feat_dim
        self.default_bucket_key = max(buckets)

        for i, feats in enumerate(features):
            if has_label:
                tgts = labels[i]
            utt_len = utt_lens[i]
            utt_id = utt_ids[i]

            for i, bkt in enumerate(buckets):
                if bkt >= utt_len:
                    i_bucket = i
                    break

            if self.has_label:
                self.data[i_bucket].append((feats, tgts))
            else:
                self.data[i_bucket].append(feats)
            self.utt_id[i_bucket].append(utt_id)
            self.utt_lens[i_bucket].append(utt_len)

        # Get the size of each bucket, so that we could sample
        # uniformly from the bucket
        bucket_sizes = [len(x) for x in self.data]

        self.batch_size = batch_size
        # convert data into ndarrays for better speed during training

        data = [np.zeros((len(x), buckets[i], self.feat_dim))
                if len(x) % self.batch_size == 0
                else np.zeros(((len(x)/self.batch_size + 1) * self.batch_size, buckets[i], self.feat_dim))
                for i, x in enumerate(self.data)]

        label = [np.zeros((len(x), buckets[i]))
                 if len(x) % self.batch_size == 0
                 else np.zeros(((len(x)/self.batch_size + 1) * self.batch_size, buckets[i]))
                 for i, x in enumerate(self.data)]

        utt_id = [[] for k in buckets]
        for i, x in enumerate(data):
            utt_id[i] = ["GAP_UTT"] * len(x)
        utt_lens = [[] for k in buckets]
        for i, x in enumerate(data):
            utt_lens[i] = [0] * len(x)


        for i_bucket in range(len(self.buckets)):
            for j in range(len(self.data[i_bucket])):
                sentence = self.data[i_bucket][j]
                if self.has_label:
                    sentence[1][delay:] = sentence[1][:-delay]
                    sentence[1][:delay] = sentence[1][0] # broadcast assignment
                    data[i_bucket][j, :len(sentence[0])] = sentence[0]
                    label[i_bucket][j, :len(sentence[1])] = sentence[1]
                else:
                    data[i_bucket][j, :len(sentence)] = sentence
                    # borrow this place to pass in sentence length. TODO: use a less hacky way.
                    label[i_bucket][j, :len(sentence)] += len(sentence)

                utt_id[i_bucket][j] = self.utt_id[i_bucket][j]
                utt_lens[i_bucket][j] = self.utt_lens[i_bucket][j]

        self.data = data
        self.label = label
        self.utt_id = utt_id
        self.utt_lens = utt_lens


        # Get the size of each bucket, so that we could sample
        # uniformly from the bucket
        bucket_sizes = [len(x) for x in self.data]

        sys.stderr.write("Summary of dataset ==================\n")
        for bkt, sz in zip(buckets, bucket_sizes):
            sys.stderr.write("bucket of len %3d : %d samples\n" % (bkt, sz))

        bucket_size_tot = float(sum(bucket_sizes))

        self.bucket_sizes = bucket_sizes
        self.make_data_iter_plan()

        self.init_states = init_states
        self.init_state_arrays = [mx.nd.zeros(x[1]) for x in init_states]

        self.provide_data = [(data_name, (batch_size, self.default_bucket_key, self.feat_dim))] + init_states
        self.provide_label = None
        if has_label:
            self.provide_label = [(label_name, (self.batch_size, self.default_bucket_key))]

    def make_data_iter_plan(self):
        "make a random data iteration plan"
        # truncate each bucket into multiple of batch-size
        bucket_n_batches = []
        for i in range(len(self.data)):
            bucket_n_batches.append(len(self.data[i]) / self.batch_size)
            self.data[i] = self.data[i][:int(bucket_n_batches[i]*self.batch_size),:]
            self.label[i] = self.label[i][:int(bucket_n_batches[i]*self.batch_size)]

        bucket_plan = np.hstack([np.zeros(n, int)+i for i, n in enumerate(bucket_n_batches)])
        np.random.shuffle(bucket_plan)

        bucket_idx_all = [np.random.permutation(len(x)) for x in self.data]

        self.bucket_plan = bucket_plan
        self.bucket_idx_all = bucket_idx_all
        self.bucket_curr_idx = [0 for x in self.data]

        self.data_buffer = []
        self.label_buffer = []
        for i_bucket in range(len(self.data)):
            data = mx.nd.zeros((self.batch_size, self.buckets[i_bucket], self.feat_dim))
            label = mx.nd.zeros((self.batch_size, self.buckets[i_bucket]))
            self.data_buffer.append(data)
            self.label_buffer.append(label)

    def __iter__(self):
        init_state_names = [x[0] for x in self.init_states]
        data_names = [self.data_name] + init_state_names
        label_names = []
        if self.has_label:
            label_names = [self.label_name]

        for i_bucket in self.bucket_plan:
            data = self.data_buffer[i_bucket]
            label = self.label_buffer[i_bucket]

            i_idx = self.bucket_curr_idx[i_bucket]
            idx = self.bucket_idx_all[i_bucket][i_idx:i_idx+self.batch_size]
            self.bucket_curr_idx[i_bucket] += self.batch_size
            data[:] = self.data[i_bucket][idx]
            label[:] = self.label[i_bucket][idx]
            data_all = [data] + self.init_state_arrays
            label_all = [label]
            utt_id = np.array(self.utt_id[i_bucket])[idx]
            utt_len = np.array(self.utt_lens[i_bucket])[idx]
            effective_sample_count = mx.nd.sum(label)
            data_batch = SimpleBatch(data_names, data_all, label_names, label_all,
                                     self.buckets[i_bucket], utt_id, utt_len,
                                     effective_sample_count=effective_sample_count)
            yield data_batch

    def reset(self):
        self.bucket_curr_idx = [0 for x in self.data]

class TruncatedSentenceIter(mx.io.DataIter):
    """DataIter used in Truncated-BPTT.

    Each sentence is split into chunks of fixed lengths. The states are
    forwarded during forward, but the backward is only computed within
    chunks. This mechanism does not require bucketing, and it sometimes
    avoid gradient exploding problems in very long sequences.

    Parameters
    ----------
    pad_zeros : bool
        Default `False`. Control the behavior of padding when we run
        out of the whole dataset. When true, we will pad with all-zeros.
        When false, will pad with a random sentence in the dataset.
        Usually, for training we would like to use `False`, but
        for testing use `True` so that the evaluation metric can
        choose to ignore the padding by detecting the zero-labels.
    """
    def __init__(self, train_sets, batch_size, init_states, truncate_len=20, delay=5,
                 feat_dim=40, data_name='data', label_name='softmax_label',
                 has_label=True, do_shuffling=True, pad_zeros=False, time_major=False):

        self.train_sets = train_sets
        self.train_sets.initialize_read()

        self.data_name = data_name
        self.label_name = label_name

        self.feat_dim = feat_dim
        self.has_label = has_label
        self.batch_size = batch_size
        self.truncate_len = truncate_len
        self.delay = delay

        self.do_shuffling = do_shuffling
        self.pad_zeros = pad_zeros

        self.time_major = time_major

        self.label = None
        if self.time_major:
            self.data = [mx.nd.zeros((truncate_len, batch_size, feat_dim))]
            if has_label:
                self.label = [mx.nd.zeros((truncate_len, batch_size))]
        else:
            self.data = [mx.nd.zeros((batch_size, truncate_len, feat_dim))]
            if has_label:
                self.label = [mx.nd.zeros((batch_size, truncate_len))]
 
        self.init_state_names = [x[0] for x in init_states]
        self.init_state_arrays = [mx.nd.zeros(x[1]) for x in init_states]

        self.provide_data = [(data_name, self.data[0].shape)] + init_states
        self.provide_label = None
        if has_label:
            self.provide_label = [(label_name, self.label[0].shape)]

        self._load_data()
        self._make_data_plan()

    def _load_data(self):
        sys.stderr.write('Loading data into memory...\n')
        self.features = []
        self.labels = []
        self.utt_ids = []

        seq_len_tot = 0.0
        while True:
            (feats, tgs, utt_id) = self.train_sets.load_next_seq()
            if utt_id is None:
                break
            if tgs is None and self.has_label:
                continue
            if feats.shape[0] == 0:
                continue

            if self.has_label and self.delay > 0:
                # delay the labels
                tgs[self.delay:] = tgs[:-self.delay]
                tgs[:self.delay] = tgs[0]  # boradcast assign
            self.features.append(feats)
            if self.has_label:
                self.labels.append(tgs+1)
            self.utt_ids.append(utt_id)
            seq_len_tot += feats.shape[0]

        sys.stderr.write('    %d utterances loaded...\n' % len(self.utt_ids))
        sys.stderr.write('    avg-sequence-len = %.0f\n' % (seq_len_tot/len(self.utt_ids)))

    def _make_data_plan(self):
        if self.do_shuffling:
            # TODO: should we group utterances of similar length together?
            self._data_plan = np.random.permutation(len(self.features))
        else:
            # we might not want to do shuffling for testing for example
            self._data_plan = np.arange(len(self.features))

    def __iter__(self):
        assert len(self._data_plan) >= self.batch_size, \
            "Total number of sentences smaller than batch size, consider using smaller batch size"
        utt_idx = self._data_plan[:self.batch_size]
        utt_inside_idx = [0] * self.batch_size

        next_utt_idx = self.batch_size
        is_pad = [False] * self.batch_size
        pad = 0
        
        if self.time_major:
            np_data_buffer = np.zeros((self.truncate_len, self.batch_size, self.feat_dim))
            np_label_buffer = np.zeros((self.truncate_len, self.batch_size))
        else:
            np_data_buffer = np.zeros((self.batch_size, self.truncate_len, self.feat_dim))
            np_label_buffer = np.zeros((self.batch_size, self.truncate_len))
 
        utt_id_buffer = [None] * self.batch_size

        data_names = [self.data_name] + self.init_state_names
        label_names = [self.label_name]

        # reset states
        for state in self.init_state_arrays:
            state[:] = 0.1

        while True:
            effective_sample_count = self.batch_size * self.truncate_len
            for i, idx in enumerate(utt_idx):
                fea_utt = self.features[idx]
                if utt_inside_idx[i] >= fea_utt.shape[0]:
                    # we have consumed this sentence

                    # reset the states
                    for state in self.init_state_arrays:
                        if self.time_major:
                            state[:, i:i+1, :] = 0.1
                        else:
                            state[i:i+1] = 0.1
                    # load new sentence
                    if is_pad[i]:
                        # I am already a padded sentence, just rewind to the
                        # beginning of the sentece
                        utt_inside_idx[i] = 0
                    elif next_utt_idx >= len(self.features):
                        # we consumed the whole dataset, simply repeat this sentence
                        # and set pad
                        pad += 1
                        is_pad[i] = True
                        utt_inside_idx[i] = 0
                    else:
                        # move to the next sentence
                        utt_idx[i] = self._data_plan[next_utt_idx]
                        idx = utt_idx[i]
                        fea_utt = self.features[idx]
                        utt_inside_idx[i] = 0
                        next_utt_idx += 1

                if is_pad[i] and self.pad_zeros:
                    np_data_buffer[i] = 0
                    np_label_buffer[i] = 0
                    effective_sample_count -= self.truncate_len
                else:
                    idx_take = slice(utt_inside_idx[i],
                                     min(utt_inside_idx[i]+self.truncate_len,
                                         fea_utt.shape[0]))
                    n_take = idx_take.stop - idx_take.start
                    if self.time_major:
                        np_data_buffer[:n_take, i, :] = fea_utt[idx_take]
                        np_label_buffer[:n_take, i] = self.labels[idx][idx_take]
                    else:
                        np_data_buffer[i, :n_take, :] = fea_utt[idx_take]
                        np_label_buffer[i, :n_take] = self.labels[idx][idx_take]
 
                    if n_take < self.truncate_len:
                        if self.time_major:
                            np_data_buffer[n_take:, i, :] = 0
                            np_label_buffer[n_take:, i] = 0
                        else:
                            np_data_buffer[i, n_take:, :] = 0
                            np_label_buffer[i, n_take:] = 0
 
                        effective_sample_count -= self.truncate_len - n_take

                    utt_inside_idx[i] += n_take

                utt_id_buffer[i] = self.utt_ids[idx]

            if pad == self.batch_size:
                # finished all the senteces
                break
            
            self.data[0][:] = np_data_buffer
            self.label[0][:] = np_label_buffer
 
            data_batch = SimpleBatch(data_names, 
                                     self.data + self.init_state_arrays,
                                     label_names, self.label, bucket_key=None,
                                     utt_id=utt_id_buffer,
                                     effective_sample_count=effective_sample_count)

            # Instead of using the 'pad' property, we use an array 'is_pad'. Because
            # our padded sentence could be in the middle of a batch. A sample is pad
            # if we are running out of the data set and they are just some previously
            # seen data to be filled for a whole batch. In prediction, those data
            # should be ignored
            data_batch.is_pad = is_pad

            yield data_batch

    def reset(self):
        self._make_data_plan()


class BucketSentenceIter(mx.io.DataIter):
    def __init__(self, train_sets, buckets, batch_size,
                 init_states, delay=5, feat_dim=40,
                 data_name='data', label_name='softmax_label', has_label=True):

        self.train_sets = train_sets
        self.train_sets.initialize_read()

        self.data_name = data_name
        self.label_name = label_name

        buckets.sort()
        i_max_bucket = len(buckets)-1
        max_bucket = buckets[i_max_bucket]

        if has_label != True:
            buckets = [i for i in range(1, max_bucket)]
            i_max_bucket = len(buckets)-1
            max_bucket = buckets[i_max_bucket]

        self.buckets = buckets
        self.data = [[] for k in buckets]
        self.utt_id = [[] for k in buckets]
        self.feat_dim = feat_dim
        self.default_bucket_key = max(buckets)
        self.has_label = has_label

        sys.stderr.write("Loading data...\n")
        T_OVERLAP = buckets[0]/2
        n = 0
        while True:
            (feats, tgts, utt_id) = self.train_sets.load_next_seq()
            if utt_id is None:
                break
            if tgts is None and self.has_label:
                continue
            if feats.shape[0] == 0:
                continue

            # we split sentence into overlapping segments if it is
            # longer than the largest bucket
            t_start = 0
            t_end = feats.shape[0]
            while t_start < t_end:
                if t_end - t_start > max_bucket:
                    t_take = max_bucket
                    i_bucket = i_max_bucket
                else:
                    for i, bkt in enumerate(buckets):
                        if bkt >= t_end-t_start:
                            t_take = t_end-t_start
                            i_bucket = i
                            break

                n += 1
                if self.has_label:
                    self.data[i_bucket].append((feats[t_start:t_start+t_take],
                                                tgts[t_start:t_start+t_take]+1))
                else:
                    self.data[i_bucket].append(feats[t_start:t_start+t_take])

                self.utt_id[i_bucket].append(utt_id)
                t_start += t_take
                if t_start >= t_end:
                    # this sentence is consumed
                    break
                t_start -= T_OVERLAP

        # Get the size of each bucket, so that we could sample
        # uniformly from the bucket
        bucket_sizes = [len(x) for x in self.data]

        self.batch_size = batch_size
        # convert data into ndarrays for better speed during training

        data = [np.zeros((len(x), buckets[i], self.feat_dim))
                if len(x) % self.batch_size == 0
                else np.zeros(((len(x)/self.batch_size + 1) * self.batch_size, buckets[i],
                               self.feat_dim))
                for i, x in enumerate(self.data)]

        label = [np.zeros((len(x), buckets[i]))
                 if len(x) % self.batch_size == 0
                 else np.zeros(((len(x)/self.batch_size + 1) * self.batch_size, buckets[i]))
                 for i, x in enumerate(self.data)]

        utt_id = [[] for k in buckets]
        for i, x in enumerate(data):
            utt_id[i] = ["GAP_UTT"] * len(x)

        for i_bucket in range(len(self.buckets)):
            for j in range(len(self.data[i_bucket])):
                sentence = self.data[i_bucket][j]
                if self.has_label:
                    sentence[1][delay:] = sentence[1][:-delay]
                    sentence[1][:delay] = sentence[1][0]  # broadcast assignment
                    data[i_bucket][j, :len(sentence[0])] = sentence[0]
                    label[i_bucket][j, :len(sentence[1])] = sentence[1]
                else:
                    data[i_bucket][j, :len(sentence)] = sentence
                    # borrow this place to pass in sentence length. TODO: use a less hacky way.
                    label[i_bucket][j, :len(sentence)] += len(sentence)

                utt_id[i_bucket][j] = self.utt_id[i_bucket][j]

        self.data = data
        self.label = label
        self.utt_id = utt_id

        # Get the size of each bucket, so that we could sample
        # uniformly from the bucket
        bucket_sizes = [len(x) for x in self.data]

        sys.stderr.write("Summary of dataset ==================\n")
        for bkt, sz in zip(buckets, bucket_sizes):
            sys.stderr.write("bucket of len %3d : %d samples\n" % (bkt, sz))

        self.bucket_sizes = bucket_sizes
        self.make_data_iter_plan()

        self.init_states = init_states
        self.init_state_arrays = [mx.nd.zeros(x[1]) for x in init_states]

        self.provide_data = [(data_name, (batch_size, self.default_bucket_key, self.feat_dim))] + \
            init_states
        self.provide_label = [(label_name, (self.batch_size, self.default_bucket_key))]

    def make_data_iter_plan(self):
        "make a random data iteration plan"
        # truncate each bucket into multiple of batch-size
        bucket_n_batches = []
        for i in range(len(self.data)):
            bucket_n_batches.append(len(self.data[i]) / self.batch_size)
            self.data[i] = self.data[i][:int(bucket_n_batches[i]*self.batch_size), :]
            self.label[i] = self.label[i][:int(bucket_n_batches[i]*self.batch_size)]

        bucket_plan = np.hstack([np.zeros(n, int)+i for i, n in enumerate(bucket_n_batches)])
        np.random.shuffle(bucket_plan)

        bucket_idx_all = [np.random.permutation(len(x)) for x in self.data]

        self.bucket_plan = bucket_plan
        self.bucket_idx_all = bucket_idx_all
        self.bucket_curr_idx = [0 for x in self.data]

        self.data_buffer = []
        self.label_buffer = []
        for i_bucket in range(len(self.data)):
            data = mx.nd.zeros((self.batch_size, self.buckets[i_bucket], self.feat_dim))
            label = mx.nd.zeros((self.batch_size, self.buckets[i_bucket]))
            self.data_buffer.append(data)
            self.label_buffer.append(label)

    def __iter__(self):
        init_state_names = [x[0] for x in self.init_states]
        data_names = [self.data_name] + init_state_names
        label_names = [self.label_name]

        for i_bucket in self.bucket_plan:
            data = self.data_buffer[i_bucket]
            label = self.label_buffer[i_bucket]

            i_idx = self.bucket_curr_idx[i_bucket]
            idx = self.bucket_idx_all[i_bucket][i_idx:i_idx+self.batch_size]
            self.bucket_curr_idx[i_bucket] += self.batch_size
            data[:] = self.data[i_bucket][idx]
            label[:] = self.label[i_bucket][idx]
            data_all = [data] + self.init_state_arrays
            label_all = [label]
            utt_id = np.array(self.utt_id[i_bucket])[idx]
            effective_sample_count = mx.nd.sum(label)
            data_batch = SimpleBatch(data_names, data_all, label_names, label_all,
                                     self.buckets[i_bucket], utt_id,
                                     effective_sample_count=effective_sample_count)
            yield data_batch

    def reset(self):
        self.bucket_curr_idx = [0 for x in self.data]