"""
@author: Payam Dibaeinia
"""

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

class InlineData(object):
    def __init__(self, seq_file, seq_names, seqExpr, tfExpr, nTrain, nValid, nTest, nBins):

        self.trainSize_ = nTrain * nBins
        self.validSize_ = nValid * nBins
        self.testSize_ = nTest * nBins
        self.idMap_ = np.array([list(seqExpr.index).index(i) for i in seq_names])
        with open(seq_file,'r') as f:
            self.allSeq_ = f.read().splitlines()
        self.allSeq_ = np.array(self.allSeq_)[np.arange(1,len(self.allSeq_), 2)]
        self.seqExpr_ = seqExpr
        self.tfExpr_ = tfExpr
        self.nBins_ = nBins
        self.idDict = {}

    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()
        if split == 'valid':
            ind += self.trainSize_
        elif split == 'test':
            ind += self.trainSize_ + self.validSize_

        sID = ind//self.nBins_
        bID = ind%self.nBins_
        eID = self.idMap_[sID]


        return sID, bID, eID

    def _encode(self,seq_list):
        """
        seq_list: list of DNA sequence as string of length L
        returns: tensor(N,L,4,1), of one-hot encoded sequences

        4 columns: A, C, G, T
        """
        N = len(seq_list)
        L = len(seq_list[0])
        ret = np.zeros((N,L,4,1))


        for si, s in enumerate(seq_list):
            for bi, b in enumerate(s):
                if b.capitalize() == 'A':
                    ret[si,bi,0,0] = 1
                elif b.capitalize() == 'C':
                    ret[si,bi,1,0] = 1
                elif b.capitalize() == 'G':
                    ret[si,bi,2,0] = 1
                elif b.capitalize() == 'T':
                    ret[si,bi,3,0] = 1

                # to makes sure dummy bases result in zero scores
                # note that we work in LLR format and then convert it to probability
                # if need (by taking exp()). Note that in LLR format, log(PWMs) are negative numbers because
                # they are probabilities
                elif b.capitalize() == 'N':
                    ret[si,bi,:,0] = 1e6

        return tf.convert_to_tensor(ret)

    def get_dataset(self, split, shuffle, 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():
            if split == 'train':
                size = self.trainSize_
                buffer = 10 * size
            elif split == 'valid':
                size = self.validSize_
                buffer = 100
            else:
                size = self.testSize_
                buffer = 100
            self._setIdDataset(split, size, shuffle, buffer, batch)

        for ids in self.idDict[split]:
            sInd, bInd, eInd = self._getIndices(ids, split)
            seq = self.allSeq_[sInd]
            seq = self._encode(seq)
            tfE = tf.convert_to_tensor(self.tfExpr_.values[:,bInd].T)
            gE = tf.convert_to_tensor(self.seqExpr_.values[eInd, bInd].reshape(-1,1))

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