text: import numpy as np
  import os
  from sklearn.preprocessing import MinMaxScaler
  import pickle
  import tensorflow as tf
  import vae_
  from sklearn.utils import shuffle
  import collections
  
  STEM_SAMPLE_POINT_NUM = 20
  
  data_size = 40*3
  batch_size = 50
  
  # settings for VAE.
  n_hidden = 100  # num of hidden units in mlp.
  dim_z = 20  # dim of latent.
  
  br_data_size = 4 * dim_z + 3
  
  # data_path = "C:/Users/Administrator/Desktop/neuron_factor_tree_test/test/data/L4_BC/data"
  #
  # sess_save_dir = "C:/Users/Administrator/Desktop/neuron_factor_tree_test/test/data/L4_BC/data/save"
  # save_dir = "C:/Users/Administrator/Desktop/neuron_factor_tree_test/test/data/L4_BC/data/save"
  
  '''
  Load both branch codes and branch id sequences.
  '''
  
  class SeqLoader():
      def __init__(self, test_path, seq_size, branch_size, batch_size, seq_type):
          self.data_path = os.path.join(test_path, "seq_data")
          self.save_path = os.path.join(test_path, "vae_save")
          self.seq_size = seq_size
          self.branch_size = branch_size
          self.batch_size = batch_size
          self.type = seq_type
  
          # load rnn data.
          training_data = os.path.join(self.data_path, "factorDB"+"_"+str(self.type)+".txt")
          if not (os.path.exists(training_data)):
              self.build_seq_data()
          else:
              self.load_processed(self.data_path)
          # self.build_seq_data()
  
          self.create_batches()
          self.reset_batch()
  
  
      # Collect all the branch codes(normalized).
      def branches_loader(self, branch_codes_file):
          '''
          A Branch : [angle, stm_scale_1, stm_scale_2, br_codes[sample_point_num, 3]].
          '''
          branches = []
          branch = []
          branch_codes = []
          br_angle = 0.0
          stm_scale = 1.0
          with open(branch_codes_file) as f:
              for line in f.readlines():
                  if line.startswith("#BRANCH:"):
                      branch = []
                      branch_codes = []
                  elif line.startswith("#ANGLE:"):
                      br_angle = float(line[:-1].split(" ")[1])
                      branch.append(br_angle)
                  elif line.startswith("#STEM:"):
                      stm_scale = float(line[:-1].split(" ")[1])
                      branch.append(stm_scale)
                  elif line.startswith("#DATA:"):
                      code_str = np.array(line[:-1].split(" ")[1:4])
                      code = list(map(float, code_str))
                      branch_codes.append(code)
                  elif line.startswith("#BRANCHEND"):
                      branch.append(branch_codes)
                      # print(branch)
                      branches.append(branch)
  
          # print(len(branches))
          return branches
  
      # Load sequence as branch ids, only load a type each time.
      def sequences_loader(self, branch_database_path):
          seqs = []
          seq = []
          current_type = 0
          with open(branch_database_path) as f:
              for line in f.readlines():
                  if line.startswith("#SEQ"):
                      if current_type==self.type and len(seq) != 0:
                          seqs.append(seq)
                      seq = []
                      current_type = int(np.array(line[:-1].split(" ")[1])) # update current seq type
  
                  elif line.startswith("#BID"):
                      b_id_str = np.array(line[:-1].split(" ")[1])
                      b_id = float(b_id_str)
                      seq.append(b_id)
  
          if current_type==self.type and len(seq)!=0:
              seqs.append(seq)
  
          # print(len(seqs))
          return seqs
  
      def process_branches(self, branches):
  
          print("Processing branches.")
  
          # load scaler_file created in branch_generator traning step.
          scaler_file = os.path.join(self.data_path, "scaler.txt")
          with open(scaler_file, "rb") as f:
              min_max_scaler = pickle.load(f)
          print("Scaler file loaded from [%s]." % scaler_file)
          self.min_max_scaler = min_max_scaler  # record the scaler.
  
          br_count = np.array([[br[0], br[1], br[2]] for br in branches])
          # print(br_count.shape)
  
          angle_stm_scaler = MinMaxScaler()
          angle_stm_scaler.fit(br_count)
  
          angle_stm_scaler_file = os.path.join(self.data_path, "angle_stm_scaler.txt")
          with open(angle_stm_scaler_file, "wb") as f:
              pickle.dump(angle_stm_scaler, f)
          print("Angle stm scaler file saved at [%s]." % angle_stm_scaler_file)
  
          # Normalixation
          data = [[angle_stm_scaler.transform(np.array([br[0], br[1], br[2]]).reshape(-1, 3)),
                   min_max_scaler.transform(br[3]).reshape(STEM_SAMPLE_POINT_NUM*2, 3,  1)]
                  for br in branches]
          # print(len(data))
          # print(data[0])
  
          return data
  
      def build_seq_data(self):
  
          print("Building sequence datas.")
  
          branch_seq_path = os.path.join(self.data_path, "branch_dataset.brd")
          branch_code_path = os.path.join(self.data_path, "branch_code.bcd")
          # Load branches.
          branches = self.branches_loader(branch_code_path)
          print("{} branches loaded.".format(len(branches)))
          self.branches = self.process_branches(branches)  # [None, [angle, scale_1, scale_2, codes]].
  
          # Load branch id sequences.
          self.sequences = self.sequences_loader(branch_seq_path)
          print("{} sequences loaded.".format(len(self.sequences)))
          if int(len(self.sequences) / self.batch_size) == 0:
              self.tensor = []
              return
          # print(self.sequences)
  
          # Sequence normalization.
          seq_ids = []
          for seq in self.sequences:
              seq.reverse()
              seq_n = []
              for id in [int(f) for f in seq]:
                  if id!=-1: # empty branches.
                      seq_n.append(id)
              if(len(seq_n)==0):
                  continue
              elif len(seq_n) < self.seq_size:
                  for i in range(self.seq_size - len(seq_n)):
                      seq_n.append(-1)
              else:  # equals or shorter than seq_size.
                  seq_n = seq_n[(len(seq_n) - self.seq_size) + 1:]
                  seq_n.append(-1) # make sure last one is -1.
  
              seq_ids.append(seq_n)
  
          # print("shape")
          # print(np.array(seq_ids).shape)
  
          # x = tf.placeholder(tf.float32, shape=[None, data_size], name='input_img')
          x = tf.placeholder(tf.float32, shape=[None, 20, 3, 1], name='input')
          latent = vae_.gaussian_MLP_encoder(x, n_hidden, dim_z, 1.0)
  
          with tf.Session() as sess:
  
              tf.global_variables_initializer().run()
              saver = tf.train.Saver(tf.global_variables())
              ckpt = tf.train.get_checkpoint_state(self.save_path)
              if ckpt:
                  print("VAE model restored for branch latent transform.")
                  saver.restore(sess, ckpt.model_checkpoint_path)
              else:
                  return
  
  
              latent_data = []
              for seq in seq_ids:
                  seq_latents = []
                  # seq_latents.append(np.array([1.0 for _ in range(br_data_size)])) # all 1.0 for soma.
                  for br in seq:
                      if br != -1:
                          _st_mean, _st_stdv = \
                              sess.run(latent, feed_dict={x: [self.branches[br][1][:STEM_SAMPLE_POINT_NUM]]})# stem 1
                          st_latent = np.append(_st_mean[0], _st_stdv[0])
                          _st_mean, _st_stdv = \
                              sess.run(latent, feed_dict={x: [self.branches[br][1][STEM_SAMPLE_POINT_NUM:]]})
                          st_latent_ = np.append(_st_mean[0], _st_stdv[0])
                          br_latent = np.append(np.array(self.branches[br][0]), st_latent)
                          br_latent = np.append(br_latent, st_latent_)
                          # print(len(br_latent))
                          # print(br_latent)
                          seq_latents.append(br_latent)
                  for i in range(self.seq_size - len(seq_latents)):
                      seq_latents.append(np.array([0.0 for _ in range(self.branch_size)])) # all 0 for leaf.
  
                  latent_data.append(seq_latents)
              # print(len(latent_data))
              # print(latent_data[0][0])
              x_data_ = np.array(latent_data, dtype=np.float32)
              print(x_data_.shape)  # (data_size, seq_size, 4*dim_z+3)
  
          print("Training data normalized.")
  
          x_data = shuffle(x_data_, random_state=42)
  
          # make up a neuronFactor-database. Each nf in DB have a id.
          factors_ = x_data.reshape(-1, self.branch_size)  # reshape to count factors.
          factors_s = [str(f) for f in factors_.tolist()]
          counter = collections.Counter(factors_s)
          count_pairs = sorted(counter.items(), key=lambda x: -x[1])
          self.factors, _ = zip(*count_pairs)
          print(len(self.factors))
          # for f in self.factors:
          #     print(f)
          self.factorDB_size = len(self.factors)
          self.factorDB = dict(zip(self.factors, range(len(self.factors))))
          db_file = os.path.join(self.data_path, "factorDB" + "_" + str(self.type) + ".txt")
          print("Got a NeuronFactor Database of size %d." % self.factorDB_size)
          tensor = np.array(list(map(self.factorDB.get, factors_s)))
          tensor_file = os.path.join(self.data_path, "tensor" + "_" + str(self.type) + "")
          self.tensor = tensor.reshape(-1, self.seq_size)  # seq_num * seq_size
  
          # save data.
          with open(db_file, "wb") as f:
              pickle.dump(self.factors, f)
              print("NeuronFactorDB saved at [%s]." % db_file)
          np.save(tensor_file, self.tensor)
          print("Factor tensors saved at [%s]." % tensor_file)
  
          print("Data preprocessed.")
  
      def load_processed(self, data_base_dir):
          db_file = os.path.join(data_base_dir, "factorDB_"+str(self.type)+".txt")
          tensor_file = os.path.join(data_base_dir, "tensor_"+str(self.type)+".npy")
  
          # Loading data.
          with open(db_file, "rb") as f:
              self.factors = pickle.load(f)
          self.factorDB_size = len(self.factors)
          self.factorDB = dict(zip(self.factors, range(len(self.factors))))
          self.tensor = np.load(tensor_file)
  
          print("[%d] training data loaded from [%s]." % (len(self.tensor), tensor_file))
  
      def create_batches(self):
          self.num_batches = int(len(self.tensor) / self.batch_size)
          if self.num_batches== 0:
              print("Don't have enough data, finish.")
              return
          # use the first self.batch_size*self.num_batches data.
          tensor = self.tensor[:self.batch_size*self.num_batches]
          x_data = tensor
          y_data = np.copy(tensor)
          for y_ in y_data:
              y_[:-1] = y_[1:]
              y_[-1] = 0
  
          self.x_batches = np.split(x_data, self.num_batches)
          self.y_batches = np.split(y_data, self.num_batches)
  
      def reset_batch(self):
          self.index_id = 0
  
      def next_batch(self):  # in test use for range.
          x, y = self.x_batches[self.index_id], \
                 self.y_batches[self.index_id]
          self.index_id += 1
          return x, y
  
      def get_seq_ter_prob(self):
          branch_seq_path = os.path.join(self.data_path, "branch_dataset.brd")
          seqs = self.sequences_loader(branch_seq_path)
          seq_len_list = [len(seq) for seq in seqs]
          seq_len_counter = collections.Counter(seq_len_list)
  
          t_prob = {}  # A dict for terminate prob.
          for l, _ in seq_len_counter.items():
              c = sum([c_ for l_, c_ in seq_len_counter.items() if l_ == l])
              # rest longer one
              r = 0  # rest longer one
              for l_, c_ in seq_len_counter.items():
                  r_ = 0
                  if l_ > l:  # for longer one
                      for i in range(l_ - l):
                          c_ = c_ / 2  # get rid of repetitive sub tree
                      r_ += c_
                  r += r_
  
              t_prob[l] = float(c) / float(c + r)
  
          print("Got termination prob.")
  
          return t_prob
  
  
  # loader = SeqLoader("C:/Users/Administrator/Desktop/TestFullNeuron", 20, 2*STEM_SAMPLE_POINT_NUM, 50, 3)
  # x, y = loader.next_batch()
  # print(x)
  # print(y)
  
  
  
  
  