"""
@author: Payam Dibaeinia
"""

import numpy as np
import pandas as pd
import tensorflow as tf


#not used in training rn.
class TFRecordData(object):
    def __init__(self):
        pass

    def build_record(self, record_path, scores, enhancer_expr, TF_expr, score_labels = []):
        """
        scores: numpy ndarray (#enahncer, length, 2, #TF)
        enhancer_expr: pandas DataFrame containing expression of all enhancers present in score array
        TF_expr: pandas DataFrame containing expression of all TFs
        score_labels: if not provided it assumes 1-to-1 map between rows of scores array and expression DataFrame
        """

        if list(score_labels) == [] and scores.shape[0] != enhancer_expr.shape[0]:
            raise ValueError('Provided scores and expressions do not match, a label mapping is required or something is wrong!')

        if list(score_labels) == []:
            score_labels = enhancer_expr.index

        self._record_writer(scores, score_labels, enhancer_expr, TF_expr, record_path)

    def _float_feature(self, val):
        return tf.train.Feature(float_list = tf.train.FloatList(value = val.reshape(-1)))

    def _encoder(self, seqScore, gE_bin, tfE_bin):
        return tf.train.Example(features = tf.train.Features(feature = {'score': self._float_feature(seqScore),
                                                                        'tfE': self._float_feature(tfE_bin),
                                                                        'gE': self._float_feature(gE_bin)})).SerializeToString()

    def _iter_data (self, scores, score_labels, enhancer_expr, TF_expr):
        c = 0
        total = scores.shape[0]
        TF_expr = TF_expr.values.T
        for s, l in zip(scores, score_labels):
            print("{} out of {}".format(c,total))
            c+=1
            expr = enhancer_expr.loc[enhancer_expr.index == l].values
            if expr.shape[0] != 1:
                raise ValueError('provided expression DataFrame contains more than one row for at least one of the enhancers!')

            for gE_bin, tfE_bin in zip(expr[0,:], TF_expr):
                yield self._encoder(s, gE_bin, tfE_bin)

    def _record_writer(self, scores, score_labels, enhancer_expr, TF_expr, record_path):

        with tf.io.TFRecordWriter(record_path) as writer:
            for r in self._iter_data(scores, score_labels, enhancer_expr, TF_expr):
                writer.write(r)


    def get_dataset(self, record_path, score_shape, tfE_shape, nBatch, shuffle = True):
        """
        if shuffle is true, we recommend providing the datasize for perfect shuffling
        """

        dataset = tf.data.TFRecordDataset([record_path]).map(lambda x : self._decoder(x, score_shape, tfE_shape))
        if shuffle:
            dataset = dataset.shuffle(1000, reshuffle_each_iteration=True)

        return dataset.batch(nBatch)

    def _decoder(self, record, score_shape, tfE_shape):
        return tf.io.parse_single_example(record,{'score': tf.io.FixedLenFeature(shape = score_shape, dtype = tf.float32),
                                                  'tfE': tf.io.FixedLenFeature(shape = tfE_shape, dtype = tf.float32),
                                                  'gE': tf.io.FixedLenFeature(shape = (1), dtype = tf.float32)})

class InlineData(object):
    def __init__(self, trainLabels, scores_train, scores_valid, scores_test, seqExpr, tfExpr, nTrain, nValid, nTest, nBins):
        self.scoreDict_ = {'train': scores_train,
                           'valid': scores_valid,
                           'test': scores_test}

        self.seqExpr_ = seqExpr
        self.tfExpr_ = tfExpr
        self.nTrain_ = nTrain
        self.nValid_ = nValid
        self.nTest_ = nTest
        self.nBins_ = nBins
        self.idDict = {}
        self.trainLabelMap_ = np.array([list(seqExpr.index).index(i) for i in trainLabels])

    def _setIdDataset(self, split, size, shuffle, buffer, batch):

        self.idDict[split] = tf.data.Dataset.range(size)
        if shuffle:
            self.idDict[split] = self.idDict[split].shuffle(buffer, reshuffle_each_iteration=True)

        self.idDict[split] = self.idDict[split].batch(batch)

    def _getIndices(self, ind, split):
        """
        ind is a tensor of data indices generated by tf.data.Dataset.range
        """
        ind = ind.numpy()
        sID = ind//self.nBins_
        bID = ind%self.nBins_

        if split == 'train':
            eID = self.trainLabelMap_[sID]
        elif split == 'valid':
            eID = sID + self.nTrain_
        elif split == 'test':
            eID = sID + self.nTrain_ + self.nValid_

        return sID, bID, eID

    def get_dataset(self, split, shuffle, buffer, batch):
        if split not in ['train', 'valid', 'test']:
            raise ValueError('Requested data is not supported. Only support train, valid and test splits.')
        #TODO move this to seq_scan

        if split not in self.idDict.keys():
            size = self.scoreDict_[split].shape[0] * self.nBins_
            self._setIdDataset(split, size, shuffle, buffer, batch)

        for ids in self.idDict[split]:
            sInd, bInd, eInd = self._getIndices(ids, split)
            score = tf.gather(self.scoreDict_[split], sInd, axis = 0)
            tfE = tf.convert_to_tensor(self.tfExpr_.values[:,bInd].T)
            gE = tf.convert_to_tensor(self.seqExpr_.values[eInd, bInd].reshape(-1,1))

            yield {'score':score, 'tfE': tfE, 'gE': gE}
