# -*- coding: utf-8 -*-
"""
Created on Thu Jul 23 16:44:22 2020

@author: Administrator
"""

import numpy as np
#import tensorflow as tf
import tensorflow.compat.v1 as tf
tf.disable_v2_behavior()
import os
import shutil
import random
import math
import scipy.io as sio
import time
import argparse
#from im2mesh.utils import libmcubes
import trimesh
from scipy.spatial import cKDTree
from plyfile import PlyData
from plyfile import PlyElement
from skimage.measure import marching_cubes_lewiner

parser = argparse.ArgumentParser()
parser.add_argument('--train',action='store_true', default=False)
parser.add_argument('--finetune',action='store_true', default=False)
parser.add_argument('--test',action='store_true', default=False)
parser.add_argument("--save_idx", type=int, default=-1)
parser.add_argument('--input_ply_file', type=str, default="test.ply")
parser.add_argument('--data_dir', type=str, default="test.ply")
parser.add_argument('--CUDA', type=int, default=0)
parser.add_argument('--OUTPUT_DIR_LOCAL', type=str, default="test.ply")
parser.add_argument('--OUTPUT_DIR_GLOBAL', type=str, default="test.ply")
a = parser.parse_args()

cuda_idx = str(a.CUDA)
os.environ["CUDA_DEVICE_ORDER"]="PCI_BUS_ID"
os.environ["CUDA_VISIBLE_DEVICES"]= cuda_idx

class_idx = '03211117'
name = 'totempole'
BS = 1
primitives = 1
POINT_NUM = 400
POINT_NUM_GT = 10000

part_vox_size = 6
OUTPUT_DIR = a.OUTPUT_DIR_LOCAL
OUTPUT_DIR_FINETUNE = a.OUTPUT_DIR_GLOBAL
LR = 0.0001
START = 0
SHAPE_NUM = 8000
BD_EMPTY = 0.05
TRAIN = a.train
bd = 0.55

if(TRAIN):
    if os.path.exists(OUTPUT_DIR):
        shutil.rmtree(OUTPUT_DIR)
        print ('test_res_dir: deleted and then created!')
    os.makedirs(OUTPUT_DIR)
    if os.path.exists(OUTPUT_DIR_FINETUNE):
        shutil.rmtree(OUTPUT_DIR_FINETUNE)
        print ('test_res_dir: deleted and then created!')
    os.makedirs(OUTPUT_DIR_FINETUNE)

def normal_points(ps_gt, ps, translation = False):
    tt =  0
    if((np.max(ps_gt[:,0])-np.min(ps_gt[:,0])))>(np.max(ps_gt[:,1])-np.min(ps_gt[:,1])):
        tt = (np.max(ps_gt[:,0])-np.min(ps_gt[:,0]))
    else:
        tt = (np.max(ps_gt[:,1])-np.min(ps_gt[:,1]))
    if(tt < (np.max(ps_gt[:,2])-np.min(ps_gt[:,2]))):
        tt = (np.max(ps_gt[:,2])-np.min(ps_gt[:,2]))
    #print('tt:',tt)
    tt = 10/(10*tt)
    ps_gt = ps_gt*tt
    ps = ps*tt
    if(translation):
        t = np.mean(ps_gt,axis = 0)
        ps_gt = ps_gt - t
        ps = ps - t
    #print('normal_gt:',np.max(ps_gt),np.min(ps_gt))
    #print('normal:',np.max(ps),np.min(ps))
    return ps_gt, ps

def fully_connected(inputs,
                    num_outputs,
                    scope,
                    use_xavier=True,
                    stddev=1e-3,
                    weight_decay=0.0,
                    activation_fn=tf.nn.relu,
                    bn=False,
                    bn_decay=None,
                    is_training=None):
  """ Fully connected layer with non-linear operation.

  Args:
    inputs: 2-D tensor BxN
    num_outputs: int

  Returns:
    Variable tensor of size B x num_outputs.
  """
  with tf.variable_scope(scope) as sc:
    num_input_units = inputs.get_shape()[-1].value
    weights = _variable_with_weight_decay('weights',
                                          shape=[num_input_units, num_outputs],
                                          use_xavier=use_xavier,
                                          stddev=stddev,
                                          wd=weight_decay)
    outputs = tf.matmul(inputs, weights)
    biases = _variable_on_cpu('biases', [num_outputs],
                             tf.constant_initializer(0.0))
    outputs = tf.nn.bias_add(outputs, biases)

    if activation_fn is not None:
      outputs = activation_fn(outputs)
    return outputs
def max_pool2d(inputs,
               kernel_size,
               scope,
               stride=[2, 2],
               padding='VALID'):
  """ 2D max pooling.

  Args:
    inputs: 4-D tensor BxHxWxC
    kernel_size: a list of 2 ints
    stride: a list of 2 ints

  Returns:
    Variable tensor
  """
  with tf.variable_scope(scope) as sc:
    kernel_h, kernel_w = kernel_size
    stride_h, stride_w = stride
    outputs = tf.nn.max_pool(inputs,
                             ksize=[1, kernel_h, kernel_w, 1],
                             strides=[1, stride_h, stride_w, 1],
                             padding=padding,
                             name=sc.name)
    return outputs
def _variable_on_cpu(name, shape, initializer, use_fp16=False):
  """Helper to create a Variable stored on CPU memory.
  Args:
    name: name of the variable
    shape: list of ints
    initializer: initializer for Variable
  Returns:
    Variable Tensor
  """
  with tf.device('/cpu:0'):
    dtype = tf.float16 if use_fp16 else tf.float32
    var = tf.get_variable(name, shape, initializer=initializer, dtype=dtype)
  return var

def _variable_with_weight_decay(name, shape, stddev, wd, use_xavier=True):
  """Helper to create an initialized Variable with weight decay.

  Note that the Variable is initialized with a truncated normal distribution.
  A weight decay is added only if one is specified.

  Args:
    name: name of the variable
    shape: list of ints
    stddev: standard deviation of a truncated Gaussian
    wd: add L2Loss weight decay multiplied by this float. If None, weight
        decay is not added for this Variable.
    use_xavier: bool, whether to use xavier initializer

  Returns:
    Variable Tensor
  """
  if use_xavier:
    initializer = tf.contrib.layers.xavier_initializer()
  else:
    initializer = tf.truncated_normal_initializer(stddev=stddev)
  var = _variable_on_cpu(name, shape, initializer)
  if wd is not None:
    weight_decay = tf.multiply(tf.nn.l2_loss(var), wd, name='weight_loss')
    tf.add_to_collection('losses', weight_decay)
  return var

def conv2d(inputs,
           num_output_channels,
           kernel_size,
           scope,
           stride=[1, 1],
           padding='SAME',
           use_xavier=True,
           stddev=1e-3,
           weight_decay=0.0,
           activation_fn=tf.nn.relu,
           bn=False,
           bn_decay=None,
           is_training=None):
  """ 2D convolution with non-linear operation.

  Args:
    inputs: 4-D tensor variable BxHxWxC
    num_output_channels: int
    kernel_size: a list of 2 ints
    scope: string
    stride: a list of 2 ints
    padding: 'SAME' or 'VALID'
    use_xavier: bool, use xavier_initializer if true
    stddev: float, stddev for truncated_normal init
    weight_decay: float
    activation_fn: function
    bn: bool, whether to use batch norm
    bn_decay: float or float tensor variable in [0,1]
    is_training: bool Tensor variable

  Returns:
    Variable tensor
  """
  with tf.variable_scope(scope) as sc:
      kernel_h, kernel_w = kernel_size
      num_in_channels = inputs.get_shape()[-1].value
      kernel_shape = [kernel_h, kernel_w,
                      num_in_channels, num_output_channels]
      kernel = _variable_with_weight_decay('weights',
                                           shape=kernel_shape,
                                           use_xavier=use_xavier,
                                           stddev=stddev,
                                           wd=weight_decay)
      stride_h, stride_w = stride
      outputs = tf.nn.conv2d(inputs, kernel,
                             [1, stride_h, stride_w, 1],
                             padding=padding)
      biases = _variable_on_cpu('biases', [num_output_channels],
                                tf.constant_initializer(0.0))
      outputs = tf.nn.bias_add(outputs, biases)

      if activation_fn is not None:
        outputs = activation_fn(outputs)
      return outputs

def distance_p2p(points_src, normals_src, points_tgt, normals_tgt):
    ''' Computes minimal distances of each point in points_src to points_tgt.

    Args:
        points_src (numpy array): source points
        normals_src (numpy array): source normals
        points_tgt (numpy array): target points
        normals_tgt (numpy array): target normals
    '''
    kdtree = KDTree(points_tgt)
    dist, idx = kdtree.query(points_src)

    if normals_src is not None and normals_tgt is not None:
        normals_src = \
            normals_src / np.linalg.norm(normals_src, axis=-1, keepdims=True)
        normals_tgt = \
            normals_tgt / np.linalg.norm(normals_tgt, axis=-1, keepdims=True)

        normals_dot_product = (normals_tgt[idx] * normals_src).sum(axis=-1)
        # Handle normals that point into wrong direction gracefully
        # (mostly due to mehtod not caring about this in generation)
        normals_dot_product = np.abs(normals_dot_product)
    else:
        normals_dot_product = np.array(
            [np.nan] * points_src.shape[0], dtype=np.float32)
    return dist, normals_dot_product

def eval_pointcloud(pointcloud, pointcloud_tgt,
                        normals=None, normals_tgt=None):
        ''' Evaluates a point cloud.

        Args:
            pointcloud (numpy array): predicted point cloud
            pointcloud_tgt (numpy array): target point cloud
            normals (numpy array): predicted normals
            normals_tgt (numpy array): target normals
        '''
        # Return maximum losses if pointcloud is empty

        pointcloud = np.asarray(pointcloud)
        pointcloud_tgt = np.asarray(pointcloud_tgt)

        # Completeness: how far are the points of the target point cloud
        # from thre predicted point cloud
        completeness, completeness_normals = distance_p2p(
            pointcloud_tgt, normals_tgt, pointcloud, normals
        )
        completeness2 = completeness**2

        completeness = completeness.mean()
        completeness2 = completeness2.mean()
        completeness_normals = np.absolute(completeness_normals).mean()

        # Accuracy: how far are th points of the predicted pointcloud
        # from the target pointcloud
        accuracy, accuracy_normals = distance_p2p(
            pointcloud, normals, pointcloud_tgt, normals_tgt
        )

        accuracy2 = accuracy**2

        accuracy = accuracy.mean()
        accuracy2 = accuracy2.mean()
        accuracy_normals = np.absolute(accuracy_normals).mean()

        # Chamfer distance
        chamferL2 = 0.5 * (completeness2 + accuracy2)
        #print(completeness,accuracy,completeness2,accuracy2)
        #print('chamferL2:',chamferL2)
        normals_correctness = (
            0.5 * completeness_normals + 0.5 * accuracy_normals
        )
        chamferL1 = 0.5 * (completeness + accuracy)
        print('chamferL2:',chamferL2,'accuracy:',accuracy,'normals_correctness:',normals_correctness,'chamferL1:',chamferL1)
        return normals_correctness, chamferL1, chamferL2

def safe_norm_np(x, epsilon=1e-12, axis=1):
    return np.sqrt(np.sum(x*x, axis=axis) + epsilon)

def safe_norm(x, epsilon=1e-12, axis=None):
  return tf.sqrt(tf.reduce_sum(x ** 2, axis=axis) + epsilon)
  #return tf.reduce_sum(x ** 2, axis=axis)

def boundingbox(x,y,z):
    return min(x),max(x),min(y),max(y),min(z),max(z)

def get_data_from_filename(filename):
    load_data = np.load(filename)
    point = np.asarray(load_data['sample_near']).reshape(-1,POINT_NUM,3)
    sample = np.asarray(load_data['sample']).reshape(-1,POINT_NUM,3)
    rt = random.randint(0,sample.shape[0]-1)
    #rt = random.randint(0,int((sample.shape[0]-1)/5))
    sample = sample[rt,:,:].reshape(BS, POINT_NUM, 3)
    point = point[rt,:,:].reshape(BS, POINT_NUM, 3)

    #print('input_points_bs:',filename)
    #print(input_points_bs)
    return point.astype(np.float32), sample.astype(np.float32)

def sample_query_points(input_ply_file):
    data = PlyData.read(a.data_dir + input_ply_file)
    v = data['vertex'].data
    v = np.asarray(v)
    print(v.shape)

    #rt = np.random.choice(v.shape, 50000, replace = False)

    points = []
    for i in range(v.shape[0]):
        points.append(np.array([v[i][0],v[i][1],v[i][2]]))
    points = np.asarray(points)
    pointcloud_s =points.astype(np.float32)
    print('pointcloud sparse:',pointcloud_s.shape[0])

    pointcloud_s_t = pointcloud_s - np.array([np.min(pointcloud_s[:,0]),np.min(pointcloud_s[:,1]),np.min(pointcloud_s[:,2])])
    pointcloud_s_t = pointcloud_s_t / (np.array([np.max(pointcloud_s[:,0]) - np.min(pointcloud_s[:,0]), np.max(pointcloud_s[:,0]) - np.min(pointcloud_s[:,0]), np.max(pointcloud_s[:,0]) - np.min(pointcloud_s[:,0])]))
    trans = np.array([np.min(pointcloud_s[:,0]),np.min(pointcloud_s[:,1]),np.min(pointcloud_s[:,2])])
    scal = np.array([np.max(pointcloud_s[:,0]) - np.min(pointcloud_s[:,0]), np.max(pointcloud_s[:,0]) - np.min(pointcloud_s[:,0]), np.max(pointcloud_s[:,0]) - np.min(pointcloud_s[:,0])])
    pointcloud_s = pointcloud_s_t

    print(np.min(pointcloud_s[:,0]), np.max(pointcloud_s[:,0]))
    print(np.min(pointcloud_s[:,1]), np.max(pointcloud_s[:,1]))
    print(np.min(pointcloud_s[:,2]), np.max(pointcloud_s[:,2]))

    sample = []
    sample_near = []
    sample_near_o = []
    sample_dis = []
    sample_vec = []
    gt_kd_tree = cKDTree(pointcloud_s)
    for i in range(int(500000/pointcloud_s.shape[0])):

        pnts = pointcloud_s
        ptree = cKDTree(pnts)
        i = 0
        sigmas = []
        for p in np.array_split(pnts,100,axis=0):
            d = ptree.query(p,51)
            sigmas.append(d[0][:,-1])

            i = i+1

        sigmas = np.concatenate(sigmas)
        sigmas_big = 0.2 * np.ones_like(sigmas)
        sigmas = sigmas

        #tt = pnts + 0.5*0.25*np.expand_dims(sigmas,-1) * np.random.normal(0.0, 1.0, size=pnts.shape)
        tt = pnts + 0.5*np.expand_dims(sigmas,-1) * np.random.normal(0.0, 1.0, size=pnts.shape)
        #tt = pnts + 1*np.expand_dims(sigmas_big,-1) * np.random.normal(0.0, 1.0, size=pnts.shape)
        sample.append(tt)
        distances, vertex_ids = gt_kd_tree.query(tt, p=2, k = 1)

        vertex_ids = np.asarray(vertex_ids)
        print('distances:',distances.shape)
        #print(vertex_ids)

        sample_near.append(pointcloud_s[vertex_ids].reshape(-1,3))
    for i in range(int(500000/pointcloud_s.shape[0])):

        pnts = pointcloud_s
        ptree = cKDTree(pnts)
        i = 0
        sigmas = []
        for p in np.array_split(pnts,100,axis=0):
            d = ptree.query(p,51)
            sigmas.append(d[0][:,-1])

            i = i+1

        sigmas = np.concatenate(sigmas)
        sigmas_big = 0.2 * np.ones_like(sigmas)
        sigmas = sigmas

        #tt = pnts + 0.5*0.25*np.expand_dims(sigmas,-1) * np.random.normal(0.0, 1.0, size=pnts.shape)
        tt = pnts + 1.0*np.expand_dims(sigmas,-1) * np.random.normal(0.0, 1.0, size=pnts.shape)
        #tt = pnts + 1*np.expand_dims(sigmas_big,-1) * np.random.normal(0.0, 1.0, size=pnts.shape)
        sample.append(tt)
        distances, vertex_ids = gt_kd_tree.query(tt, p=2, k = 1)

        vertex_ids = np.asarray(vertex_ids)
        print('distances:',distances.shape)
        #print(vertex_ids)

        sample_near.append(pointcloud_s[vertex_ids].reshape(-1,3))

    sample = np.asarray(sample).reshape(-1,3)
    sample_near = np.asarray(sample_near).reshape(-1,3)
    np.savez_compressed(a.data_dir + input_ply_file , sample = sample, sample_near=sample_near,pointcloud_s = pointcloud_s, trans = trans, scal = scal)
    sample_all = sample.reshape(-1,3)
    sample_near_all = sample_near.reshape(-1,3)
    sample_part = [[] for i in range(part_vox_size*part_vox_size*part_vox_size)]
    sample_near_part = [[] for i in range(part_vox_size*part_vox_size*part_vox_size)]
    bd_max_x = np.max(pointcloud_s[:,0])
    bd_max_y = np.max(pointcloud_s[:,1])
    bd_max_z = np.max(pointcloud_s[:,2])
    bd_min_x = np.min(pointcloud_s[:,0])
    bd_min_y = np.min(pointcloud_s[:,1])
    bd_min_z = np.min(pointcloud_s[:,2])
    for l in range(sample_near_all.shape[0]):
        ex = sample_near_all[l,0] - bd_min_x
        ix = int(math.floor(ex/((bd_max_x- bd_min_x)/(part_vox_size))))
        #print(ex,ix)
        ey = sample_near_all[l,1] - bd_min_y
        iy = int(math.floor(ey/((bd_max_y- bd_min_y)/(part_vox_size))))
        ez = sample_near_all[l,2] - bd_min_z
        iz = int(math.floor(ez/((bd_max_z- bd_min_z)/(part_vox_size))))
        ix = np.clip(ix,0,part_vox_size-1)
        iy = np.clip(iy,0,part_vox_size-1)
        iz = np.clip(iz,0,part_vox_size-1)
        #print(ix,iy,iz)
        sample_part[ix*(part_vox_size)*(part_vox_size)+iy*(part_vox_size)+iz].append(sample_all[l])
        sample_near_part[ix*(part_vox_size)*(part_vox_size)+iy*(part_vox_size)+iz].append(sample_near_all[l])

    for iv in range(len(sample_part)):
            #print(np.asarray(sample[iv]).shape)
            np.savez(a.data_dir + input_ply_file + '_' + str(iv), sample = sample_part[iv],sample_near = sample_near_part[iv])
if(a.train):
    sample_query_points(a.input_ply_file)
mm = 0
files = []
files_path = []

files.append(a.input_ply_file)
files_path.append(a.data_dir + a.input_ply_file)

points_all = []
samples_all = []

if(a.train):
    for fi in range(len(files_path)):
        print(files_path[fi])
        for i in range(part_vox_size*part_vox_size*part_vox_size):
        #for i in range(100):
            if(os.path.exists(files_path[fi] + '_{}.npz'.format(i))):
                print(i)
                load_data = np.load(files_path[fi] + '_{}.npz'.format(i))
                sample_near = np.asarray(load_data['sample_near'])
                sampler = np.asarray(load_data['sample'])
                print(sample_near.shape[0])
                if(sample_near.shape[0]>=POINT_NUM):
                    print(sample_near.shape[0])
                    sample_near,sampler = normal_points(sample_near,sampler,True)
                    tt = int(math.floor(sample_near.shape[0]*1.0/POINT_NUM))
                    tt = (tt*POINT_NUM)
                    #print(sample_near[0:tt,:].shape)
                    points_all.append(sample_near[0:tt,:])
                    samples_all.append(sampler[0:tt,:])
            #print(points_all[i].shape)

SHAPE_NUM = len(files)
#SHAPE_NUM = 26
print('SHAPE_NUM:',SHAPE_NUM)

points_target = tf.placeholder(tf.float32, shape=[BS,POINT_NUM,3])
input_points_3d = tf.placeholder(tf.float32, shape=[BS,POINT_NUM,3])
normal_gt = tf.placeholder(tf.float32, shape=[BS,None,3])
points_target_num = tf.placeholder(tf.int32, shape=[1,1])
points_input_num = tf.placeholder(tf.int32, shape=[1,1])
points_cd = tf.placeholder(tf.float32, shape=[BS,None,3])

def local_decoder(feature,input_points_3d):
    with tf.variable_scope('local', reuse=tf.AUTO_REUSE):
        feature_f = tf.nn.relu(tf.layers.dense(feature,512))
        net = tf.nn.relu(tf.layers.dense(input_points_3d, 512))
        net = tf.concat([net,feature_f],2)
        print('net:',net)
        with tf.variable_scope('decoder', reuse=tf.AUTO_REUSE):
            for i in range(8):
                with tf.variable_scope("resnetBlockFC_%d" % i ):
                    b_initializer=tf.constant_initializer(0.0)
                    w_initializer = tf.random_normal_initializer(mean=0.0,stddev=np.sqrt(2) / np.sqrt(512))
                    net = tf.layers.dense(tf.nn.relu(net),512,kernel_initializer=w_initializer,bias_initializer=b_initializer)

        b_initializer=tf.constant_initializer(-0.5)
        w_initializer = tf.random_normal_initializer(mean=2*np.sqrt(np.pi) / np.sqrt(512), stddev = 0.000001)
        print('net:',net)
        sdf = tf.layers.dense(tf.nn.relu(net),1,kernel_initializer=w_initializer,bias_initializer=b_initializer)
        print('sdf',sdf)

        grad = tf.gradients(ys=sdf, xs=input_points_3d)
        print('grad',grad)
        print(grad[0])
        normal_p_lenght = tf.expand_dims(safe_norm(grad[0],axis = -1),-1)
        print('normal_p_lenght',normal_p_lenght)
        grad_norm = grad[0]/normal_p_lenght
        print('grad_norm',grad_norm)
        return sdf,grad_norm

input_points_3d_global = tf.placeholder(tf.float32, shape=[BS,None,3])
points_target_global = tf.placeholder(tf.float32, shape=[BS,None,3])
feature_global = tf.placeholder(tf.float32, shape=[BS,None,SHAPE_NUM])
#with tf.variable_scope('pointnet', reuse=tf.AUTO_REUSE):
#    input_image = tf.expand_dims(points_target_global,-1)
#    net = conv2d(input_image, 64, [1,3], padding='VALID', stride = [1,1], is_training = True, scope = 'conv1')
#    net = conv2d(input_image, 128, [1,3], padding='VALID', stride = [1,1], is_training = True, scope = 'conv2')
#    net = conv2d(input_image, 1024, [1,3], padding='VALID', stride = [1,1], is_training = True, scope = 'conv3')
#    net = max_pool2d(net,[POINT_NUM,1], padding = 'VALID', scope = 'maxpool')
#    net = tf.reshape(net,[1,-1])
#    net = fully_connected(net, 512, is_training = True, scope = 'fc1')
#    net = fully_connected(net, 256, is_training = True, scope = 'fc2')
#    feature_global = net
#feature_global = tf.tile(tf.expand_dims(feature_global,1),[1,POINT_NUM,1])
def global_decoder(feature_global_f,input_points_3d_global_f):
    with tf.variable_scope('global', reuse=tf.AUTO_REUSE):
        feature_g = tf.nn.relu(tf.layers.dense(feature_global_f,512))
        net_g = tf.nn.relu(tf.layers.dense(input_points_3d_global_f, 512))
        #print(net_g,feature_g)
        net_g = tf.concat([net_g,feature_g],2)
        for i in range(8):
            net_g = tf.layers.dense(tf.nn.relu(net_g),512)

        feature_output = tf.layers.dense(tf.nn.relu(net_g),SHAPE_NUM)
        d_output = tf.layers.dense(tf.nn.relu(net_g),3)
        sdf_g,grad_norm_g = local_decoder(feature_output,input_points_3d_global_f+d_output)
        g_points_g = input_points_3d_global_f - sdf_g * grad_norm_g
        return g_points_g, sdf_g

g_points_g, sdf_g = global_decoder(feature_global,input_points_3d_global)
loss_g = tf.reduce_mean(tf.norm((points_target_global-g_points_g), axis=-1))

with tf.variable_scope('pointnet', reuse=tf.AUTO_REUSE):
    input_image = tf.expand_dims(points_target,-1)
    net = conv2d(input_image, 64, [1,3], padding='VALID', stride = [1,1], is_training = True, scope = 'conv1')
    net = conv2d(input_image, 128, [1,3], padding='VALID', stride = [1,1], is_training = True, scope = 'conv2')
    net = conv2d(input_image, 1024, [1,3], padding='VALID', stride = [1,1], is_training = True, scope = 'conv3')
    net = max_pool2d(net,[POINT_NUM,1], padding = 'VALID', scope = 'maxpool')
    net = tf.reshape(net,[1,-1])
    net = fully_connected(net, 512, is_training = True, scope = 'fc1')
    net = fully_connected(net, 256, is_training = True, scope = 'fc2')
    feature = net
feature = tf.tile(tf.expand_dims(feature,1),[1,POINT_NUM,1])
sdf,grad_norm = local_decoder(feature,input_points_3d)
g_points = input_points_3d - sdf * grad_norm

loss = tf.reduce_mean(tf.norm((points_target, g_points), axis=-1))

t_vars = tf.trainable_variables()
optim = tf.train.AdamOptimizer(learning_rate=LR, beta1=0.9)
loss_grads_and_vars = optim.compute_gradients(loss, var_list=t_vars)
loss_optim = optim.apply_gradients(loss_grads_and_vars)

global_vars = tf.get_collection(tf.GraphKeys.TRAINABLE_VARIABLES, scope='global')
loss_grads_and_vars_g = optim.compute_gradients(loss_g, var_list=global_vars)
#loss_grads_and_vars_g = optim.compute_gradients(loss_g, var_list=t_vars)
loss_optim_g = optim.apply_gradients(loss_grads_and_vars_g)

config = tf.ConfigProto(allow_soft_placement=False)
saver_restore = tf.train.Saver(var_list=t_vars)
saver = tf.train.Saver(max_to_keep=2000000)

with tf.Session(config=config) as sess:
    feature_bs_all = []
    for i in range(SHAPE_NUM):
        tt = []
        for j in range(int(POINT_NUM)):
            t = np.zeros(SHAPE_NUM)

            t[i] = 1
            tt.append(t)
        feature_bs_all.append(tt)
    feature_bs_all = np.asarray(feature_bs_all)
    #print(feature_bs_all,feature_bs_all.shape)
    if(TRAIN):
        print('train start')
        sess.run(tf.global_variables_initializer())
        start_time = time.time()
        POINT_NUM_GT_bs = np.array(POINT_NUM_GT).reshape(1,1)
        points_input_num_bs = np.array(POINT_NUM).reshape(1,1)
        print('data shape:',len(points_all))
        for bi in range(500):
            epoch_index = np.random.choice(len(points_all)-1, len(points_all)-1, replace = False)
            for epoch in epoch_index:

                points = points_all[epoch].reshape(-1,POINT_NUM,3)
                samples = samples_all[epoch].reshape(-1,POINT_NUM,3)

                rt = random.randint(0,samples.shape[0]-1)
                sample = samples[rt,:].reshape(BS, POINT_NUM, 3)
                point = points[rt,:].reshape(BS, POINT_NUM, 3)

                _, loss_c,g_points_g_c = sess.run([loss_optim,loss,g_points],feed_dict={points_target_num:POINT_NUM_GT_bs,points_input_num:points_input_num_bs,
                                     points_target:point,input_points_3d:sample})

            if(bi%100 == 0):
                print('model:',bi,'epoch:',epoch,'loss:',loss_c)
                saver.save(sess, os.path.join(OUTPUT_DIR, "model"), global_step=bi)

        saver.save(sess, os.path.join(OUTPUT_DIR, "model"), global_step=bi)
    if(a.finetune):
        print('finuetune')
        POINT_NUM_GT_bs = np.array(POINT_NUM_GT).reshape(1,1)
        points_input_num_bs = np.array(POINT_NUM).reshape(1,1)
        points_all = []
        samples_all = []
        for epoch in range(1):
            print('epoch:',epoch)
            sess.run(tf.global_variables_initializer())
            start_time = time.time()
            checkpoint = tf.train.get_checkpoint_state(OUTPUT_DIR).all_model_checkpoint_paths
            print(checkpoint[a.save_idx])

            saver.restore(sess, checkpoint[a.save_idx])
            print(files_path[0] + '.npz')
            load_data = np.load(files_path[0] + '.npz')

            points = load_data['sample_near'].reshape(-1,3)
            samples = load_data['sample'].reshape(-1,3)
            SP_NUM = points.shape[0]
            for bi in range(100010):
                feature_bs = feature_bs_all[0]
                rt = np.random.choice(SP_NUM, POINT_NUM, replace = False)
                #rt = random.randint(0,samples.shape[0]-1)
                sample = samples[rt,:].reshape(BS, POINT_NUM, 3)
                point = points[rt,:].reshape(BS, POINT_NUM, 3)

                feature_bs_t = feature_bs.reshape(BS,POINT_NUM,SHAPE_NUM)
                _, loss_c = sess.run([loss_optim_g,loss_g],feed_dict={feature_global:feature_bs_t,points_target_num:POINT_NUM_GT_bs,points_input_num:points_input_num_bs,
                                     points_target_global:point,input_points_3d_global:sample})

                if(bi%100000 == 0):
                    print('model:',bi,'epoch:',epoch,'loss:',loss_c)
                    saver.save(sess, os.path.join(OUTPUT_DIR_FINETUNE, "model"), global_step=bi)
            #saver.save(sess, os.path.join(OUTPUT_DIR_FINETUNE, "model"), global_step=epoch)

    if(a.test):
        """ feature_bs = []
        for j in range(vox_size*vox_size):
            t = np.zeros(SHAPE_NUM)
            t[0] = 1
            feature_bs.append(t)
        feature_bs = np.asarray(feature_bs)
        sdf_c = sess.run([sdf_g],feed_dict={input_points_3d_global:input_points_2d_bs_t,feature_global:feature_bs_t,
                                 points_target_num:POINT_NUM_GT_bs,points_input_num:points_input_num_bs}) """
        print('test start')

        s = np.arange(-bd,bd, (2*bd)/128)

        print(s.shape[0])
        vox_size = s.shape[0]
        POINT_NUM_GT_bs = np.array(vox_size).reshape(1,1)
        points_input_num_bs = np.array(POINT_NUM).reshape(1,1)

        POINT_NUM_GT_bs = np.array(vox_size*vox_size).reshape(1,1)

        sess.run(tf.global_variables_initializer())
        checkpoint = tf.train.get_checkpoint_state(OUTPUT_DIR_FINETUNE).all_model_checkpoint_paths
        print(checkpoint[a.save_idx])
        saver.restore(sess, checkpoint[a.save_idx])

        #saver.restore(sess, a.out_dir + 'model-0')

        point_sparse = np.load(a.data_dir + a.input_ply_file + '.npz')['pointcloud_s']

        input_points_2d_bs = []

        bd_max = [np.max(point_sparse[:,0]), np.max(point_sparse[:,1]), np.max(point_sparse[:,2])]
        bd_min = [np.min(point_sparse[:,0]), np.min(point_sparse[:,1]),np.min(point_sparse[:,2])]
        bd_max  = np.asarray(bd_max) + 0.05
        bd_min = np.asarray(bd_min) - 0.05
        sx = np.arange(bd_min[0], bd_max[0], (bd_max[0] - bd_min[0])/vox_size)
        sy = np.arange(bd_min[1], bd_max[1], (bd_max[1] - bd_min[1])/vox_size)
        sz = np.arange(bd_min[2], bd_max[2], (bd_max[2] - bd_min[2])/vox_size)
        print(bd_max)
        print(bd_min)
        for i in sx:
            for j in sy:
                for k in sz:
                    input_points_2d_bs.append(np.asarray([i,j,k]))
        input_points_2d_bs = np.asarray(input_points_2d_bs)
        input_points_2d_bs = input_points_2d_bs.reshape((vox_size,vox_size,vox_size,3))

        vox = []
        feature_bs = []
        for j in range(vox_size*vox_size):
            t = np.zeros(SHAPE_NUM)
            t[0] = 1
            feature_bs.append(t)
        feature_bs = np.asarray(feature_bs)
        for i in range(input_points_2d_bs.shape[0]):

            input_points_2d_bs_t = input_points_2d_bs[i,:,:,:]
            input_points_2d_bs_t = input_points_2d_bs_t.reshape(BS, vox_size*vox_size, 3)
            feature_bs_t = feature_bs.reshape(BS,vox_size*vox_size,SHAPE_NUM)
            sdf_c = sess.run([sdf_g],feed_dict={input_points_3d_global:input_points_2d_bs_t,feature_global:feature_bs_t,
                                 points_target_num:POINT_NUM_GT_bs,points_input_num:points_input_num_bs})
            vox.append(sdf_c)

        vox = np.asarray(vox)
        #vis_single_points(moved_points, 'moved_points.ply')
        #print('vox',np.min(vox),np.max(vox),np.mean(vox))
        vox = vox.reshape((vox_size,vox_size,vox_size))
        vox_max = np.max(vox.reshape((-1)))
        vox_min = np.min(vox.reshape((-1)))
        print('max_min:',vox_max,vox_min,np.mean(vox))

        #threshs = [0.001,0.0015,0.002,0.0025,0.005]
        threshs = [0.005]
        for thresh in threshs:
            print(np.sum(vox>thresh),np.sum(vox<thresh))

            if(np.sum(vox>0.0)<np.sum(vox<0.0)):
                thresh = -thresh
            #vertices, triangles = libmcubes.marching_cubes(vox, thresh)
            vertices, triangles, _, _ = marching_cubes_lewiner(vox, thresh)
            if(vertices.shape[0]<10 or triangles.shape[0]<10):
                print('no sur---------------------------------------------')
                continue
            if(np.sum(vox>0.0)>np.sum(vox<0.0)):
                triangles_t = []
                for it in range(triangles.shape[0]):
                    tt = np.array([triangles[it,2],triangles[it,1],triangles[it,0]])
                    triangles_t.append(tt)
                triangles_t = np.asarray(triangles_t)
            else:
                triangles_t = triangles
                triangles_t = np.asarray(triangles_t)

            vertices -= 0.5
            # Undo padding
            vertices -= 1
            # Normalize to bounding box
            vertices /= np.array([vox_size-1, vox_size-1, vox_size-1])
            vertices = (bd_max-bd_min) * vertices + bd_min
            mesh = trimesh.Trimesh(vertices, triangles_t,
                            vertex_normals=None,
                            process=False)

            loc_data = np.load(a.data_dir + a.input_ply_file + '.npz')
            vertices = vertices * loc_data['scal'] + loc_data['trans']
            mesh = trimesh.Trimesh(vertices, triangles_t,
                                vertex_normals=None,
                                process=False)
            mesh.export(OUTPUT_DIR_FINETUNE +  '/PCL_' + a.input_ply_file + '_'+ str(thresh) + '.off')


#!/usr/bin/env python3
# -*- coding: utf-8 -*-
"""
Created: 2018
Author: A. P. Naik
Description: 'data' submodule of spam package. See README for details and usage
examples.

Attributes
----------
SPARCGalaxy : class
    Main object containing all relevant data for each SPARC galaxy.
names_full : list of strings, length 147
    List of names of SPARC galaxies in 'full' sample, i.e. 147 galaxies
    remaining after first 4 data cuts described in Naik et al. (2019).
names_standard : list of strings, length 85
    List of names of SPARC galaxies in 'standard' sample, i.e. 85 galaxies
    remaining after all data cuts described in Naik et al. (2019). Difference
    between 'standard' and 'full' samples are that in the 'standard' case,
    environmentally screened galaxies have additionally been cut from the
    sample.
"""
import os as _os
import numpy as _np
from scipy.constants import parsec as _pc

# physical constants
_kpc = 1e+3*_pc
_Mpc = 1e+6*_pc
_Msun = 1.989e+30

class SPARCGalaxy:
    """
    Class containing all relevant data for a given galaxy. All data come from
    the SPARC database (http://astroweb.cwru.edu/SPARC/), with the following
    exceptions:

    - gas_radius : calculated in spam.data.fit_gas_disc.py
    - hernquist_radius : calculated in spam.data.fit_stellar_bulge.py
    - hernquist_rho_0 : ditto
    - stellar_expdisc_sigma_0 : calculated in spam.data.fit_stellar_disc.py
    - stellar_expdisc_R_d : ditto
    - ext_potential : calculated via the screening map of Desmond et al.
    - ext_potential_lower : ditto
    - ext_potential_upper : ditto

    Parameters
    ----------
    name : str
        Name of galaxy matching name in SPARC database, e.g. F574-1 or CamB.

    Attributes
    ----------
    name : str
        As above.
    hubble_type : str
        Hubble classification of galaxy.
    distance : float
        Distance to galaxy. UNITS: m
    distance_err : float
        Error on distance to galaxy. UNITS: m
    distance_method : int, {1, 2, 3, 4, 5}
        Method used to determine distance to galaxy (see SPARC database for
        meanings of numbers).
    inclination : float
        Inclination of galaxy. UNITS: degrees
    inclination_err : float
        Error on inclination. UNITS: degrees
    luminosity_tot : float
        Total luminosity of galaxy at 3.6mu. UNITS: 10^9 L_sun.
    luminosity_err : float
        Error on total luminosity. UNITS: 10^9 L_sun.
    disc_scale : float
        Scale length of disc fit to photometry data. UNITS: m
    disc_SB : float
        Central surface brightness of disc fit to photometry data. UNITS: m
    HI_mass : float
        Total mass of HI gas. UNITS: kg
    Q_flag : int
        Quality flag (see SPARC database).
    StellarBulge : bool
        Whether a bulge component is detected.
    R : 1D numpy.ndarray
        Radii of rotation curve measurements. UNITS: kpc
    v : 1D numpy.ndarray, shape same as R
        Rotation curve measurements. UNITS: km/s
    v_err : 1D numpy.ndarray, shape same as R
        Errors on rotation curve. UNITS: km/s
    v_gas : 1D numpy.ndarray, shape same as R
        Gas contribution to rotation curve. UNITS: km/s
    v_disc : 1D numpy.ndarray, shape same as R
        Stellar disc contribution to rotation curve, assuming mass-to-light
        ratio of 1 M_sun/L_sun. UNITS: km/s
    v_bul : 1D numpy.ndarray, shape same as R
        Stellar bulge contribution to rotation curve, assuming mass-to-light
        ratio of 1 M_sun/L_sun. Zero everywhere if StellarBulge is False.
        UNITS: km/s
    coords_RA : float
        Right ascension of galaxy. UNITS: degrees
    coords_DEC : float
        Declination of galaxy. UNITS: degrees
    gas_radius : float
        Best fit radius of gas disc, calculated in spam.data.fit_gas_disc.py.
        UNITS: m
    hernquist_radius : float
        Best fit radius of Hernquist bulge, calculated in
        spam.data.fit_stellar_bulge.py. UNITS: m
    hernquist_rho_0 : float
        Best fit central density of Hernquist bulge, calculated in
        spam.data.fit_stellar_bulge.py. UNITS: kg/m^3
    stellar_expdisc_sigma_0 : float
        Best fit central density of stellar disc, calculated in
        spam.data.fit_stellar_disc.py. UNITS: kg/m^2
    stellar_expdisc_R_d : float
        Best fit scale length of stellar disc, calculated in
        spam.data.fit_stellar_disc.py. UNITS: m
    ext_potential : float
        Maximum posterior external potential (specifically, log10(phi/c^2) )
        calculated via the screening map of Desmond et al. (see Naik et al.,
        2019 for details and refs).
    ext_potential_lower : float
        1 sigma lower bound on external potential.
    ext_potential_upper : ditto
        1 sigma upper bound on external potential.
    """
    def __init__(self, name):

        self.name = name
        datadir = _os.path.dirname(_os.path.realpath(__file__))+"/SPARCData"

        # loading metadata
        listfile = open(datadir+"/metadata.txt", 'r')
        data = listfile.readlines()
        listfile.close()

        names = []
        for i in range(len(data)):
            names.append(data[i].split()[0])
        ind = names.index(self.name)

        htypes = {0: 'S0', 1: 'Sa', 2: 'Sab', 3: 'Sb', 4: 'Sbc', 5: 'Sc',
                  6: 'Scd', 7: 'Sd', 8: 'Sdm', 9: 'Sm', 10: 'Im', 11: 'BCD'}
        self.hubble_type = htypes[int(data[ind].split()[1])]
        self.distance = float(data[ind].split()[2])*_Mpc  # metres
        self.distance_err = float(data[ind].split()[3])*_Mpc  # metres
        self.distance_method = int(data[ind].split()[4])
        self.inclination = float(data[ind].split()[5])  # degrees
        self.inclination_err = float(data[ind].split()[6])  # degrees
        self.luminosity_tot = float(data[ind].split()[7])  # 1e+9 Lsun
        self.luminosity_err = float(data[ind].split()[8])  # 1e+9 Lsun
        self.disc_scale = float(data[ind].split()[11])*_kpc  # metres
        self.disc_SB = float(data[ind].split()[12])/_pc**2  # Lsun/m^2
        self.HI_mass = float(data[ind].split()[13])*1e+9*_Msun  # kg
        self.Q_flag = int(data[ind].split()[17])

        # loading main SPARC data
        self.filename = datadir+"/data/"+name+"_rotmod.dat"
        gal_file = open(self.filename, 'r')
        data = gal_file.readlines()
        gal_file.close()
        self.R = _np.zeros((len(data[3:]),))
        self.v = _np.zeros((len(data[3:]),))
        self.v_err = _np.zeros((len(data[3:]),))
        self.v_gas = _np.zeros((len(data[3:]),))
        self.v_disc = _np.zeros((len(data[3:]),))
        self.v_bul = _np.zeros((len(data[3:]),))
        for i in range(len(data[3:])):
            self.R[i] = float(data[3:][i].split()[0])
            self.v[i] = float(data[3:][i].split()[1])
            self.v_err[i] = float(data[3:][i].split()[2])
            self.v_gas[i] = float(data[3:][i].split()[3])
            self.v_disc[i] = float(data[3:][i].split()[4])
            self.v_bul[i] = float(data[3:][i].split()[5])
        if (self.v_bul == 0).all():
            self.StellarBulge = False
        else:
            self.StellarBulge = True

        # loading coords
        coordfile = open(datadir+"/coords.txt", 'r')
        data = coordfile.readlines()[1:]
        coordfile.close()
        assert data[ind].split()[0] == self.name
        self.coords_RA = float(data[ind].split()[2])
        self.coords_DEC = float(data[ind].split()[3])

        # loading gas radius
        gasfile = open(datadir+"/gas_radii.txt", 'r')
        data = gasfile.readlines()
        gasfile.close()
        assert data[ind].split()[0] == self.name
        self.gas_radius = float(data[ind].split()[1])

        # loading hernquist parameters
        if self.StellarBulge:
            hernquistfile = open(datadir+"/hernquist_parameters.txt", 'r')
            data = hernquistfile.readlines()
            hernquistfile.close()
            assert data[ind].split()[0] == self.name
            self.hernquist_rho_0 = float(data[ind].split()[1])
            self.hernquist_radius = float(data[ind].split()[2])
        else:
            self.hernquist_rho_0 = None
            self.hernquist_radius = None

        # loading stellar disc fit parameters
        discparfile = open(datadir+"/stellar_disc_parameters.txt", 'r')
        data = discparfile.readlines()
        discparfile.close()
        assert data[ind].split()[0] == self.name
        self.stellar_expdisc_sigma_0 = float(data[ind].split()[1])  # kg/m^2
        self.stellar_expdisc_R_d = float(data[ind].split()[2])  # metres

        # loading external potential data
        potential_dir = datadir+"/SPARC_potentials"
        col1 = _np.array([], dtype=_np.float64)
        col2 = _np.array([], dtype=_np.float64)
        col3 = _np.array([], dtype=_np.float64)
        for i in range(20):
            file = open(potential_dir+"/SPARC_screen_"+str(i)+".dat", 'r')
            data = file.readlines()
            file.close()
            assert data[ind].split()[0] == self.name
            col1 = _np.append(col1, float(data[ind].split()[1]))
            col2 = _np.append(col2, float(data[ind].split()[2]))
            col3 = _np.append(col3, float(data[ind].split()[3]))
        self.ext_potential_lower = col1
        self.ext_potential = col2
        self.ext_potential_upper = col3

        return

# getting list of galaxy names
names_full = []
names_standard = []
_datadir = _os.path.dirname(_os.path.realpath(__file__))+"/SPARCData"
_namefile = open(_datadir+"/names_full.txt", 'r')
_data = _namefile.readlines()
_namefile.close()
names_full = []
for _i in range(len(_data)):
    names_full.append(_data[_i].split()[0])
_namefile = open(_datadir+"/names_standard.txt", 'r')
_data = _namefile.readlines()
_namefile.close()
names_standard = []
for _i in range(len(_data)):
    names_standard.append(_data[_i].split()[0])

__all__ = ['SPARCGalaxy', 'names_full', 'names_standard']