#!/usr/bin/env python

import rospy
import numpy as np # http://docs.scipy.org/doc/numpy/reference/routines.html
from copy import copy
import tf
import os
from flexbe_core import EventState, Logger
from rosbag import Bag


"""
Created on 01/05/2015

@author: Moritz Schappler

Sources: 
[1] Kroenig: Kraft-/Drehmomentsensoren-Kalibrierung

TODOs:
Move chain loop inside the opening of the bag file
"""

class CalculateForceTorqueCalibration(EventState):
	"""
	Calculate a Calibration Dataset for the force torque sensor

	-- calibration_chain 			string[]	Names of the chain of the ft sensors to be calibrated.
	-- settlingtime		float		settling time after each pose
	-- 	static_calibration_data	dict()	dictionary, where keys are the entries of calibration_chain and elements are lists with 4 entries: Mass, center of mass from previous calibrations
	
	#> trajectories_command JointTrajectory{}	A dictionary where the keys are ['left_arm', 'right_arm', ...] and each has a trajectory as the value.
	#> bag_filename		 			string[]	path to the bag file to be read for calibration

	># ft_calib_data 	float[] 	calibration data for the given ft sensors.

	<= done 						Successfully calculated the calibration of all ft sensors.
	<= failed 						error in the calculation.

	"""

	def __init__(self, calibration_chain, settlingtime, static_calibration_data):
		"""Constructor"""
		super(CalculateForceTorqueCalibration, self).__init__(outcomes=['done', 'failed'],
												input_keys=['bag_filename', 'trajectories_command'], 
												output_keys=['ft_calib_data'])

		self._failed = False
		self._done = False
		
		# inputs
		self._calibration_chain = calibration_chain
		self._settlingtime = settlingtime
		self._static_calibration_data = static_calibration_data

		# prepare tmpp variable for output
		self._ft_calib_data = dict() # dictionary like userdata.ft_calib_data
		for chain in calibration_chain:
			self._ft_calib_data[chain] = np.zeros((10,1))

	def execute(self, userdata):
		if self._failed:
			return 'failed'
		if self._done:
			return 'done'

	def on_enter(self, userdata):
		self._failed = False
		try:
			# Initialization
		
			# userdata
			self._bag_filename = userdata.bag_filename
			self._trajectories_command = userdata.trajectories_command
			
			# joint definitions
			l_arm_range = range(16,23);
			r_arm_range = range(23,30);
			atlasJointNames = [
			    'back_bkz', 'back_bky', 'back_bkx', 'neck_ry',
			    'l_leg_hpz', 'l_leg_hpx', 'l_leg_hpy', 'l_leg_kny', 'l_leg_aky', 'l_leg_akx',
			    'r_leg_hpz', 'r_leg_hpx', 'r_leg_hpy', 'r_leg_kny', 'r_leg_aky', 'r_leg_akx',
			    'l_arm_shz', 'l_arm_shx', 'l_arm_ely', 'l_arm_elx', 'l_arm_wry', 'l_arm_wrx', 'l_arm_wry2',
			    'r_arm_shz', 'r_arm_shx', 'r_arm_ely', 'r_arm_elx', 'r_arm_wry', 'r_arm_wrx', 'r_arm_wry2']
			joint_position_cmd = [0]*len(l_arm_range)
			# starting and stopping time of measurement data to be taken into account for cailbration
			time_start = 0
			time_end = 0
			time_tf = 0
			
			# gravity vector in world frame
			g_world = np.resize(np.array([0,0,-9.81]),(3,1))
		
			# take this number of data points for each pose.
			max_data_points_per_pose = 100
			
			# rotation of the force torque sensor in world frame (quaternions)
			ft_rotation = [0,0,0,0]
			
			# loop over all ft sensors to calibrate
			
			for chain in self._calibration_chain: 
				tf_chain=['/pelvis', 'ltorso', 'mtorso', 'utorso']
				if chain == 'left_arm':
					joint_range = l_arm_range
					tfname = 'l_hand'
					tf_chain.extend(['l_clav', 'l_scap', 'l_uarm', 'l_larm', 'l_ufarm', 'l_lfarm', 'l_hand'])
				elif chain == 'right_arm':
					joint_range = r_arm_range
					tf_chain.extend(['r_clav', 'r_scap', 'r_uarm', 'r_larm', 'r_ufarm', 'r_lfarm', 'r_hand'])
					
					tfname = 'r_hand'
				else:
					Logger.logwarn('CalculateForceTorqueCalibration: Undefined chain %s', chain)

				# initialize transformations
				tf_data = [] # frame to frame transformation with numerical indexes from the tf_chain
				tf_data_cum = [] # world to frame transformation with numerical indexes from the tf_chain

				for i in range(len(tf_chain)):
					tf_data.append([0,0,0,0])
					tf_data_cum.append([0,0,0,0])
									
				# get number of poses from the commanded trajectories. 2 Points per Pose
				number_of_poses = len(self._trajectories_command[chain].points) / 2
				# define information matrix for calibration
				InfMat = np.zeros((6*max_data_points_per_pose*number_of_poses, 10))
				InfMat_i = 0 # Index in this Matrix
				# definie measurement vector for calibration
				MeasVec = np.zeros((6*max_data_points_per_pose*number_of_poses, 1))
				
				# commanded joint trajectory for current chain
				# all commanded positions in this trajectory
				current_traj_cmd = self._trajectories_command[chain] 
				
				# flag if the time interval was already set:
				timesetflag = False
				
				point_index_cmd = 0 # start with index 0
				transformation_available = False
				
				# read time series from bag file
				# check in desired input trajectory, at which time a position is reached
				# take the data in a defined time period after the commanded new position
				bag_from_robot = Bag(os.path.expanduser(self._bag_filename)) # open bag file
				Logger.loginfo('CalculateForceTorqueCalibration: Calibrate %s from %s. Using %d different positions from trajectory command. Expecting 2 commanded positions per pose' % (chain, self._bag_filename, len(current_traj_cmd.points)) )
				for topic, msg, t in bag_from_robot.read_messages(topics=['/flor/controller/atlas_state', '/tf', '/flor/controller/joint_command']):
					# Check if all desired poses have been reached
					if point_index_cmd > len(current_traj_cmd.points)-1:
						break
					current_position_cmd = current_traj_cmd.points[point_index_cmd].positions
					####################################################
					# Check tf message: remember last transformation from world to ft sensor
					if topic == '/tf':
			  			# loop over all transformations inside this tf message
						for j in range(len(msg.transforms)):
							data = msg.transforms[j]
							tr = data.transform # current transformation
							header = data.header
							# check if frame matches on of the saved frames
							for i in range(len(tf_chain)):
								if data.child_frame_id == tf_chain[i]:
									tf_data[i] = np.array([tr.rotation.x, tr.rotation.y, tr.rotation.z, tr.rotation.w])

						# calculate the newest transformation to the force torque sensor
						tf_data_cum[0] = tf_data[0]