# Author: xiaotaw@qq.com (Any bug report is welcome)
# Time: Aug 2016
# Addr: Shenzhen
# Description: define functions and parameters related to input data


from __future__ import absolute_import
from __future__ import division
from __future__ import print_function

import os
import numpy
import cPickle
import random # for random.shuffle()

data_dir = "data_files"
mgfp_dir = "mgfp_files" 

class DataSet(object):
  """dataset class, contains compds and labels, 
     and also provides method to generate a next batch of data. 
  """
  def __init__(self, compds, labels, dtype=numpy.float32):
    assert compds.shape[0] == labels.shape[0], "shape don't match"
    self.compds = compds.astype(dtype)
    self.labels = labels
    self.size = compds.shape[0]
    self.begin = 0
    self.end = 0

  def generate_batch(self, batch_size):
    assert self.compds.shape[0] == self.labels.shape[0], "shape don't match"
    assert batch_size <= self.size, "too big batch_size"
    self.end = self.begin + batch_size
    if self.end > self.size:
      self.shuffle()
      self.begin = 0
      self.end = batch_size
    compds_batch = self.compds[self.begin: self.end]
    labels_batch = self.labels[self.begin: self.end]
    self.begin = self.end
    return compds_batch, labels_batch
  
  def shuffle(self):
    perm = range(self.size)
    random.shuffle(perm)
    self.compds = self.compds[perm]
    self.labels = self.labels[perm]

def dense_to_one_hot(labels_dense, num_classes=2):
  """Convert class labels from scalars to one-hot vectors.
  """
  num_labels = labels_dense.shape[0]
  index_offset = numpy.arange(num_labels) * num_classes
  labels_one_hot = numpy.zeros((num_labels, num_classes))
  labels_one_hot.flat[index_offset + labels_dense.ravel().astype(int)] = 1
  # print(labels_one_hot)
  return labels_one_hot

def get_inputs(target):
  """read dataset from file according to target name
     note: target.compds are 2048 length mgfp code files, which were generated by 'cPickle.dump'.
  Args:
    target: a string of the protein kinase target name
  """
  with open(os.path.join(data_dir, target + ".compds"), "r") as compds_file, open(os.path.join(data_dir, target + ".labels"), "r") as labels_file:
    compds = cPickle.load(compds_file)
    labels = cPickle.load(labels_file)
    return DataSet(compds, labels, )

def get_inputs_pseudo(target, is_pos):
  """read dataset from file according to target name,
     'pseudo' means that labels are manually generated rather than read from file.
     note: target.compds are 2048 length mgfp code files, which were generated by 'cPickle.dump'.
  Args:
    target: a string of the protein kinase target name.
    is_pos: if true, labels are ones, else labels are zeros.
  """
  with open(os.path.join(data_dir, target + ".compds"), "r") as compds_file:
    compds = cPickle.load(compds_file)
    numpy.clip(compds, 0, 1, out=compds)
    if is_pos:
      labels_dense = numpy.ones(compds.shape[0], dtype=numpy.int32)
    else:
      labels_dense = numpy.zeros(compds.shape[0], dtype=numpy.int32)
    labels = dense_to_one_hot(labels_dense)
    return DataSet(compds, labels)

def 


@@ -337,16 +337,42 @@
 ffle()%0A%0A
+import tensorflow as tf%0A%0A%0A
 data_dir
@@ -413,16 +413,37 @@
 iles%22 %0A%0A
+record_bytes = 2048%0A%0A
 class Da
@@ -1875,16 +1875,27 @@
 t_inputs
+_by_cpickle
 (target)
@@ -1959,38 +1959,35 @@
    note: target.
-compds
+pkl
  are 2048 length
@@ -2155,36 +2155,34 @@
  target + %22.
-compds
+pkl
 %22), %22r
+b
 %22) as compds
@@ -2369,18 +2369,16 @@
 , labels
-, 
 )%0A%0Adef g
@@ -2387,16 +2387,27 @@
 _inputs_
+by_cpickle_
 pseudo(t
@@ -2838,20 +2838,18 @@
 + %22.
-compds
+pkl
 %22), %22r
+b
 %22) a
@@ -3189,10 +3189,645 @@
 s)%0A%0Adef 
+get_compds_by_filequeue(batch_size):%0A  %22%22%22%0A  %22%22%22%0A  target_list = %5B%22cdk2%22, %22egfr_erbB1%22, %22gsk3b%22, %22hgfr%22, %22map_k_p38a%22, %0A                 %22tpk_lck%22, %22tpk_src%22, %22vegfr2%22, %22pubchem_neg_sample%22%5D%0A  reader = tf.FixedLengthRecordReader(record_bytes=record_bytes)%0A  fq = %7B%7D%0A  value = %7B%7D%0A  records = %7B%7D%0A  for t in target_list:%0A    fq%5Bt%5D = tf.train.string_input_producer(%5B%22data_files/%22 + t + %22.bin%22%5D)%0A    _, value%5Bt%5D = reader.read(fq%5Bt%5D)%0A    records%5Bt%5D = tf.decode_raw(value%5Bt%5D, tf.uint8)%0A%0A  compds = tf.train.batch(records, batch_size=batch_size, num_threads=1, capacity=1000 + 3 * batch_size, shapes=%5B(2048,)%5D * len(records))%0A%0A  return compds%0A
 %0A%0A

add get_input_by_filequeue