#! /usr/bin/env python3
import roslib
import rospy
from std_msgs.msg import Header
from std_msgs.msg import String
from geometry_msgs.msg import Point
import sys
# sys.path.append('/home/tw/Badminton_traj_ws/src/badminton_trajectory')
# from msg import Info

import pyrealsense2 as rs

import getopt
import numpy as np
import os
from glob import glob
import piexif
from keras.preprocessing.image import ImageDataGenerator  #array_to_img, img_to_array, load_img
from keras.utils import image_utils
import pandas as pd
from sklearn.model_selection import train_test_split
from keras.models import *
from keras.layers import *
from TrackNet3 import TrackNet3
import keras.backend as K
from keras import optimizers
import tensorflow as tf
import cv2
from os.path import isfile, join
from PIL import Image
import time
import queue

BATCH_SIZE=1
HEIGHT=288
WIDTH=512
#HEIGHT=360
#WIDTH=640
sigma=2.5
mag=1

pipeline = rs.pipeline()  # 定义流程pipeline
config = rs.config()  # 定义配置config
# config.enable_stream(rs.stream.depth, 848, 480, rs.format.z16, 30)
# config.enable_stream(rs.stream.color, 848, 480, rs.format.bgr8, 30)
config.enable_stream(rs.stream.depth, 1280, 720, rs.format.z16, 30)
config.enable_stream(rs.stream.color, 1280, 720, rs.format.bgr8, 30)
profile = pipeline.start(config)  # 流程开始
align_to = rs.stream.color  # 与color流对齐
align = rs.align(align_to)

def get_aligned_images():
    frames = pipeline.wait_for_frames()  # 等待获取图像帧
    aligned_frames = align.process(frames)  # 获取对齐帧
    aligned_depth_frame = aligned_frames.get_depth_frame()  # 获取对齐帧中的depth帧
    color_frame = aligned_frames.get_color_frame()  # 获取对齐帧中的color帧

    ############### 相机参数的获取 #######################
    intr = color_frame.profile.as_video_stream_profile().intrinsics  # 获取相机内参
    depth_intrin = aligned_depth_frame.profile.as_video_stream_profile(
    ).intrinsics  # 获取深度参数（像素坐标系转相机坐标系会用到）
    '''camera_parameters = {'fx': intr.fx, 'fy': intr.fy,
                         'ppx': intr.ppx, 'ppy': intr.ppy,
                         'height': intr.height, 'width': intr.width,
                         'depth_scale': profile.get_device().first_depth_sensor().get_depth_scale()
                         }'''

    # 保存内参到本地
    # with open('./intrinsics.json', 'w') as fp:
    #json.dump(camera_parameters, fp)
    #######################################################

    depth_image = np.asanyarray(aligned_depth_frame.get_data())  # 深度图（默认16位）
    depth_image_8bit = cv2.convertScaleAbs(depth_image, alpha=0.03)  # 深度图（8位）
    depth_image_3d = np.dstack(
        (depth_image_8bit, depth_image_8bit, depth_image_8bit))  # 3通道深度图
    color_image = np.asanyarray(color_frame.get_data())  # RGB图，返回ndarray的形式

    # 返回相机内参、深度参数、彩色图、深度图、齐帧中的depth帧
    return intr, depth_intrin, color_image, depth_image, aligned_depth_frame


def genHeatMap(w, h, cx, cy, r, mag):
	if cx < 0 or cy < 0:
		return np.zeros((h, w))
	x, y = np.meshgrid(np.linspace(1, w, w), np.linspace(1, h, h))
	heatmap = ((y - (cy + 1))**2) + ((x - (cx + 1))**2)
	heatmap[heatmap <= r**2] = 1
	heatmap[heatmap > r**2] = 0
	return heatmap*mag

#time: in milliseconds
def custom_time(time):
	remain = int(time / 1000)
	ms = (time / 1000) - remain
	s = remain % 60
	s += ms
	remain = int(remain / 60)
	m = remain % 60
	remain = int(remain / 60)
	h = remain
	#Generate custom time string
	cts = ''
	if len(str(h)) >= 2:
		cts += str(h)
	else:
		for i in range(2 - len(str(h))):
			cts += '0'
		cts += str(h)
	
	cts += ':'

	if len(str(m)) >= 2:
		cts += str(m)
	else:
		for i in range(2 - len(str(m))):
			cts += '0'
		cts += str(m)

	cts += ':'

	if len(str(int(s))) == 1:
		cts += '0'
	cts += str(s)

	return cts

#Loss function
def custom_loss(y_true, y_pred):
	loss = (-1)*(K.square(1 - y_pred) * y_true * K.log(K.clip(y_pred, K.epsilon(), 1)) + K.square(y_pred) * (1 - y_true) * K.log(K.clip(1 - y_pred, K.epsilon(), 1)))
	return K.mean(loss)

def publish_image(real_x, real_y, real_z, classification, confidence):
    detect_result=Info()
    rate = rospy.Rate(30)
    header = Header(stamp=rospy.Time.now())
    header.frame_id = 'map'

    detect_result.x = real_x
    detect_result.y = real_y
    detect_result.z = real_z
    detect_result.classification = classification
    detect_result.confidence = confidence

    pub.publish(detect_result)
    rate.sleep()
    # rospy.loginfo("x:%.3f,  y:%.3f,  z:%.3f,  classification:%s: ", detect_result.x, detect_result.y,  detect_result.z, detect_result.classification)


if __name__ == "__main__":
	# rospy.init_node('rs_tracknet')
	# pub = rospy.Publisher("/detect_result_out", Info, queue_size=10)
	load_weights = "/home/tw/Badminton_traj_ws/src/badminton_trajectory/model906_30"

	# 设置轨迹点个数
	q = queue.deque()
	for i in range(0,24):
		q.appendleft(None)



	model = load_model(load_weights, custom_objects={'custom_loss':custom_loss})
	model.summary()
	print('Beginning predicting......')

	start = time.time()
	try:
		while True:
			intr_1, depth_intrin_1, color_image_1, depth_image_1, aligned_depth_frame_1 = get_aligned_images()
			# cap = cv2.VideoCapture(0)
			# print(videoName)
			if not depth_image_1.any() or not color_image_1.any():
				continue
			# Convert images to numpy arrays
			# Apply colormap on depth image (image must be converted to 8-bit per pixel first)
			depth_colormap = cv2.applyColorMap(cv2.convertScaleAbs(
				depth_image_1, alpha=0.03), cv2.COLORMAP_JET)
			# Stack both images horizontally
			images_1 = np.hstack((color_image_1, depth_colormap))
			intr_2, depth_intrin_2, color_image_2, depth_image_2, aligned_depth_frame_2 = get_aligned_images()
			if not depth_image_2.any() or not color_image_2.any():
				continue
			# Convert images to numpy arrays
			# Apply colormap on depth image (image must be converted to 8-bit per pixel first)
			depth_colormap = cv2.applyColorMap(cv2.convertScaleAbs(
				depth_image_2, alpha=0.03), cv2.COLORMAP_JET)
			# Stack both images horizontally
			images_2 = np.hstack((color_image_2, depth_colormap))
			intr_3, depth_intrin_3, color_image_3, depth_image_3, aligned_depth_frame_3 = get_aligned_images()
			if not depth_image_3.any() or not color_image_3.any():
				continue
			# Convert images to numpy arrays
			# Apply colormap on depth image (image must be converted to 8-bit per pixel first)
			depth_colormap = cv2.applyColorMap(cv2.convertScaleAbs(
				depth_image_3, alpha=0.03), cv2.COLORMAP_JET)
			# Stack both images horizontally
			images_3 = np.hstack((color_image_3, depth_colormap))
			# success, image1 = cap.read()
			# print(success)
			# frame_time1 = custom_time(cap.get(cv2.CAP_PROP_POS_MSEC))
			# success, image2 = cap.read()
			# frame_time2 = custom_time(cap.get(cv2.CAP_PROP_POS_MSEC))
			# success, image3 = cap.read()
			# frame_time3 = custom_time(cap.get(cv2.CAP_PROP_POS_MSEC))

			ratio = color_image_1.shape[0] / HEIGHT

			size = (int(WIDTH*ratio), int(HEIGHT*ratio))
			# fps = 30


			# out = cv2.VideoWriter(videoName[:-4]+'_predict_show'+videoName[-4:], fourcc, fps, size)

			count = 0

			
			unit = []
			#Adjust BGR format (cv2) to RGB format (PIL)
			x1 = color_image_1[...,::-1]
			x2 = color_image_2[...,::-1]
			x3 = color_image_3[...,::-1]
			
			#Convert np arrays to PIL images
			x1 = image_utils.array_to_img(x1)
			x2 = image_utils.array_to_img(x2)
			x3 = image_utils.array_to_img(x3)

			#Resize the images
			x1 = x1.resize(size = (WIDTH, HEIGHT))
			x2 = x2.resize(size = (WIDTH, HEIGHT))
			x3 = x3.resize(size = (WIDTH, HEIGHT))
			#Convert images to np arrays and adjust to channels first
			x1 = np.moveaxis(image_utils.img_to_array(x1), -1, 0)
			x2 = np.moveaxis(image_utils.img_to_array(x2), -1, 0)
			x3 = np.moveaxis(image_utils.img_to_array(x3), -1, 0)
				#Create data
			unit.append(x1[0])
			unit.append(x1[1])
			unit.append(x1[2])
			unit.append(x2[0])
			unit.append(x2[1])
			unit.append(x2[2])
			unit.append(x3[0])
			unit.append(x3[1])
			unit.append(x3[2])
			unit=np.asarray(unit)	
			unit = unit.reshape((1, 9, HEIGHT, WIDTH))
			unit = unit.astype('float32')
			unit /= 255
			y_pred = model.predict(unit, batch_size=BATCH_SIZE)
			y_pred = y_pred > 0.5
			y_pred = y_pred.astype('float32')
			h_pred = y_pred[0]*255
			h_pred = h_pred.astype('uint8')
			for i in range(3):
				if i == 0:
					image = color_image_1
					depth_intrin = depth_intrin_1
					aligned_depth_frame = aligned_depth_frame_1
				elif i == 1:
					image = color_image_2
					depth_intrin = depth_intrin_2
					aligned_depth_frame = aligned_depth_frame_2
				elif i == 2:
					image = color_image_3
					depth_intrin = depth_intrin_3
					aligned_depth_frame = aligned_depth_frame_3


				if np.amax(h_pred[i]) > 0:
				# 	out.write(image)
				# else:	
				#h_pred
					(cnts, _) = cv2.findContours(h_pred[i].copy(), cv2.RETR_EXTERNAL, cv2.CHAIN_APPROX_SIMPLE)  # 寻找目标轮廓,cv2.RETR_EXTERNAL表示只检测外轮廓
					rects = [cv2.boundingRect(ctr) for ctr in cnts]  # 返回几个轮廓点坐标和对应最小外接矩形宽高    # cv2.CHAIN_APPROX_SIMPLE压缩水平方向，垂直方向，对角线方向的元素，只保留该方向的终点坐标
					# print(rects)
					max_area_idx = 0
					max_area = rects[max_area_idx][2] * rects[max_area_idx][3]
					for i in range(len(rects)):
						area = rects[i][2] * rects[i][3]
						if area > max_area:
							max_area_idx = i
							max_area = area
					target = rects[max_area_idx]
					(cx_pred, cy_pred) = (int(ratio*(target[0] + target[2] / 2)), int(ratio*(target[1] + target[3] / 2)))
					dis = aligned_depth_frame.get_distance(cx_pred, cy_pred)
					center = (cx_pred, cy_pred)
                    # camera_xyz = rs.rs2_deproject_pixel_to_point(
                    #     depth_intrin, (ux, uy), dis)  # 计算相机坐标系的xyz
					camera_xyz = rs.rs2_deproject_pixel_to_point(depth_intrin, center, dis)
					camera_xyz = np.round(np.array(camera_xyz), 3)  # 转成3位小数
					camera_xyz = camera_xyz.tolist()
					cv2.circle(image, (cx_pred, cy_pred), 5, (0,0,255), -1)
					cv2.putText(image, str(camera_xyz), (cx_pred+20, cy_pred+10), 0, 1,
                                [225, 255, 255], thickness=2, lineType=cv2.LINE_AA)#标出坐标

					# if cx_pred != 0 and cy_pred != 0:
					# 	q.appendleft([cx_pred, cy_pred])
					# 	q.pop()
					# else:
					# 	q.appendleft(None)
					# 	q.pop()
					# image_cp = np.copy(image)
					# for i in range(0,24):
					# 	if q[i] is not None:
					# 		draw_x = q[i][0]
					# 		draw_y = q[i][1]
					# 		cv2.circle(image_cp, (draw_x, draw_y), 5, (0,0,255), -1)
					cv2.namedWindow('detection', flags=cv2.WINDOW_NORMAL |
									cv2.WINDOW_KEEPRATIO | cv2.WINDOW_GUI_EXPANDED)
					cv2.imshow('detection', image)
				else:
					cv2.namedWindow('detection', flags=cv2.WINDOW_NORMAL |
									cv2.WINDOW_KEEPRATIO | cv2.WINDOW_GUI_EXPANDED)
					cv2.imshow('detection', image)
				key = cv2.waitKey(1)
				if key & 0xFF == ord('q') or key == 27:
					cv2.destroyAllWindows()
					break
            # Press esc or 'q' to close the image window
			

				# success, image1 = cap.read()
				# frame_time1 = custom_time(cap.get(cv2.CAP_PROP_POS_MSEC))
				# success, image2 = cap.read()
				# frame_time2 = custom_time(cap.get(cv2.CAP_PROP_POS_MSEC))
				# success, image3 = cap.read()
				# frame_time3 = custom_time(cap.get(cv2.CAP_PROP_POS_MSEC))
			end = time.time()
			print('Prediction time:', end-start, 'secs')
			print('Done......')
	finally:
        # Stop streaming
		pipeline.stop()
	rospy.spin()