text: from __future__ import print_function
  
  import numpy as np
  import tensorflow as tf
  from rdkit import Chem
  from rdkit.Chem import AllChem
  import tftraj.rmsd as rmsd
  import copy
  from tensorboardX import SummaryWriter
  from tf_rmsd import tf_centroid, tf_centroid_masked, tf_kabsch_rmsd_masked, tf_kabsch_rmsd
  import pdb
  import rmsd
  import glob
  import copy
  import os
  import shutil
  import pickle as pkl
  
  class Model(object):
  
      def __init__(self, data, n_max, dim_node, dim_edge, dim_h, dim_f, \
                  batch_size, val_num_samples, \
                  mpnn_steps=5, alignment_type='default', tol=1e-5, \
                  use_X=True, use_R=True, virtual_node=False, seed=0, \
                  refine_steps=0, refine_mom=0.99, prior_T=1):
  
          # set random seed
          np.random.seed(seed)
          tf.set_random_seed(seed)
  
          # hyper-parameters
          self.data = data
          self.mpnn_steps = mpnn_steps
          self.n_max, self.dim_node, self.dim_edge, self.dim_h, self.dim_f, self.batch_size = n_max, dim_node, dim_edge, dim_h, dim_f, batch_size
          self.val_num_samples = val_num_samples
          self.tol = tol
          self.virtual_node = virtual_node
          self.refine_steps = refine_steps
          self.refine_mom = refine_mom
          self.use_X = use_X
          self.use_R = use_R
  
          if alignment_type == 'linear':
              self.msd_func = self.linear_transform_msd
          elif alignment_type == 'kabsch':
              self.msd_func = self.kabsch_msd
          elif alignment_type == 'default':
              self.msd_func = self.mol_msd
  
          # variables
          self.G = tf.Graph()
          self.G.as_default()
  
          # allow the tensorflow graph to be flexible for number of samples in batch
          # will be useful for validation when we use multiple samples
          self.node = tf.placeholder(tf.float32, [self.batch_size, self.n_max, self.dim_node])
          self.mask = tf.placeholder(tf.float32, [self.batch_size, self.n_max, 1]) # node yes = 1, no = 0
          self.edge = tf.placeholder(tf.float32, [self.batch_size, self.n_max, self.n_max, self.dim_edge])
          self.pos = tf.placeholder(tf.float32, [self.batch_size, self.n_max, 3])
          self.pos_to_proximity = self._pos_to_proximity(self.pos)
          self.proximity = tf.placeholder(tf.float32, [self.batch_size, self.n_max, self.n_max])
          if self.virtual_node:
              self.true_masks = tf.placeholder(tf.float32, [self.batch_size, self.n_max, 1])
              mask = self.true_masks
          else:
              mask = self.mask
          self.trn_flag = tf.placeholder(tf.bool)
  
          self.n_atom = tf.reduce_sum( tf.transpose(self.mask, [0, 2, 1]), 2) #[batch_size, 1]
  
          self.node_embed = self._embed_node(self.node)
          self.edge_2 = tf.concat([self.edge, tf.tile( tf.reshape(self.n_atom, [self.batch_size, 1, 1, 1]), [1, self.n_max, self.n_max, 1] )], 3)
  
          # p(Z|G) -- prior of Z
          self.priorZ_edge_wgt = self._edge_nn(self.edge_2, name = 'priorZ', reuse = False) #[batch_size, n_max, n_max, dim_h, dim_h]
          self.priorZ_hidden = self._MPNN(self.priorZ_edge_wgt, self.node_embed, name = 'priorZ', reuse = False)
          self.priorZ_out = self._g_nn(self.priorZ_hidden, self.node_embed, 2 * self.dim_h, name = 'priorZ', reuse = False)
          self.priorZ_mu, self.priorZ_lsgms = tf.split(self.priorZ_out, [self.dim_h, self.dim_h], 2)
          self.priorZ_sample = self._draw_sample(self.priorZ_mu, self.priorZ_lsgms, T=prior_T)
  
          # q(Z|R(X),G) -- posterior of Z, used R insted of X as input for simplicity, should be updated
          if use_R:
              self.postZ_edge_wgt = self._edge_nn(tf.concat([self.edge_2, tf.reshape(self.proximity, [self.batch_size, self.n_max, self.n_max, 1])], 3),  name = 'postZ', reuse = False)
          else:
              self.postZ_edge_wgt = self._edge_nn(self.edge_2,  name = 'postZ', reuse = False) #[batch_size, n_max, n_max, dim_h, dim_h]
  
          if use_X:
              self.postZ_hidden = self._MPNN(self.postZ_edge_wgt, self._embed_node(tf.concat([self.node, self.pos], 2)), name = 'postZ', reuse = False)
          else:
              self.postZ_hidden = self._MPNN(self.postZ_edge_wgt, self.node_embed, name = 'postZ', reuse = False)
  
          self.postZ_out = self._g_nn(self.postZ_hidden, self.node_embed, 2 * self.dim_h, name = 'postZ', reuse = False)
          self.postZ_mu, self.postZ_lsgms = tf.split(self.postZ_out, [self.dim_h, self.dim_h], 2)
          self.postZ_sample = self._draw_sample(self.postZ_mu, self.postZ_lsgms)
  
          # p(X|Z,G) -- posterior of X
          self.X_edge_wgt = self._edge_nn(self.edge_2, name = 'postX', reuse = False) #[batch_size, n_max, n_max, dim_h, dim_h]
          self.X_hidden = self._MPNN(self.X_edge_wgt, self.postZ_sample + self.node_embed, name = 'postX', reuse = False)
          self.X_pred = self._g_nn(self.X_hidden, self.node_embed, 3, name = 'postX', reuse = False, mask=mask)
  
          # p(X|Z,G) -- posterior of X without sampling from latent space
          # used for iterative refinement of predictions
          # det stands for deterministic
          self.X_edge_wgt_det = self._edge_nn(self.edge_2, name = 'postX', reuse = True) #[batch_size, n_max, n_max, dim_h, dim_h]
          self.X_hidden_det = self._MPNN(self.X_edge_wgt_det, self.postZ_mu + self.node_embed, name = 'postX', reuse = True)
          self.X_pred_det = self._g_nn(self.X_hidden_det, self.node_embed, 3, name = 'postX', reuse = True, mask=mask)
  
          # Prediction of X with p(Z|G) in the test phase
          self.PX_edge_wgt = self._edge_nn(self.edge_2, name = 'postX', reuse = True) #[batch_size, n_max, n_max, dim_h, dim_h]
          self.PX_hidden = self._MPNN(self.PX_edge_wgt, self.priorZ_sample + self.node_embed, name = 'postX', reuse = True)
          self.PX_pred = self._g_nn(self.PX_hidden, self.node_embed, 3, name = 'postX', reuse = True, mask=mask)
  
          self.saver = tf.train.Saver()
          self.sess = tf.Session()
  
  
      def test(self, D1_v, D2_v, D3_v, D4_v, D5_v, MS_v, load_path = None, \
                  tm_v=None, debug=False, savepred_path=None, savepermol=False, useFF=False):
          if load_path is not None:
              self.saver.restore( self.sess, load_path )
  
          # val batch size is different from train batch size
          # since we use multiple samples
          val_batch_size = int(self.batch_size / self.val_num_samples)
          n_batch_val = int(len(D1_v)/val_batch_size)
          assert ((self.batch_size % self.val_num_samples) == 0)
          assert (len(D1_v) % val_batch_size == 0)
  
          val_size = D1_v.shape[0]
          valscores_mean = np.zeros(val_size)
          valscores_std = np.zeros(val_size)
  
          if savepred_path != None:
              if not savepermol:
                  pred_v = np.zeros((len(D1_v), self.val_num_samples, self.n_max, 3))
  
          print ("testing model...")
          for i in range(n_batch_val):
              if debug:
                  print (i, n_batch_val)
              start_ = i * val_batch_size
              end_ = start_ + val_batch_size
  
              node_val = np.repeat(D1_v[start_:end_], self.val_num_samples, axis=0)
              mask_val = np.repeat(D2_v[start_:end_], self.val_num_samples, axis=0)
              edge_val = np.repeat(D3_v[start_:end_], self.val_num_samples, axis=0)
              proximity_val = np.repeat(D4_v[start_:end_], self.val_num_samples, axis=0)
  
              dict_val = {self.node: node_val, self.mask:mask_val, self.edge:edge_val, \
                          self.trn_flag: False }
  
              if self.virtual_node:
                  true_masks_val = np.repeat(tm_v[start_:end_], self.val_num_samples, axis=0)
                  dict_val[self.true_masks] = true_masks_val
                  D5_batch = self.sess.run(self.PX_pred, feed_dict=dict_val)
              else:
                  D5_batch = self.sess.run(self.PX_pred, feed_dict=dict_val)
  
              if savepred_path != None:
                  if not savepermol:
                      pred_v[start_:end_] = D5_batch.reshape(val_batch_size, self.val_num_samples, self.n_max, 3)
  
              # iterative refinement of posterior
              D5_batch_pred = copy.deepcopy(D5_batch)
              for r in range(self.refine_steps):
                  if self.use_X:
                      dict_val[self.pos] = D5_batch_pred
                  if self.use_R:
                      pred_proximity = self.sess.run(self.pos_to_proximity, \
                                          feed_dict={self.pos: D5_batch_pred, \
                                                      self.mask:mask_val})
                      dict_val[self.proximity] = pred_proximity
                  D5_batch = self.sess.run(self.X_pred, feed_dict=dict_val)
                  D5_batch_pred = \
                      self.refine_mom * D5_batch_pred + (1-self.refine_mom) * D5_batch
  
              valres=[]
              for j in range(D5_batch_pred.shape[0]):
                  ms_v_index = int(j / self.val_num_samples) + start_
                  res = self.getRMS(MS_v[ms_v_index], D5_batch_pred[j], useFF)
                  valres.append(res)
  
              valres = np.array(valres)
              valres = np.reshape(valres, (val_batch_size, self.val_num_samples))
              valres_mean = np.mean(valres, axis=1)
              valres_std = np.std(valres, axis=1)
  
              valscores_mean[start_:end_] = valres_mean
              valscores_std[start_:end_] = valres_std
  
              # save results per molecule if request
              if savepermol:
                  pred_curr = copy.deepcopy(D5_batch_pred).reshape(val_batch_size, self.val_num_samples, self.n_max, 3)
                  for tt in range(0, val_batch_size):
                      save_dict_tt = {'rmsd': valres[tt], 'pred': pred_curr[tt]}
                      pkl.dump(save_dict_tt, \
                          open(os.path.join(savepred_path, 'mol_{}_neuralnet.p'.format(tt+start_)), 'wb'))
  
          print ("val scores: mean is {} , std is {}".format(np.mean(valscores_mean), np.mean(valscores_std)))
          if savepred_path != None:
              if not savepermol:
                  print ("saving neural net predictions into {}".format(savepred_path))
                  pkl.dump(pred_v, open(savepred_path, 'wb'))
  
          return np.mean(valscores_mean), np.mean(valscores_std)
  
      def getRMS(self, prb_mol, ref_pos, useFF=False):
  
          def optimizeWithFF(mol):
  
              molf = Chem.AddHs(mol, addCoords=True)
              AllChem.MMFFOptimizeMolecule(molf)
              molf = Chem.RemoveHs(molf)
  
              return molf
  
          n_est = prb_mol.GetNumAtoms()
  
          ref_cf = Chem.rdchem.Conformer(n_est)
          for k in range(n_est):
              ref_cf.SetAtomPosition(k, ref_pos[k].tolist())
  
          ref_mol = copy.deepcopy(prb_mol)
          ref_mol.RemoveConformer(0)
          ref_mol.AddConformer(ref_cf)
  
          if useFF:
              try:
                  res = AllChem.AlignMol(prb_mol, optimizeWithFF(ref_mol))
              except:
                  res = AllChem.AlignMol(prb_mol, ref_mol)
          else:
              res = AllChem.AlignMol(prb_mol, ref_mol)
  
          return res
  
  
      def train(self, D1_t, D2_t, D3_t, D4_t, D5_t, MS_t, D1_v, D2_v, D3_v, D4_v, D5_v, MS_v,\
              load_path = None, save_path = None, train_event_path = None, valid_event_path = None,\
              log_train_steps=100, tm_trn=None, tm_val=None, w_reg=1e-3, debug=False, exp=None):
          if exp is not None:
              data_path = exp.get_data_path(exp.name, exp.version)
              save_path = os.path.join(data_path, 'checkpoints/model.ckpt')
              event_path = os.path.join(data_path, 'event/')
              print(save_path, flush=True)
              print(event_path, flush=True)
          # SummaryWriter
          if not debug:
              train_summary_writer = SummaryWriter(train_event_path)
              valid_summary_writer = SummaryWriter(valid_event_path)
  
          # objective functions
          cost_KLDZ = tf.reduce_mean( tf.reduce_sum( self._KLD(self.postZ_mu, self.postZ_lsgms, self.priorZ_mu, self.priorZ_lsgms), [1, 2]) ) # posterior | prior
          cost_KLD0 = tf.reduce_mean( tf.reduce_sum( self._KLD_zero(self.priorZ_mu, self.priorZ_lsgms), [1, 2]) ) # prior | N(0,1)
  
          mask = self.true_masks if self.virtual_node else self.mask
          cost_X = tf.reduce_mean( self.msd_func(self.X_pred, self.pos, mask) )
  
          cost_op = cost_X + cost_KLDZ + w_reg * cost_KLD0 #hyperparameters!
          train_op = tf.train.AdamOptimizer(learning_rate=3e-4).minimize(cost_op)
  
          self.sess.run(tf.global_variables_initializer())
          self.sess.graph.finalize()
          if load_path is not None:
              self.saver.restore( self.sess, load_path )
  
          # session
          n_batch = int(len(D1_t)/self.batch_size)
          n_batch_val = int(len(D1_v)/self.batch_size)
          np.set_printoptions(precision=5, suppress=True)
  
          # training
          print('::: start training')
          num_epochs = 2500
          valaggr_mean = np.zeros(num_epochs)
          valaggr_std = np.zeros(num_epochs)
  
          for epoch in range(num_epochs):
  
              [D1_t, D2_t, D3_t, D4_t, D5_t] = self._permutation([D1_t, D2_t, D3_t, D4_t, D5_t])
  
              trnscores = np.zeros((n_batch, 4))
              for i in range(n_batch):
                  start_ = i * self.batch_size
                  end_ = start_ + self.batch_size
  
                  if self.virtual_node:
                      trnresult = self.sess.run([train_op, cost_op, cost_X, cost_KLDZ, cost_KLD0],
                                                feed_dict={self.node: D1_t[start_:end_], self.mask: D2_t[start_:end_],
                                                           self.edge: D3_t[start_:end_],
                                                           self.proximity: D4_t[start_:end_],
                                                           self.pos: D5_t[start_:end_],
                                                           self.true_masks: tm_trn[start_:end_],
                                                           self.trn_flag: True})
                  else:
                      trnresult = self.sess.run([train_op, cost_op, cost_X, cost_KLDZ, cost_KLD0],
                                      feed_dict = {self.node: D1_t[start_:end_], self.mask: D2_t[start_:end_],
                                                   self.edge: D3_t[start_:end_], self.proximity: D4_t[start_:end_],
                                                   self.pos: D5_t[start_:end_], self.trn_flag: True})
  
                  trnresult = trnresult[1:]
                  if debug:
                      print (i, n_batch)
                      print(trnresult, flush=True)
  
                  # log results
                  curr_iter = epoch * n_batch + i
  
                  if not debug:
                      if curr_iter % log_train_steps == 0:
                          train_summary_writer.add_scalar("train/cost_op", trnresult[0], curr_iter)
                          train_summary_writer.add_scalar("train/cost_X", trnresult[1], curr_iter)
                          train_summary_writer.add_scalar("train/cost_KLDZ", trnresult[2], curr_iter)
                          train_summary_writer.add_scalar("train/cost_KLD0", trnresult[3], curr_iter)
  
                  assert np.sum(np.isnan(trnresult)) == 0
                  trnscores[i,:] = trnresult
              print(np.mean(trnscores,0), flush=True)
              exp_dict = {}
              if exp is not None:
                  exp_dict['training epoch id'] = epoch
                  exp_dict['train_score'] = np.mean(trnscores,0)
  
              valscores_mean, valscores_std = self.test(D1_v, D2_v, D3_v, D4_v, D5_v, MS_v, \
                                              load_path=None, tm_v=tm_val, debug=debug)
  
              valaggr_mean[epoch] = valscores_mean
              valaggr_std[epoch] = valscores_std
  
              if not debug:
                  valid_summary_writer.add_scalar("val/valscores_mean", valscores_mean, epoch)
                  valid_summary_writer.add_scalar("val/min_valscores_mean", np.min(valaggr_mean[0:epoch+1]), epoch)
                  valid_summary_writer.add_scalar("val/valscores_std", valscores_std, epoch)
                  valid_summary_writer.add_scalar("val/min_valscores_std", np.min(valaggr_std[0:epoch+1]), epoch)
  
              #print('::: training epoch id', epoch, ':: --- val : ', np.mean(valscores, 0), '--- min : ', np.min(valaggr[0:epoch+1]), flush=True)
              #print('::: training epoch id', epoch, ':: --- val mean {} std {} : ', valscores_mean, valscores_std, '--- min mean {} std {} : ', np.min(valaggr_mean[0:epoch+1]), flush=True)
              print ('::: training epoch id {} :: --- val mean={} , std={} ; --- best val mean={} , std={} '.format(\
                      epoch, valscores_mean, valscores_std, np.min(valaggr_mean[0:epoch+1]), np.min(valaggr_std[0:epoch+1])))
              if exp is not None:
                  exp_dict['val mean'] = valscores_mean
                  exp_dict['std'] = valscores_std
                  exp_dict['best val mean'] = np.min(valaggr_mean[0:epoch+1])
                  exp_dict['std of best val mean'] = np.min(valaggr_std[0:epoch+1])
                  exp.log(exp_dict)
                  exp.save()
  
              if save_path is not None and not debug:
                  self.saver.save( self.sess, save_path )
              # keep track of the best model as well in the separate checkpoint
              # it is done by copying the checkpoint
              if valaggr_mean[epoch] == np.min(valaggr_mean[0:epoch+1]) and not debug:
                  for ckpt_f in glob.glob(save_path + '*'):
                      model_name_split = ckpt_f.split('/')
                      model_path = '/'.join(model_name_split[:-1])
                      model_name = model_name_split[-1]
                      best_model_name = model_name.split('.')[0] + '_best.' + '.'.join(model_name.split('.')[1:])
                      full_best_model_path = os.path.join(model_path, best_model_name)
                      full_model_path = ckpt_f
                      shutil.copyfile(full_model_path, full_best_model_path)
  
      def do_mask(self, vec, m):
          return tf.boolean_mask(vec, tf.reshape(tf.greater(m, tf.constant(0.5)), [self.n_max,]) )
  
      def kabsch_msd(self, frames, targets, masks):
          losses = []
          for i in range(self.batch_size):
              frame = frames[i]
              target = targets[i]
              mask = masks[i]
              target_cent = target - tf_centroid_masked(target, mask, self.tol)
              frame_cent = frame - tf_centroid_masked(frame, mask, self.tol)
              losses.append(tf_kabsch_rmsd_masked(tf.stop_gradient(target_cent), frame_cent, mask, self.tol))
  
          loss = tf.stack(losses, 0)
          return loss
  
      def mol_msd(self, frames, targets, masks):
          frames -= tf.reduce_mean(frames, axis = 1, keepdims = True)
          targets -= tf.reduce_mean(targets, axis = 1, keepdims = True)
  
          loss = tf.stack([rmsd.squared_deviation( self.do_mask(frames[i], masks[i]), self.do_mask(targets[i], masks[i]) ) for i in range(self.batch_size)], 0)
          return loss / tf.reduce_sum(masks, axis=[1,2])
  
      def linear_transform_msd(self, frames, targets, masks):
          def linearly_transform_frames(padded_frames, padded_targets):
              s, u, v = tf.svd(padded_frames)
              tol = 1e-7
              atol = tf.reduce_max(s) * tol
              s = tf.boolean_mask(s, s > atol)
              s_inv = tf.diag(1. / s)
              pseudo_inverse = tf.matmul(v, tf.matmul(s_inv, u, transpose_b=True))
  
              weight_matrix = tf.matmul(padded_targets, pseudo_inverse)
              transformed_frames = tf.matmul(weight_matrix, padded_frames)
              return transformed_frames
  
          padding = tf.constant([[0, 0], [0, 0], [0, 1]])
          padded_frames = tf.pad(frames, padding, 'constant', constant_values=1)
          padded_targets = tf.pad(targets, padding, 'constant', constant_values=1)
  
          mask_matrices = []
          for i in range(self.batch_size):
              mask_matrix = tf.diag(tf.reshape(masks[i], [-1]))
              mask_matrices.append(mask_matrix)
          #mask_matrix = tf.diag(tf.reshape(masks, [self.batch_size, -1]))
          mask_tensor = tf.stack(mask_matrices)
          masked_frames = tf.matmul(mask_tensor, padded_frames)
          masked_targets = tf.matmul(mask_tensor, padded_targets)
          transformed_frames = []
          for i in range(self.batch_size):
              transformed_frames.append(linearly_transform_frames(masked_frames[i], masked_targets[i]))
          transformed_frames = tf.stack(transformed_frames)
          #transformed_frames = linearly_transform_frames(masked_frames, masked_targets)
          loss = tf.losses.mean_squared_error(transformed_frames, masked_targets)
  
          return loss
  
      def _permutation(self, set):
  
          permid = np.random.permutation(len(set[0]))
          for i in range(len(set)):
              set[i] = set[i][permid]
  
          return set
  
  
      def _draw_sample(self, mu, lsgms, T=1):
  
          epsilon = tf.random_normal(tf.shape(lsgms), 0., 1.)
          sample = tf.multiply(tf.exp(0.5 * lsgms) * T, epsilon)
          sample = tf.add(mu, sample)
          sample = tf.multiply(sample, self.mask)
  
          return sample
  
  
      def _embed_node(self, inp): #[batch_size, n_max, dim_node]
  
          inp = tf.reshape(inp, [self.batch_size * self.n_max, int(inp.shape[2])])
  
          inp = tf.layers.dense(inp, self.dim_h, activation = tf.nn.sigmoid)
          inp = tf.layers.dense(inp, self.dim_h, activation = tf.nn.tanh)
  
          inp = tf.reshape(inp, [self.batch_size, self.n_max, self.dim_h])
          inp = tf.multiply(inp, self.mask)
  
          return inp
  
  
      def _edge_nn(self, inp, name='', reuse=True): #[batch_size, n_max, n_max, dim_edge]
  
          with tf.variable_scope('edge_nn'+name, reuse=reuse):
  
              inp = tf.reshape(inp, [self.batch_size * self.n_max * self.n_max, int(inp.shape[3])])
  
              inp = tf.layers.dense(inp, 2 * self.dim_h, activation = tf.nn.sigmoid)
              inp = tf.layers.dense(inp, self.dim_h * self.dim_h, activation = tf.nn.tanh)
  
              inp = tf.reshape(inp, [self.batch_size, self.n_max, self.n_max, self.dim_h, self.dim_h])
  
          return inp
  
  
      def _msg_nn(self, wgt, node, name='', reuse=True):
  
          wgt = tf.reshape(wgt, [self.batch_size * self.n_max, self.n_max * self.dim_h, self.dim_h])
          node = tf.reshape(node, [self.batch_size * self.n_max, self.dim_h, 1])
  
          msg = tf.matmul(wgt, node)
          msg = tf.reshape(msg, [self.batch_size, self.n_max, self.n_max, self.dim_h])
          msg = tf.transpose(msg, perm = [0, 2, 3, 1])
          msg = tf.reduce_mean(msg, 3) / self.n_max
  
          return msg
  
  
      def _update_GRU(self, msg, node, name='', reuse=True, mask=None):
  
          if mask is None: mask=self.mask
          with tf.variable_scope('update_GRU'+name, reuse=reuse):
  
              msg = tf.reshape(msg, [self.batch_size * self.n_max, 1, self.dim_h])
              node = tf.reshape(node, [self.batch_size * self.n_max, self.dim_h])
  
              cell = tf.nn.rnn_cell.GRUCell(self.dim_h)
              _, node_next = tf.nn.dynamic_rnn(cell, msg, initial_state = node)
  
              node_next = tf.reshape(node_next, [self.batch_size, self.n_max, self.dim_h])
              node_next = tf.multiply(node_next, mask)
  
          return node_next
  
  
      def _MPNN(self, edge_wgt, node_hidden_0, name='', reuse=True, true_mask=False):
  
          for i in range(self.mpnn_steps): #hyperparameters!
  
              mv_0 = self._msg_nn(edge_wgt, node_hidden_0)
              if true_mask and i == self.mpnn_steps - 1:
                  node_hidden_0 = self._update_GRU(mv_0, node_hidden_0, name=name, reuse=(i + reuse) != 0, mask=self.true_masks)
              else:
                  node_hidden_0 = self._update_GRU(mv_0, node_hidden_0, name=name, reuse=(i+reuse)!=0)#[batch_size, n_max, dim_h]
  
          return node_hidden_0
  
  
      def _g_nn(self, inp, node, outdim, name='', reuse=True, mask=None): #[batch_size, n_max, -]
  
          if mask is None: mask = self.mask
          with tf.variable_scope('g_nn'+name, reuse=reuse):
  
              inp = tf.concat([inp, node], 2)
  
              inp = tf.reshape(inp, [self.batch_size * self.n_max, int(inp.shape[2])])
              inp = tf.layers.dropout(inp, rate = 0.2, training = self.trn_flag)
              inp = tf.layers.dense(inp, self.dim_f, activation = tf.nn.sigmoid)
              inp = tf.layers.dropout(inp, rate = 0.2, training = self.trn_flag)
              #inp = tf.layers.dense(inp, self.dim_f, activation = tf.nn.sigmoid)
              inp = tf.layers.dense(inp, outdim)
  
              inp = tf.reshape(inp, [self.batch_size, self.n_max, outdim])
              inp = tf.multiply(inp, mask)
  
          return inp
  
  
      def _pos_to_proximity(self, pos, reuse=True): #[batch_size, n_max, 3]
  
          with tf.variable_scope('pos_to_proximity', reuse=reuse):
  
              pos_1 = tf.expand_dims(pos, axis = 2)
              pos_2 = tf.expand_dims(pos, axis = 1)
  
              pos_sub = tf.subtract(pos_1, pos_2)
              proximity = tf.square(pos_sub)
              proximity = tf.reduce_sum(proximity, 3)
              proximity = tf.sqrt(proximity + 1e-5)
  
              proximity = tf.reshape(proximity, [self.batch_size, self.n_max, self.n_max])
              proximity = tf.multiply(proximity, self.mask)
              proximity = tf.multiply(proximity, tf.transpose(self.mask, perm = [0, 2, 1]))
  
              proximity = tf.matrix_set_diag(proximity, [[0] * self.n_max] * self.batch_size)
  
          return proximity
  
  
      def _KLD(self, mu0, lsgm0, mu1, lsgm1):# [batch_size, n_max, dim_h]
  
          var0 = tf.exp(lsgm0)
          var1 = tf.exp(lsgm1)
          a = tf.div( var0 + 1e-5, var1 + 1e-5)
          b = tf.div( tf.square( tf.subtract(mu1, mu0) ), var1 + 1e-5)
          c = tf.log( tf.div(var1 + 1e-5, var0 + 1e-5 ) + 1e-5)
  
          kld = 0.5 * tf.reduce_sum(a + b - 1 + c, 2, keepdims = True) * self.mask
  
          return kld
  
  
      def _KLD_zero(self, mu0, lsgm0):# [batch_size, n_max, dim_h]
  
          a = tf.exp(lsgm0) + tf.square(mu0)
          b = 1 + lsgm0
  
          kld = 0.5 * tf.reduce_sum(a - b, 2, keepdims = True) * self.mask
  
          return kld
  