#!/usr/bin/python

import numpy
from PIL import Image
import argparse

def get_camera_params_perspective(calibration_file_path):
    print("loading sensor parameters from file %s" % calibration_file_path)
    f = open(calibration_file_path)
    calibration_file = f.readlines()

    #reading matrix for extrinsic (hand-eye) and intrinisic sensor calibration
    for line in calibration_file:
        if line[0] != '#' and line[0] != '\n':
            key, value = line.split("=")
            if key == "camera.0.R":
                sensor2tcp_rot = numpy.fromstring(value.translate(None, "[;]"),dtype=float,sep=' ')
            if key == "camera.0.T":
                sensor2tcp_trans = numpy.fromstring(value.translate(None, "[;]"),dtype=float,sep=' ')
            if key == "camera.0.A":
                sensor_intrinsic = numpy.fromstring(value.translate(None, "[;]"),dtype=float,sep=' ')

    #converting to numpy matrix
    sensor2tcp = numpy.matrix(( numpy.append(sensor2tcp_rot[0:3], sensor2tcp_trans[0]), numpy.append(sensor2tcp_rot[3:6], sensor2tcp_trans[1]), numpy.append(sensor2tcp_rot[6:9], sensor2tcp_trans[2]), (0., 0., 0., 1.) ))

    #calculating pixel offset and delta of sensor based on perspective camera system                
    focal_length_x = sensor_intrinsic[0]
    focal_length_y = sensor_intrinsic[4]
    principal_x = sensor_intrinsic[2]
    principal_y = sensor_intrinsic[5]
    offset_x = - (principal_x / focal_length_x)
    offset_y = - (principal_y / focal_length_y)
    delta_x = 1.0 / focal_length_x
    delta_y = 1.0 / focal_length_y
    return sensor2tcp, offset_x, offset_y, delta_x, delta_y

def create_pointcloud(depth_image, tcp2world, calibration_file_path):
    sensor2tcp, offset_x, offset_y, delta_x, delta_y = get_camera_params_perspective(calibration_file_path)
    
    #transformation from world to sensor
    sensor2world = tcp2world.dot(sensor2tcp)

    #initializing pointcloud for all depth values
    pointcloud = numpy.empty(shape=(depth_image.size[0] * depth_image.size[1],3))

    #loading image
    depth_image_values = depth_image.load()

    #raw point in sensor pixels
    raw = numpy.empty(3)
    
    #transform depth value to world coordinate system
    raw[1] = offset_y
    for y in range(depth_image.size[1]):
        raw[0] = offset_x
        for x in range(depth_image.size[0]):
            if depth_image_values[x,y] != 0:
                #convert to millimeters
                depth = depth_image_values[x,y] * 0.124987

                # z axis is depth
                raw[2] = depth

                #calculate measurement as point in sensor coordinate system (multiply x and y based on camera pixels with depth
                sensor_point = numpy.array([raw[0] * raw[2], raw[1] * raw[2], raw[2], 1.0])

                #transform point to world coordinate system
                world_point = sensor2world.dot(sensor_point)
                pointcloud[depth_image.size[0] * y + x] = world_point[0,0:3]

                raw[0] += delta_x
        raw[1] += delta_y

    # save rows with at least one non zero value
    print("saving pointcloud to file pointcloud.txt")
    idxs = numpy.any(pointcloud != 0, axis=1) 
    numpy.savetxt("pointcloud.txt", pointcloud[idxs, :], fmt='%.4f')     


if __name__ == '__main__':

    parser = argparse.ArgumentParser(description='Converting a depth image from the indu dataset to a pointcloud as python array.')
    parser.add_argument('depth_image_path')
    parser.add_argument('pose_file_path')
    parser.add_argument('calibration_file_path')
    args = parser.parse_args()
    
    #loading depth image from file                                                                                                                                                                                 
    depth_image = Image.open(args.depth_image_path)

    #loading tcp2world homogenous matrix transformation from file
    tcp2world_array = numpy.loadtxt(args.pose_file_path)
    tcp2world = numpy.matrix( ( (tcp2world_array[0:4]),(tcp2world_array[4:8]),(tcp2world_array[8:12]), (0., 0., 0., 1.) ) )

    create_pointcloud(depth_image, tcp2world, args.calibration_file_path)