# examples/Python/Advanced/global_registration.py
import open3d as o3d
import numpy as np
import copy
import sys
import glob
import os



def obj_pcd(file):
    #getting new path
    path=file.replace('\\','/')
    path=path.replace('\\','/')
    base_path=path.rsplit('/',1)
    name=base_path[1].split('.')
    print(name)
    new_path=base_path[0]+'/'+name[0]+'.pcd'


    #convert obj->pcd
    file=open(path,'r')
    list_norm=[]
    list_xyz=[]
    for i in file:
        if not(i.find('vn ')):
            string=i.strip("vn").strip('\n').replace(" ",",").strip("").split(",")
            del string[0]
            list_norm.append(string)
        if not(i.find('v ')):
            string = i.strip("v").strip('\n').replace(" ",",").strip("").split(",")
            del string[0]
            list_xyz.append(string)
    file.close()
    list_xyz=np.float32(list_xyz)
    list_norm=np.float32(list_norm)

    file_new=open(new_path,'w+')
    if len(list_norm)==0:
        file_new.write("# .PCD v0.7 - Point Cloud Data file format\n\
VERSION 0.7\n\
FIELDS x y z\n\
SIZE 4 4 4\n\
TYPE F F F\n\
COUNT 1 1 1\n\
WIDTH %s\n\
HEIGHT 1\n\
VIEWPOINT 0 0 0 1 0 0 0\n\
POINTS %s\n\
DATA ascii\n\
"%(str(len(list_xyz)),str(len(list_xyz))))
        for i in range(0,len(list_xyz)):
            file_new.write("%f %f %f\n"%(list_xyz[i][0],list_xyz[i][1],list_xyz[i][2]))
        file_new.close()
    else:
        file_new.write("# .PCD v0.7 - Point Cloud Data file format\n\
VERSION 0.7\n\
FIELDS x y z\n\
SIZE 4 4 4\n\
TYPE F F F\n\
COUNT 1 1 1\n\
WIDTH %s\n\
HEIGHT 1\n\
VIEWPOINT 0 0 0 1 0 0 0\n\
POINTS %s\n\
DATA ascii\n\
" % (str(len(list_xyz)), str(len(list_xyz))))
        for i in range(0, len(list_xyz)):
            file_new.write("%f %f %f %f %f %f\n" % (list_xyz[i][0], list_xyz[i][1], list_xyz[i][2],list_norm[i][0], list_norm[i][1], list_norm[i][2]))
        file_new.close()

def get_mesh(pc):
    path=pc.rsplit('/',maxsplit=1)
    folder_path=path[0]
    file=open(pc,'r')
    file_w=open(folder_path+'/mesh_data.pcd','w+')
    list=[]
    for i in file:
        if not (i.find('f')):
           list.append(i)
           file_w.write(str(i))
    file.close()
    file_w.close()

def write_obj(vert,faces):
    p_ver=vert
    p_faces=faces
    p_folder_=p_ver.rsplit('/',maxsplit=1)
    p_folder=p_folder_[0]
    count=p_folder_[1].rsplit('.',maxsplit=1)
    file_v=open(p_ver,'r')
    file_f=open(p_faces,'r')
    file_r=open(p_folder+'/'+str(count[0])+'.obj','w+')
    c=0
    for i in file_v:
        c +=1
        if c>11:
            file_r.write('v'+' '+str(i))
    for j in file_f:
        file_r.write(str(j))
    file_v.close()
    file_f.close()
    file_r.close()

def draw_registration_result(source, target, transformation):
    source_temp = copy.deepcopy(source)
    target_temp = copy.deepcopy(target)
    source_temp.paint_uniform_color([1, 0.706, 0])
    target_temp.paint_uniform_color([0, 0.651, 0.929])
    source_temp.transform(transformation)
    o3d.visualization.draw_geometries([source_temp, target_temp])


def preprocess_point_cloud(pcd, voxel_size):
    print(":: Downsample with a voxel size %.3f." % voxel_size)
    pcd_down = pcd.voxel_down_sample(voxel_size=voxel_size)


    radius_normal = voxel_size * 5 #orignal 2
    print(":: Estimate normal with search radius %.3f." % radius_normal)
    pcd_down.estimate_normals(o3d.geometry.KDTreeSearchParamHybrid(radius=radius_normal, max_nn=3))#orignsl 30

    radius_feature = voxel_size * 5 #orignal 5
    print(":: Compute FPFH feature with search radius %.3f." % radius_feature)
    pcd_fpfh = o3d.registration.compute_fpfh_feature(
        pcd_down,
        o3d.geometry.KDTreeSearchParamHybrid(radius=radius_feature, max_nn=3)) #orignal 100
    return pcd_down, pcd_fpfh


def prepare_dataset(voxel_size):
    print(":: Load two point clouds and disturb initial pose.")
    """
    source = o3d.io.read_point_cloud(r"\stick_template.pcd")
    target = o3d.io.read_point_cloud(r"\pruned_cloud.pcd")
    """
    source = o3d.io.read_point_cloud(r"D:\skelton-approach\10-12-vSS2\template\stick_template.pcd")
    target = o3d.io.read_point_cloud(r"D:\skelton-approach\10-12-vSS2\4.Skelton\1_Spruned.obj")
    trans_init = np.asarray([[0.0, 0.0, 1.0, 0.0], [1.0, 0.0, 0.0, 0.0],
                             [0.0, 1.0, 0.0, 0.0], [0.0, 0.0, 0.0, 1.0]])
    source.transform(trans_init)
    draw_registration_result(source, target, np.identity(4))

    source_down, source_fpfh = preprocess_point_cloud(source, voxel_size)
    target_down, target_fpfh = preprocess_point_cloud(target, voxel_size)
    print(source_down)
    print(source_fpfh)
    return source, target, source_down, target_down, source_fpfh, target_fpfh


def execute_global_registration(source_down, target_down, source_fpfh,
                                target_fpfh, voxel_size):
    distance_threshold = voxel_size * 3.5 #orignal 1.5 playing----orignal 2
    print(":: RANSAC registration on downsampled point clouds.")
    print("   Since the downsampling voxel size is %.3f," % voxel_size)
    print("   we use a liberal distance threshold %.3f." % distance_threshold)
    result = o3d.registration.registration_ransac_based_on_feature_matching(
        source_down, target_down, source_fpfh, target_fpfh, distance_threshold,
        o3d.registration.TransformationEstimationPointToPoint(False), 4, [
            o3d.registration.CorrespondenceCheckerBasedOnEdgeLength(0.9),
            o3d.registration.CorrespondenceCheckerBasedOnDistance(
                distance_threshold)
        ], o3d.registration.RANSACConvergenceCriteria(4000000, 500))
    return result


def refine_registration(source, target, source_fpfh, target_fpfh, voxel_size,rs):
    distance_threshold = voxel_size * 2 #orignal 0.4 playing-----orignal 1
    print(":: Point-to-plane ICP registration is applied on original point")
    print("   clouds to refine the alignment. This time we use a strict")
    print("   distance threshold %.3f." % distance_threshold)
    source_normal=source.estimate_normals()
    target_normal=target.estimate_normals()
    result = o3d.registration.registration_icp(
        source, target, distance_threshold, rs.transformation,
        o3d.registration.TransformationEstimationPointToPoint(),o3d.registration.ICPConvergenceCriteria(max_iteration = 2000))
    print("-------------------------------")
    return result



def geometry(temp,pc):
    geometry.counter +=1
    template=temp
    pointcloud=pc

    print(":: Load two point clouds and disturb initial pose.")
    template=template.replace('\\','/')
    pointcloud = pointcloud.replace('\\', '/')
    source = o3d.io.read_point_cloud(template)
    target = o3d.io.read_point_cloud(pointcloud)

    voxel_size = 0.05
    trans_init = np.asarray([[0.0, 0.0, 1.0, 0.0], [1.0, 0.0, 0.0, 0.0],
                             [0.0, 1.0, 0.0, 0.0], [0.0, 0.0, 0.0, 1.0]])
    source.transform(trans_init)

    #draw_registration_result(source, target, np.identity(4))

    source_down, source_fpfh = preprocess_point_cloud(source, voxel_size)
    target_down, target_fpfh = preprocess_point_cloud(target, voxel_size)
    #return source, target, source_down, target_down, source_fpfh, target_fpfh



    #voxel_size = 0.05  # means 5cm for the dataset (0.05)
    #source, target, source_down, target_down, source_fpfh, target_fpfh = \
    #       prepare_dataset(voxel_size)

    result_ransac = execute_global_registration(source_down, target_down,
                                                source_fpfh, target_fpfh,
                                                voxel_size)

    print(result_ransac)
    draw_registration_result(source_down, target_down,
                             result_ransac.transformation)
    """
    result=source.transform(result_ransac.transformation)
    pointcloud=pointcloud.rsplit('/',maxsplit=2)
    path=pointcloud[0]+'/alignedCloud/%s.pcd'%(str(geometry.counter))
    o3d.io.write_point_cloud(path,result,write_ascii=True)
    write_obj(path,pointcloud[0]+'/stickdata/mesh_data.pcd')
    """
    # Inverse the transformation matrix for template

    # Show the results

    #Select the best one

    #Save the Aligned template to the folder
    

    result_icp = refine_registration(source, target, source_fpfh, target_fpfh,
                                     voxel_size,result_ransac)
    print(result_icp)
    draw_registration_result(source, target, result_icp.transformation)
    result = source.transform(result_icp.transformation)
    pointcloud = pointcloud.rsplit('/', maxsplit=2)
    path = pointcloud[0] + '/alignedCloud/%s.pcd' % (str(geometry.counter))
    o3d.io.write_point_cloud(path, result, write_ascii=True)
    #write_obj(path, pointcloud[0] + '/stickdata/mesh_data.pcd')


geometry.counter=0
base_path=r"D:\acquisitions\2019-04-05\aidan_orignal_1_slap"    #sys.arg[1]
base_path=base_path.replace('\\','/')


#Getting template- assuming just one .obj file in folder
list_temp=[]
for j in glob.glob(base_path+'/stickdata/*template.obj'):
    list_temp.append(j.replace('\\','/'))

#Getting Point clouds
list_points=[]
path_point=base_path+'/pruned_pointcloud'
count=1
for i in glob.glob(path_point+'/*.obj'):
    list_points.append(i.replace('\\','/'))
    count +=1
    if count==7:
        break

#writing .pcd files
obj_pcd(list_temp[0])
for i in range (0,min(6,len(list_points))):
    obj_pcd(list_points[i])

#getting .pcd files locations
list_temp = []
for j in glob.glob(base_path + '/stickdata/*.pcd'):    #add features to delete already existng .pcd files
    list_temp.append(j.replace('\\', '/'))
list_points = []
path_point = base_path + '/pruned_pointcloud'
count = 1
for i in glob.glob(path_point + '/*.pcd'):
    list_points.append(i.replace('\\', '/'))
    count += 1
    if count == 7:
        break

#creating aligned cloud directory:'
dir=base_path+'/alignedCloud'
try:
    os.makedirs(dir)
except OSError:
    print("dir not created")

#Writing aligned clouds
for i in range(0,min(6,len(list_points))):
    print(list_points[i])
    print(list_temp[0])
    geometry(list_points[i],list_temp[0])




















