import glob
import os
import pandas as pd
import argparse
import numpy as np
import cv2
import time
import json

from fire import Fire
from tqdm import tqdm
import random

from detector import Detector
from pose_estimation import PoseNet
#from tracker import SiamFCTracker
from tracker_old import SiamFCTracker
from embeddingNet import EmbeddingNet
from match import Matcher
from model.nms.nms_wrapper import nms

class Track_And_Detect:
	effective_track_thresh = 0.5
	effective_detection_thresh = 0.5
	effective_keypoints_thresh = 0.6
	effective_keypoints_number = 8
	iou_match_thresh = 0.5
	nms_thresh = 0.5
	oks_thresh = 0.8
	embedding_match_thresh = 2
	feature_length = 2048
	
	oks_flag = True
	tracker_flag = True
	tracker_update_flag = True
	new_embedding_flag = True
	descrease_tracker_flag = True
	
	def __init__(self, gpu_id=0, track_model='/export/home/zby/SiamFC/models/output/siamfc_35_old.pth',
								 #track_model = '/export/home/zby/SiamFC/data/models/siamfc_cpm_368_5.pth',
								 pose_model='/export/home/zby/SiamFC/data/models/final_new.pth.tar',
								 embedding_model='/export/home/zby/SiamFC/models/embedding_model.pth'):
		if self.tracker_flag:
			self.tracker = SiamFCTracker(gpu_id, track_model)
		self.posenet = PoseNet(gpu_id, pose_model)
		if not self.new_embedding_flag:
			self.embedder = EmbeddingNet(gpu_id, embedding_model)
		self.matcher = Matcher()
		print('----------------------------------------')
		print('Flag parameters are set as follow:')
		print('Tracker flag: {}'.format(self.tracker_flag))
		print('Tracker update flag: {}'.format(self.tracker_update_flag))
		print('Decrease tracker flag: {}'.format(self.descrease_tracker_flag))
		print('New embedding(with pose) flag: {}'.format(self.new_embedding_flag))
		print('----------------------------------------')
		
	#initialize the first frame of this video
	def init_tracker(self, frame, bbox_list):
		self.new_id_flag=0
		self.track_id_dict= dict()
		if self.tracker_flag:
			self.tracker.clear_data()
		#conver bgr(opencv) to rgb
		rgb_frame = frame[:,:,::-1]
		bbox_list, keypoint_list = self.oks_filter(bbox_list, frame)
		#print(bbox_list)
		for bbox in bbox_list:
			self.create_id(frame, rgb_frame, bbox)
		bbox_list=[]
		#pose_list=[]
		for id,item in self.track_id_dict.items():
			bbox = item['bbox_and_score'] + [id]
			bbox_list.append(bbox)
		#	pose_position, pose_value, pose_heatmap = self.pose_detect(frame, bbox)
		#	pose_info = np.hstack(pose_postion, pose_value)
		return bbox_list
		
	def oks_filter(self, det_list, frame):
		keypoint_list = []
		for bbox in det_list:
			center, scale = self.posenet.x1y1x2y2_to_cs(bbox[0:4])
			area = np.prod(scale*200,1)
			pred = np.zeros((15,3), dtype=np.float32)
			pose_positions, pose_vals, pose_heatmaps = self.pose_detect(frame, bbox)
			#print(pose_vals)
			#posa_vals = np.expand_dims(pose_vals, axis=1)
			pred[:,0:2] = pose_positions
			pred[:,2] = pose_vals
			score_all,valid_num = 0, 0
			for i in range(15):
				score_i = pose_vals[i]
				if score_i >= 0.2:
					score_all += score_i
					valid_num += 1
			if valid_num!=0:
				new_score = score_all/valid_num *bbox[4]
			else:
				new_score = 0
			keypoint_dict={'score':new_score, 'area':area, 'keypoints':pred}
			keypoint_list.append(keypoint_dict)
		keep = self.matcher.oks_nms(keypoint_list, thresh= self.oks_thresh)
		new_det_list = [det_list[i] for i in keep]
		new_keypoint_list = [keypoint_list[i] for i in keep]
		return new_det_list, new_keypoint_list
		
	def create_id(self, frame, rgb_frame, bbox):
		score = bbox[4]
		bbox = bbox[0:4]
		track_id = self.new_id_flag
		if self.new_embedding_flag:
			feature = self.embedding(frame, bbox)
		else:
			feature = self.embedder.embedding(frame, bbox)
		self.track_id_dict[track_id]={'bbox_and_score':bbox+[score],'feature':feature, 'frame_flag':1, 'exist':True}
		if self.tracker_flag:
			self.update_tracker(rgb_frame, bbox, track_id)
		self.new_id_flag += 1
		
	def update_id(self, frame, rgb_frame, det_bbox, track_id):
		bbox, score = det_bbox[0:4], det_bbox[4]
		if self.new_embedding_flag:
			feature = np.array(self.embedding(frame, bbox))
		else:
			feature = np.array(self.embedder.embedding(frame, bbox))
		
		former_track_dict = self.track_id_dict[track_id]
		former_frame_flag, former_feature = former_track_dict['frame_flag'], np.array(former_track_dict['feature'])
		now_frame_flag = former_frame_flag+1
		now_feature = feature.tolist()
		self.track_id_dict[track_id]={'bbox_and_score':det_bbox,'feature':now_feature, 'frame_flag':now_frame_flag, 'exist':True}
		if self.tracker_flag and self.tracker_update_flag:
			self.update_tracker(rgb_frame, bbox, track_id)
		
	def multi_track(self, frame):
		rgb_frame = frame[:,:,::-1]
		bbox_list = []
		for id in self.track_id_dict:
			if self.track_id_dict[id]['exist'] == False:
				continue
			bbox, score = self.tracker.track_id(rgb_frame, id)
			bbox_list.append(bbox +[score] +[id])
			self.track_id_dict[id]['bbox'] = bbox
			# if self.tracker_update_flag:
				# self.update_tracker(rgb_frame, bbox, id)
		return bbox_list
		
	def match_detection_embedding(self, detections, frame):
		rgb_frame = frame[:,:,::-1]
		#get detection feature
		if self.oks_flag:
			detections, keypoint_list = self.oks_filter(detections, frame)
		
		det_feature_list =[]
		for det_bbox in detections:
			if self.new_embedding_flag:
				det_id_feature = self.embedding(frame, det_bbox) + det_bbox
			else:
				det_id_feature = self.embedder.embedding(frame, det_bbox) + det_bbox
			det_feature_list.append(det_id_feature)
			
		#get feature for former trackers
		database_feature_list = []
		for database_id in self.track_id_dict:
			database_id_feature = self.track_id_dict[database_id]['feature'] + [database_id]
			database_feature_list.append(database_id_feature)
			
		#match the detection and tracklist
		embedding_matches, \
		embedding_unmatched_detections,\
		embedding_unmatched_trackers = self.matcher.associate_detections_to_trackers_embedding(det_feature_list, 
																								database_feature_list,
																								distance_threshold = self.embedding_match_thresh)
		
		for match in embedding_matches:
			det_index, track_index = match
			det_bbox = det_feature_list[det_index][2048:]
			update_id = database_feature_list[track_index][2048]
			self.update_id(frame, rgb_frame, det_bbox, update_id) 
			
		for new_index in embedding_unmatched_detections:
			det_bbox = det_feature_list[new_index][2048:]
			det_score = det_bbox[4]
			pose_position, pose_value, pose_heatmap = self.pose_detect(frame, det_bbox)
			#print('judge')
			if det_score >= self.effective_detection_thresh and np.sum(pose_value >= self.effective_keypoints_thresh) >= self.effective_keypoints_number:
				#print(True)
				self.create_id(frame, rgb_frame, det_bbox)
			
		#delete unuseful trackers
		for delete_index in embedding_unmatched_trackers:
			delete_id = database_feature_list[delete_index][2048]
			del self.track_id_dict[delete_id]
			if self.tracker_flag:
				self.tracker.delete_id(delete_id)
		
		bbox_list = []
		for id,item in self.track_id_dict.items():
			if item['exist']==True:
				bbox_list.append(item['bbox_and_score']+[id])
		#print(bbox_list)
		return bbox_list

	def match_detection_iou(self, detections, frame):
		rgb_frame = frame[:,:,::-1]
		
		if self.oks_flag:
			detections, keypoint_list = self.oks_filter(detections, frame)
		#get feature for former trackers
		database_bbox_list = []
		for database_id in self.track_id_dict:
			database_id_bbox = self.track_id_dict[database_id]['bbox_and_score']+[database_id]
			database_bbox_list.append(database_id_bbox)
		
		matches, unmatched_detections, unmatched_trackers = self.matcher.associate_detections_to_trackers_iou(detections, database_bbox_list, iou_threshold = self.iou_match_thresh)
		
		#update the matched trackers with detection bbox
		for match in matches:
			det_index, track_index = match
			det_bbox = detections[det_index]
			update_id = database_bbox_list[track_index][5]
			self.track_id_dict[update_id]['bbox_and_score'] = det_bbox[0:5]
			
		#create new index for unmatched_detections
		for new_index in unmatched_detections:
			det_bbox = detections[new_index]
			det_score = det_bbox[4]
			pose_position, pose_value, pose_heatmap = self.pose_detect(frame, det_bbox)
			if det_score >= self.effective_detection_thresh and np.sum(pose_value >= self.effective_keypoints_thresh) >= self.effective_keypoints_number:
				self.create_id(frame, rgb_frame, det_bbox)
			
		#delete unuseful index for unmatched_trackers
		for delete_index in unmatched_trackers:
			delete_id = database_bbox_list[delete_index][5]
			del self.track_id_dict[delete_id]
			if self.tracker_flag:
				self.tracker.delete_id(delete_id)
			
		bbox_list = []
		for id,item in self.track_id_dict.items():
			if item['exist']==True:
				bbox_list.append(item['bbox_and_score']+[id])
		#print(bbox_list)
		return bbox_list
		
	def match_detection_iou_embedding(self, detections, frame):
		rgb_frame = frame[:,:,::-1]
			
		if self.oks_flag:
			detections, keypoint_list = self.oks_filter(detections, frame)
		
		#get feature for former trackers
		database_bbox_list = []
		for database_id in self.track_id_dict:
			database_id_bbox = self.track_id_dict[database_id]['bbox_and_score']+[database_id]
			database_bbox_list.append(database_id_bbox)
		matches, unmatched_detections, unmatched_trackers = self.matcher.associate_detections_to_trackers_iou(detections, database_bbox_list, iou_threshold = self.iou_match_thresh)
		
		#update the matched trackers with detection bbox
		for match in matches:
			det_index, track_index = match
			det_bbox = detections[det_index]
			update_id = database_bbox_list[track_index][5]
			self.update_id(frame, rgb_frame, det_bbox, update_id)
			
		#create new index for unmatched_detections
		det_feature_list =[]
		for new_index in unmatched_detections:
			det_bbox = detections[new_index]
			det_score = det_bbox[4]
			if self.new_embedding_flag:
				det_id_feature = self.embedding(frame, det_bbox) + det_bbox
			else:
				det_id_feature = self.embedder.embedding(frame, det_bbox) + det_bbox
			det_feature_list.append(det_id_feature)
			
		track_feature_list = []
		for delete_index in unmatched_trackers:
			track_bbox = database_bbox_list[delete_index]
			track_score, delete_id = track_bbox[4], track_bbox[5]
			delete_id_feature = self.track_id_dict[delete_id]['feature'] + [delete_id]
			track_feature_list.append(delete_id_feature)
			
		embedding_matches, \
		embedding_unmatched_detections,\
		embedding_unmatched_trackers = self.matcher.associate_detections_to_trackers_embedding(det_feature_list, 
																								track_feature_list,
																								distance_threshold = self.embedding_match_thresh)
		
		#update matched embedding detections and former tracking feature
		for match in embedding_matches:
			det_index, track_index = match
			det_bbox = det_feature_list[det_index][self.feature_length:]
			update_id = track_feature_list[track_index][self.feature_length]
			self.update_id(frame, rgb_frame, det_bbox, update_id) 
			
		#create new id for unmatched detections
		for new_index in embedding_unmatched_detections:
			det_bbox = det_feature_list[new_index][self.feature_length:]
			det_score = det_bbox[4]
			pose_position, pose_value, pose_heatmap = self.pose_detect(frame, det_bbox)
			if det_score >= self.effective_detection_thresh and np.sum(pose_value >= self.effective_keypoints_thresh) >= self.effective_keypoints_number:
				self.create_id(frame, rgb_frame, det_bbox)
			
		#delete unuseful index for unmatched_trackers
		for delete_index in embedding_unmatched_trackers:
			delete_id = track_feature_list[delete_index][self.feature_length]
			del self.track_id_dict[delete_id]
			if self.tracker_flag:
				self.tracker.delete_id(delete_id)
			
		bbox_list = []
		for id,item in self.track_id_dict.items():
			if item['exist']==True:
				bbox_list.append(item['bbox_and_score']+[id])
		return bbox_list
		
	def match_detection_tracking_oks_iou_embedding(self, detections, track_list, frame):
		rgb_frame = frame[:,:,::-1]
			
		#print(detections)
		for track in track_list:
			track_score = track[4]
			if track_score >= self.effective_track_thresh:
				if self.descrease_tracker_flag:
					track_score -= 0.35
				detections.append(track[0:4]+[track_score])
		#print(detections)
		if self.oks_flag:
			detections, keypoint_list = self.oks_filter(detections, frame)
		#decrease the tracking score
			
		#get feature for former trackers
		database_bbox_list = []
		for database_id in self.track_id_dict:
			database_id_bbox = self.track_id_dict[database_id]['bbox_and_score']+[database_id]
			database_bbox_list.append(database_id_bbox)
		
		matches, unmatched_detections, unmatched_trackers = self.matcher.associate_detections_to_trackers_iou(detections, database_bbox_list, iou_threshold = self.iou_match_thresh)
		
		#update the matched trackers with detection bbox
		for match in matches:
			det_index, track_index = match
			det_bbox = detections[det_index]
			update_id = database_bbox_list[track_index][5]
			self.update_id(frame, rgb_frame, det_bbox, update_id)
			
		#create new index for unmatched_detections
		det_feature_list =[]
		for new_index in unmatched_detections:
			det_bbox = detections[new_index]
			det_score = det_bbox[4]
			if self.new_embedding_flag:
				det_id_feature = self.embedding(frame, det_bbox) + det_bbox
			else:
				det_id_feature = self.embedder.embedding(frame, det_bbox) + det_bbox
			det_feature_list.append(det_id_feature)
			
		track_feature_list = []
		for delete_index in unmatched_trackers:
			track_bbox = database_bbox_list[delete_index]
			track_score, delete_id = track_bbox[4], track_bbox[5]
			delete_id_feature = self.track_id_dict[delete_id]['feature'] + [delete_id]
			track_feature_list.append(delete_id_feature)
			
		embedding_matches, \
		embedding_unmatched_detections,\
		embedding_unmatched_trackers = self.matcher.associate_detections_to_trackers_embedding(det_feature_list, 
																								track_feature_list,
																								distance_threshold = self.embedding_match_thresh)
		
		#update matched embedding detections and former tracking feature
		for match in embedding_matches:
			det_index, track_index = match
			det_bbox = det_feature_list[det_index][2048:]
			update_id = track_feature_list[track_index][2048]
			self.update_id(frame, rgb_frame, det_bbox, update_id) 
			
		#create new id for unmatched detections
		for new_index in embedding_unmatched_detections:
			det_bbox = det_feature_list[new_index][2048:]
			det_score = det_bbox[4]
			pose_position, pose_value, pose_heatmap = self.pose_detect(frame, det_bbox)
			if det_score >= self.effective_detection_thresh and np.sum(pose_value >= self.effective_keypoints_thresh) >= self.effective_keypoints_number:
				self.create_id(frame, rgb_frame, det_bbox)
			
		#delete unuseful index for unmatched_trackers
		for delete_index in embedding_unmatched_trackers:
			delete_id = track_feature_list[delete_index][2048]
			del self.track_id_dict[delete_id]
			if self.tracker_flag:
				self.tracker.delete_id(delete_id)
			
		bbox_list = []
		for id,item in self.track_id_dict.items():
			if item['exist']==True:
				bbox_list.append(item['bbox_and_score']+[id])
		return bbox_list
		
		
	def match_detection_tracking_iou_embedding(self, detections, track_list, frame):
		rgb_frame = frame[:,:,::-1]
		
		if self.oks_flag:
			detections, keypoint_list = self.oks_filter(detections, frame)
		
		matches, unmatched_detections, unmatched_trackers = self.matcher.associate_detections_to_trackers_iou(detections, track_list, iou_threshold = self.iou_match_thresh)
		
		#update the matched trackers with detection bbox
		for match in matches:
			det_index, track_index = match
			det_bbox = detections[det_index]
			update_id = track_list[track_index][5]
			self.update_id(frame, rgb_frame, det_bbox, update_id)
			
		#create new index for unmatched_detections
		det_feature_list =[]
		for new_index in unmatched_detections:
			det_bbox = detections[new_index]
			det_score = det_bbox[4]
			if self.new_embedding_flag:
				det_id_feature = self.embedding(frame, det_bbox) + det_bbox
			else:
				det_id_feature = self.embedder.embedding(frame, det_bbox) + det_bbox
			det_feature_list.append(det_id_feature)
			
		track_feature_list = []
		for delete_index in unmatched_trackers:
			track_bbox = track_list[delete_index]
			track_score, delete_id = track_bbox[4], track_bbox[5]
			delete_id_feature = self.track_id_dict[delete_id]['feature'] + [delete_id]
			track_feature_list.append(delete_id_feature)
			
		#match the detection and tracklist
		embedding_matches, \
		embedding_unmatched_detections,\
		embedding_unmatched_trackers = self.matcher.associate_detections_to_trackers_embedding(det_feature_list, 
																								track_feature_list,
																								distance_threshold = self.embedding_match_thresh)
		
		for match in embedding_matches:
			det_index, track_index = match
			det_bbox = det_feature_list[det_index][2048:]
			update_id = track_feature_list[track_index][2048]
			self.update_id(frame, rgb_frame, det_bbox, update_id) 
			
		for new_index in embedding_unmatched_detections:
			det_bbox = det_feature_list[new_index][2048:]
			det_score = det_bbox[4]
			pose_position, pose_value, pose_heatmap = self.pose_detect(frame, det_bbox)
			if det_score >= self.effective_detection_thresh and np.sum(pose_value >= self.effective_keypoints_thresh) >= self.effective_keypoints_number:
				self.create_id(frame, rgb_frame, det_bbox)
			
		#delete unuseful index for unmatched_trackers
		for delete_index in embedding_unmatched_trackers:
			delete_id = track_feature_list[delete_index][2048]
			del self.track_id_dict[delete_id]
			if self.tracker_flag:
				self.tracker.delete_id(delete_id)
		
		bbox_list = []
		for id,item in self.track_id_dict.items():
			if item['exist']==True:
				bbox_list.append(item['bbox_and_score']+[id])
		return bbox_list

	def match_detection_tracking_oks_embedding(self, detections, track_list, frame):
		rgb_frame = frame[:,:,::-1]
			
		#print(detections)
		for track in track_list:
			track_score = track[4]
			if track_score >= self.effective_track_thresh:
				if self.descrease_tracker_flag:
					track_score -= 0.35
				detections.append(track[0:4]+[track_score])
		#print(detections)
		if self.oks_flag:
			detections, keypoint_list = self.oks_filter(detections, frame)	
			
		#create new index for unmatched_detections
		det_feature_list =[]
		for det in detections:
			det_bbox, det_score = det[0:4], det[4]
			if self.new_embedding_flag:
				det_id_feature = self.embedding(frame, det_bbox) + det
			else:
				det_id_feature = self.embedder.embedding(frame, det_bbox) + det
			det_feature_list.append(det_id_feature)
			
		track_feature_list = []
		for track_id in self.track_id_dict:
			track_id_feature = self.track_id_dict[track_id]['feature'] + [track_id]
			track_feature_list.append(track_id_feature)
			
		embedding_matches, \
		embedding_unmatched_detections,\
		embedding_unmatched_trackers = self.matcher.associate_detections_to_trackers_embedding(det_feature_list, 
																								track_feature_list,
																								distance_threshold = self.embedding_match_thresh)
		
		#update matched embedding detections and former tracking feature
		for match in embedding_matches:
			det_index, track_index = match
			det_bbox = det_feature_list[det_index][2048:]
			update_id = track_feature_list[track_index][2048]
			self.update_id(frame, rgb_frame, det_bbox, update_id)
			
		#create new id for unmatched detections
		for new_index in embedding_unmatched_detections:
			det_bbox = det_feature_list[new_index][2048:]
			det_score = det_bbox[4]
			pose_position, pose_value, pose_heatmap = self.pose_detect(frame, det_bbox)
			if det_score >= self.effective_detection_thresh and np.sum(pose_value >= self.effective_keypoints_thresh) >= self.effective_keypoints_number:
				self.create_id(frame, rgb_frame, det_bbox)
			
		#delete unuseful index for unmatched_trackers
		for delete_index in embedding_unmatched_trackers:
			delete_id = track_feature_list[delete_index][2048]
			del self.track_id_dict[delete_id]
			if self.tracker_flag:
				self.tracker.delete_id(delete_id)
			
		bbox_list = []
		for id,item in self.track_id_dict.items():
			if item['exist']==True:
				bbox_list.append(item['bbox_and_score']+[id])
		return bbox_list

	#bbox must be format of x1y1x2y2
	def update_tracker(self, rgb_frame, bbox, track_id):
		self.tracker.update_data_dict(rgb_frame, bbox, track_id)
		
	def pose_detect(self, im, bbox):
		return self.posenet.detect_pose(im, bbox)
		
	def embedding(self, frame, bbox):
		feature = self.posenet.embedding(frame,bbox)
		return feature	

if __name__ == "__main__":
    #track_and_detect(model_path='/export/home/zby/SiamFC/models/output/siamfc_35.pth')
	mytrack = Track_And_Detect(gpu_id=[0,0,0])
	
