text: from __future__ import absolute_import
  from __future__ import division
  from __future__ import print_function
  
  from keras.layers import Lambda
  from keras import backend as K
  import tensorflow as tf
  import numpy as np
  import os
  import matplotlib.pyplot as plt
  from general_utils import plot_3d_point_cloud
  # from evaluate import get_iou, get_emd
  import sys
  sys.path.append("external")
  from in_out import snc_category_to_synth_id, create_dir, PointCloudDataSet, load_all_point_clouds_under_folder
  from general_utils import plot_3d_point_cloud
  # from structural_losses.tf_nndistance import nn_distance
  #from structural_losses.tf_approxmatch import approx_match, match_cost
  # from tf_ops.emd import tf_auctionmatch
  # from tf_ops.sampling import tf_sampling
  # from tf_ops.CD import tf_nndistance
  
  
  def save_weights(model, model_name, weights_dir, latent_dim, prefix='model'):
      weights_dir = os.path.join(os.getcwd(), weights_dir)
      if not os.path.isdir(weights_dir):
          os.makedirs(weights_dir)
      filename = '%s-%s-weights-%d.h5' % (prefix, model_name, latent_dim)
      path = os.path.join(weights_dir, filename)
      model.save_weights(path)
  
      # filename = '%s-%s-model-%d.h5' % (prefix, model_name, latent_dim)
      # path = os.path.join(weights_dir, filename)
      # model.save(path)
  
  
  # reparameterization trick
  # instead of sampling from Q(z|X), sample eps = N(0,I)
  # then z = z_mean + sqrt(var)*eps
  def sampling(args):
      """Reparameterization trick by sampling fr an isotropic unit Gaussian.
  
      # Arguments:
          args (tensor): mean and log of variance of Q(z|X)
  
      # Returns:
          z (tensor): sampled latent vector
      """
  
      z_mean, z_log_var = args
      batch = K.shape(z_mean)[0]
      dim = K.int_shape(z_mean)[1]
      # by default, random_normal has mean=0 and std=1.0
      epsilon = K.random_normal(shape=(batch, dim))
      return z_mean + K.exp(0.5 * z_log_var) * epsilon
  
  
  def save_images(images, datasource, gray=False, folder="tmp"):
      path = os.path.dirname(os.path.realpath(__file__))
      path = os.path.join(path, folder)
      i = 0
      for image in images:
          filename = "image-%d.png" % i
          image_path = os.path.join(path, filename)
          if gray:
              shape = image.shape 
              plt.imsave(image_path, np.reshape(image, (shape[0], shape[0])), cmap='gray')
          else:
              datasource.save_img(image, image_path)
          i += 1
  
  
  def plot_vox(vox, title="vox"):
      fig = plt.figure()
      plt.title(title)
      ax = fig.gca(projection='3d')
      ax.voxels(vox, edgecolor='b')
      ax.view_init(azim=240, elev=10)
      plt.show()
      # plt.close('all')
  
  
  def ptcloud_to_vox(ptcloud, dim=32):
      vox = np.zeros([dim, dim, dim])
      for idx in ptcloud:
          vox[tuple(idx)] = 1
  
      return vox
  
  def get_all_iou(posterior_net, ptcloud_ae):
      im, pc, vol = posterior_net.test_source.next_batch()
  
  
  def plot_image2ptcloud_results(image2ptcloud):
      im, pc, vol = image2ptcloud.test_source.next_batch()
      print("pc: ", pc.shape)
      save_images(im, posterior_net.test_source)
      
      pre = image2ptcloud.predict(im)
      pre *= 0.5
      iou = 0.0
      for i in range(len(latent)):
          grd = pc[i]
          print("Ground: ", i)
          plot_3d_point_cloud(grd[:, 0], 
                              grd[:, 1], 
                              grd[:, 2], title='point cloud', in_u_sphere=True);
  
          print("Prediction: ", i)
          plot_3d_point_cloud(pre[i][:, 0], 
                              pre[i][:, 1], 
                              pre[i][:, 2], title='point cloud', in_u_sphere=True);
  
  
  def get_chamfer(gt, pred, batch_size=32, version=1):
      if version==0:
          p1top2, _, p2top1, _ = nn_distance(pred, gt)
          cd = K.mean(p1top2, axis=1) + K.mean(p2top1, axis=1)
          return K.eval(cd)
  
      p1top2 , _, p2top1, _ = tf_nndistance.nn_distance(gt, pred)
      #p1top2 is for each element in gt, the cloest distance to this element
      cd = p1top2 #  + p2top1
      cd = K.mean(cd, axis=1)
      return K.eval(cd)
  
  
  def get_emd(gt, pred, batch_size=32, version=1):
       matchl_out, matchr_out = tf_auctionmatch.auction_match(pred, gt)
       matched_out = tf_sampling.gather_point(gt, matchl_out)
       emd = tf.reshape((pred - matched_out) ** 2, shape=(batch_size, -1))
       emd = tf.reduce_mean(emd, axis=1, keepdims=True)
       emd = K.eval(emd)
       return emd
  
  
  def eval_posterior_net_results(posterior_net, ptcloud_ae, emd_metric=True, version=1):
      metric = None
      batch = 0
      cost = get_emd
      if not emd_metric:
          cost = get_chamfer
      while True:
          im, pc, view = posterior_net.test_source.next_batch()
          if im is None:
              break
  
          batch += posterior_net.batch_size
          latent = posterior_net.posterior_net.predict(im)
          pred = ptcloud_ae.decoder.predict(latent)
          # ground truth is [-0.5, 0.5]
          # prediction is [-1.0, 1.0]
          # rescale
          pred *= 0.5
          metric_ = cost(pc, pred, version=version)
          if metric is None:
              metric = metric_
          else:
              metric = np.append(metric, metric_, axis=0)
          if emd_metric:
              print("EMD", metric.shape, ": ", metric.mean())
          else:
              print("Sqrt Chamfer ver", str(version), metric.shape[0], ": ", np.sqrt(metric).mean())
  
      if emd_metric:
          print("Mean EMD: ", metric.mean())
      else:
          print("Mean Sqrt Chamfer: ", np.sqrt(metric).mean())
  
  
  def plot_posterior_net_results(posterior_net, ptcloud_ae=None, image_ae=None, im2pc=True):
      im, pc, view = posterior_net.test_source.next_batch()
      print("pc: ", pc.shape)
      # save_images(im, posterior_net.test_source)
     
      if im2pc:
          latent = posterior_net.posterior_net.predict(im)
          pre = ptcloud_ae.decoder.predict(latent)
          pre *= 0.5
          cost = get_chamfer(pc, pre)
          cost = np.sqrt(cost)
          print("Batch chamfer: ", cost.mean())
          for i in range(len(latent)):
              grd = pc[i]
              print("Ground: ", i)
              plot_3d_point_cloud(grd[:, 0], 
                                  grd[:, 1], 
                                  grd[:, 2], title='point cloud', in_u_sphere=True);
  
              print("Prediction: ", i)
              print("Chamfer: ", cost[i])
              plot_3d_point_cloud(pre[i][:, 0], 
                                  pre[i][:, 1], 
                                  pre[i][:, 2], title='point cloud', in_u_sphere=True);
      else:
          latent, v = posterior_net.posterior_net.predict([pc, view])
          pre  = image_ae.decoder.predict(latent)
          for i in range(len(latent)):
              grd = pc[i]
              print("Prediction: ", i)
              print("View ground: ", view[i], " View pred: ", v[i])
              img = pre[i]
              s = img.shape
              img = np.reshape(img, (s[0], s[0]))
              save_image(img, i, folder="predictions")
              plot_3d_point_cloud(grd[:, 0], 
                                  grd[:, 1], 
                                  grd[:, 2], title='point cloud', in_u_sphere=True);
  
  
  def save_image(image, index, folder="tmp"):
      path = os.path.dirname(os.path.realpath(__file__))
      path = os.path.join(path, folder)
      filename = "image-%d.png" % index
      image_path = os.path.join(path, filename)
      print(image_path)
      plt.imsave(image_path, image, cmap='gray')
      # datasource.save_img(image, image_path, option='gray')
  
  
  def plot_decoded_ptcloud(ae, test_source):
  
      im, pc = test_source.next_batch()
      data_dir = 'ptcloud_out'
      save_dir = os.path.join(os.getcwd(), data_dir)
      if not os.path.isdir(save_dir):
          os.makedirs(save_dir)
  
      for i in range(len(pc)):
          test = pc[i]
          plot_3d_point_cloud(test[:, 0], 
                              test[:, 1], 
                              test[:, 2], in_u_sphere=True);
  
          test = pc[i] / 0.5
          shape = (1,) + test.shape
          test = np.reshape(test, shape)
          reconstruction = 0.5 * ae.predict(test)
          plot_3d_point_cloud(reconstruction[0][:, 0], 
                              reconstruction[0][:, 1], 
                              reconstruction[0][:, 2], in_u_sphere=True);
  
  
  def plot_ae_results(models,
                      test_data,
                      gray=True,
                      model_name="ae_test"):
  
      encoder, decoder, ae = models
      data_dir = 'predictions'
      save_dir = os.path.join(os.getcwd(), data_dir)
      if not os.path.isdir(save_dir):
          os.makedirs(save_dir)
  
      outputs = ae.predict(test_data)
      nitems = outputs.shape[0]
  
      for i in range(nitems):
          filename = 'pred-ae-%d.png' % i
          path = os.path.join(data_dir, filename)
          if gray:
              shape = outputs[i].shape 
              plt.imsave(path, np.reshape(outputs[i], (shape[0], shape[0])), cmap='gray')
          else:
              plt.imsave(path, outputs[i])
  
          filename = 'grnd-ae-%d.png' % i
          if gray:
              plt.imsave(path, np.reshape(test_data[i], (shape[0], shape[0])), cmap='gray')
          else:
              path = os.path.join(data_dir, filename)
          plt.imsave(path, test_data[i])
  
  
  def plot_vae_results(models, test_data, latent_dim=1024, gray=False, data_dir="predictions"):
  
      encoder, decoder, vae = models
      save_dir = os.path.join(os.getcwd(), data_dir)
      if not os.path.isdir(save_dir):
          os.makedirs(save_dir)
  
      outputs = vae.predict(test_data)
      nitems = outputs.shape[0]
  
      for i in range(nitems):
          filename = 'pred-vae-%d.png' % i
          path = os.path.join(data_dir, filename)
          # plt.imsave(path, outputs[i])
          image = outputs[i]
          shape = image.shape 
          plt.imsave(path, np.reshape(image, (shape[0], shape[0])), cmap='gray')
  
          filename = 'grnd-vae-%d.png' % i
          path = os.path.join(data_dir, filename)
          # plt.imsave(path, test_data[i])
          image = test_data[i]
          plt.imsave(path, np.reshape(image, (shape[0], shape[0])), cmap='gray')
  
          lshape = (1, latent_dim)
          z_sample = np.random.uniform(-1, 1, lshape)
          reco = decoder.predict(z_sample)
          filename = 'reco-vae-%d.png' % i
          path = os.path.join(data_dir, filename)
          # plt.imsave(path, reco[0])
          image = reco[0]
          plt.imsave(path, np.reshape(image, (shape[0], shape[0])), cmap='gray')
  