#!/usr/bin/env python
import numpy as np
import rospy, tf, sys, time
from geometry_msgs.msg import Pose, PoseStamped
from sensor_msgs.msg import Image
from sensor_msgs.msg import CameraInfo
import tf2_ros
import moveit_commander
import moveit_msgs.msg
import moveit_msgs.srv
import geometry_msgs.msg
import cv2
import pyrealsense2 as rs
import copy

# camera matrix of realsense415 with half resolution
camera_info = CameraInfo()
camera_info.K = [931.6937866210938, 0.0, 624.7894897460938, 0.0, 931.462890625, 360.5186767578125, 0.0, 0.0, 1.0]
camera_info.header.frame_id = 'camera_color_optical_frame'
camera_info.height = 720
camera_info.width = 1280

# TODO: Construct 3D calibration grid across workspace
# User options (change me)
# --------------- Setup options ---------------
original_point = [0.2131, -0.0445, 0.3986]
bias = 0.01
workspace_limits = np.asarray([[original_point[0], original_point[0]], 
                               [original_point[1]-bias, original_point[1]+bias], 
                               [original_point[2]-bias/2.0, original_point[2]+bias/2.0]]) # Cols: min max, Rows: x y z (define workspace limits in robot coordinates)
calib_grid_step = 0.002
chessboard_size = (3,3)
refine_criteria = (cv2.TERM_CRITERIA_EPS + cv2.TERM_CRITERIA_MAX_ITER, 30, 0.001)
# ---------------------------------------------
gridspace_x = np.linspace(workspace_limits[0][0], workspace_limits[0][1], 1 + (workspace_limits[0][1] - workspace_limits[0][0])/calib_grid_step)
gridspace_y = np.linspace(workspace_limits[1][0], workspace_limits[1][1], 1 + (workspace_limits[1][1] - workspace_limits[1][0])/calib_grid_step)
gridspace_z = np.linspace(workspace_limits[2][0], workspace_limits[2][1], 1 + (workspace_limits[2][1] - workspace_limits[2][0])/calib_grid_step)
calib_grid_x, calib_grid_y, calib_grid_z = np.meshgrid(gridspace_x, gridspace_y, gridspace_z)
num_calib_grid_pts = calib_grid_x.shape[0]*calib_grid_x.shape[1]*calib_grid_x.shape[2]
calib_grid_x.shape = (num_calib_grid_pts,1)
calib_grid_y.shape = (num_calib_grid_pts,1)
calib_grid_z.shape = (num_calib_grid_pts,1)
calib_grid_pts = np.concatenate((calib_grid_x, calib_grid_y, calib_grid_z), axis=1)

# realsense
points = rs.points()
pipeline= rs.pipeline()
config = rs.config()
config.enable_stream(rs.stream.color, 1280, 720, rs.format.bgr8, 30)
config.enable_stream(rs.stream.depth, 1280, 720, rs.format.z16, 30)
profile = pipeline.start(config)
depth_sensor = profile.get_device().first_depth_sensor()
depth_scale = depth_sensor.get_depth_scale()
print("Depth Scale is: " , depth_scale)
align_to = rs.stream.color
align = rs.align(align_to)

# TODO: calculate the position (x, y, z) of the chessboard center
def image_callback(color_image, depth_frame,color_intrinsics, iter):
    # convert the color image to gray scale and find the chessboard corners
    gray_image = cv2.cvtColor(color_image, cv2.COLOR_BGR2GRAY)
    chessboard_found, corners = cv2.findChessboardCorners(gray_image, (3,3), None)

    #imgpoints:   stored the corners
    #observed_pt: stored the postion in camera coordinate
    objpoints = [] 
    imgpoints = []
    observed_pt = []
    
    if chessboard_found:
        # refine the corners if chessboard is found
        corners2 = cv2.cornerSubPix(gray_image, corners,(11,11),(-1,-1),refine_criteria)
        imgpoints.append(corners2)

        # Get observed chessboard center 3D point in camera space
        u = int(corners2[4][0][0])
        v = int(corners2[4][0][1])
        #depth = depth_frame.get_distance(u, v)
        window_length = 5
        depth_sum = 0
        for a in range(2*window_length):
            for b in range(2*window_length):
                uu = u - window_length + a
                vv = v - window_length + b
                depth_sum = depth_sum + depth_frame.get_distance(uu, vv)

        depth = depth_sum/((2.0*window_length)*(2.0*window_length))

        # Transform u,v,depth to x,y,z
        if depth!=0:
            x,y,z=rs.rs2_deproject_pixel_to_point(color_intrinsics, [u, v], depth)
            observed_pt.append(x)
            observed_pt.append(y)
            observed_pt.append(z)

        # Draw and display the corners
        img = cv2.drawChessboardCorners(color_image, chessboard_size, corners2, chessboard_found)
        cv2.imshow('img', img)
    else:
        cv2.imshow('img', color_image)
    cv2.waitKey(10)

    return observed_pt

measured_pts = []
observed_pts = []
i = 0
depth_sum=0
while(i<10):
    print("Points Sampling - #",i)

    # type the measured position
#    input_data = raw_input("Please enter the position:").split(' ')
#    measured_pt = [float(input_data[0]), float(input_data[1]), float(input_data[2])]
#    print("The measured position is ",measured_pt)

    # take color and depth images
    frames = pipeline.wait_for_frames()
    aligned_frames = align.process(frames)
    aligned_depth_frame = aligned_frames.get_depth_frame()
    color_frame = aligned_frames.get_color_frame()
    color_intrinsics = color_frame.profile.as_video_stream_profile().intrinsics
    color_image = np.asanyarray(color_frame.get_data())
#    cv2.imwrite('calibration/images/color_image_%02d.png' % i, color_image)

    # calculate the observed position
    observed_pt = image_callback(color_image, aligned_depth_frame, color_intrinsics, iter)
    print("The observed position is ",observed_pt)

    judge = raw_input("Please confirme the positions: Y/n?")
    time.sleep(1)
    if judge=='Y' or judge=='y':
#        measured_pts.append(measured_pt)
#        observed_pts.append(observed_pt)
        depth_sum=depth_sum+observed_pt[2]
        i=i+1
    else:
        print("Resampling...")
print("Average depth is : ",depth_sum/10)
#np.savez("calibration/calibration_data_3.npz",observed_pts,measured_pts)
pipeline.stop()
print("Collecting data Finished and data saved!")
