#!/usr/bin/env python

# Collect PR2 sleeve data for model estimation

import math, numpy as np
import matplotlib.pyplot as pp

import scipy.linalg as lin
import scipy.ndimage as ni

import roslib; roslib.load_manifest('sandbox_tapo_darpa_m3')
import rospy
import tf
import os
import optparse

import hrl_lib.util as ut
import hrl_lib.transforms as tr
import hrl_lib.matplotlib_util as mpu 
import pickle

from hrl_haptic_manipulation_in_clutter_msgs.msg import SkinContact
from hrl_haptic_manipulation_in_clutter_msgs.msg import TaxelArray as TaxelArray_Meka

import sys

def callback(data, callback_args):
    #rospy.loginfo('Getting data!')

    # Fixing Transforms    
    tf_lstnr = callback_args
    sc = SkinContact()
    sc.header.frame_id = '/odom_combined' #'/torso_lift_link' # has to be this and no other coord frame.
    sc.header.stamp = data.header.stamp

    tf_lstnr.waitForTransform(sc.header.frame_id, data.header.frame_id, rospy.Time(0), rospy.Duration(40.0))
    
    try:
    	tf_lstnr.waitForTransform(sc.header.frame_id, data.header.frame_id, rospy.Time(0), rospy.Duration(40.0))
    	t1, q1 = tf_lstnr.lookupTransform(sc.header.frame_id,
                                      data.header.frame_id,
                                      rospy.Time(0))
                                     
    	t1 = np.matrix(t1).reshape(3,1)
    	r1 = tr.quaternion_to_matrix(q1)

    	# Gathering Force Data
    	force_vectors = np.row_stack([data.values_x, data.values_y, data.values_z])
    	fmags_instant = ut.norm(force_vectors)
    	threshold = 0.01
    	fmags_tuned = fmags_instant - threshold
    	fmags_tuned[np.where(fmags_tuned<0)]=0
    	fmags_instant_tuned = fmags_tuned
    	global fmags
    	for i in range(len(fmags_instant_tuned)):
            fmags[i].append(fmags_instant_tuned[i])

        # Gathering deformation data for model estimation "Z-coordinate in taxel frame"
        global global_deformation
        for i in range(len(fmags_instant_tuned)):
            global_deformation[i].append(data.centers_z[i])
        
    	# Gathering Contact Data for Haptic Mapping
    	global global_contact_vector
    	for i in range(len(fmags_instant_tuned)):
            global_contact_vector[i].append(r1*((np.column_stack([data.centers_x[i], data.centers_y[i], data.centers_z[i]])).T) + t1)
    except (tf.Exception, tf.LookupException, tf.ConnectivityException):
        pass

def processdata():
    global sleeve
    str_process = 'Processing data for ' + sleeve + '!'
    rospy.loginfo(str_process)
    global fmags
    global global_deformation
    global global_contact_vector
    
    for key in fmags:
        temp_force_store = []
        temp_contact_motion = []
        init_contact = 0
        init_contact_store = 0.0
        max_temp = max(fmags[key])
        #if max_temp > 0.0:
            #print key
            #print '#####'
        for i in range(len(fmags[key])):
            if fmags[key][i] > 0.0:
                init_contact = init_contact + 1
                temp_force_store.append(fmags[key][i])
                if init_contact == 1:
                    #print "Started Contact !"
                    init_contact_store = global_contact_vector[key][i]
                    temp_contact_motion.append(0.0)
                    init_deformation_store = global_deformation[key][i]
                    temp_deformation.append(0.0)
                else:
                    temp_contact_motion.append(abs(lin.norm(global_contact_vector[key][i] - init_contact_store)))
                    temp_deformation.append(global_deformation[key][i]-init_deformation_store)
            else:
                if len(temp_force_store) > 0:
                    #print "Broke Contact !"
                    savedata(temp_force_store, temp_contact_motion, temp_deformation)
                    temp_force_store = []
                    temp_contact_motion = []
                    temp_deformation = []
                    init_contact = 0
                    init_contact_store = 0.0
                    

def savedata(force, motion, deformation):
    global trial_index
    global directory
    global sleeve 
    time = []
    contact_area = []
    if len(force) > 10:
        str_save = 'Saving data for ' + sleeve + '!'
        rospy.loginfo(str_save)
        time_len = len(force)
        while len(time) < time_len:
            if len(time) == 0:
                time.append(0.0)
                contact_area.append(1.0)
            else:
                time.append(time[len(time)-1] + 0.01)
                contact_area.append(1.0)
        
        if not os.path.exists(directory):
            os.makedirs(directory)
        ut.save_pickle([time, force, contact_area, motion, deformation], directory + '/trial_' + sleeve + np.str(trial_index) +'.pkl')
        trial_index = trial_index + 1
    #else:
        #print "Too few samples, Not saving the data"

def plotdata():
    rospy.loginfo('Plotting data!')
    global directory
    global trial_index
    global sleeve
    for trial_num in range(1, trial_index):
        ta = ut.load_pickle(directory + '/trial_' + sleeve + np.str(trial_num) +'.pkl')
        
        mpu.figure(4*trial_num-3)
        pp.title('Time-Varying Force')
        pp.xlabel('Time (s)')
        pp.ylabel('Max Force')
        pp.plot(ta[0], ta[1])
        pp.grid('on') 

        mpu.figure(4*trial_num-2)
        pp.title('Time-Varying Contact')
        pp.xlabel('Time (s)')
        pp.ylabel('No. of Contact Regions')
        pp.plot(ta[0], ta[2])
        pp.grid('on') 

        mpu.figure(4*trial_num-1)
        pp.title('Point Tracker')
        pp.xlabel('Time (s)')
        pp.ylabel('Contact Point Distance')
        pp.plot(ta[0], ta[3])
        pp.grid('on')
        
        mpu.figure(4*trial_num)
        pp.title('Deformation')
        pp.xlabel('Time (s)')
        pp.ylabel('Deformation Perpendicular to Taxel')
        pp.plot(ta[0], ta[4])
        pp.grid('on')
    
    
def getdata(node_name, topic_subscriber, sleeve):
    str_init = 'Initializing the Node for ' + sleeve + '!'
    str_sub  = 'Waiting to Subscribe to the Skin Message for ' + sleeve
    rospy.loginfo(str_init)
    rospy.init_node(node_name, anonymous=True)
    tf_lstnr = tf.TransformListener()
    rospy.loginfo(str_sub)
    rospy.Subscriber(topic_subscriber, TaxelArray_Meka, callback, callback_args = (tf_lstnr))
    rospy.spin()


if __name__ == '__main__':

    p = optparse.OptionParser()
    
    p.add_option('--upperarm', action='store_true', dest='upperarm', help='node for the upperarm taxels of PR2 sleeve')
    p.add_option('--forearm', action='store_true', dest='forearm', help='node for the forearm taxels of PR2 sleeve')
    p.add_option('--gripper_left_link', action='store_true', dest='gripper_left_link', help='node for the gripper left link taxels of PR2 sleeve')
    p.add_option('--gripper_right_link', action='store_true', dest='gripper_right_link', help='node for the gripper right link taxels of PR2 sleeve')
    p.add_option('--gripper_palm', action='store_true', dest='gripper_palm', help='node for the gripper palm taxels of PR2 sleeve')
    p.add_option('--fingertip_left', action='store_true', dest='fingertip_left', help='node for the fingertip left taxels of PR2 pps')
    p.add_option('--fingertip_right', action='store_true', dest='fingertip_right', help='node for the fingertip right taxels of PR2 pps')
    p.add_option('--exp_num', action='store', dest='exp_num', default = None, help='specify the experiment number', type='float')
    
    opt, args = p.parse_args()

    if opt.upperarm and opt.exp_num:
        node_name = "PR2_upperarm_collect_data"
        topic_subscriber = "/pr2_fabric_upperarm_sensor/taxels/forces"
        topic_publisher = 'demo_upperarm_taxelarray_with_cost'
        num_taxels = 4
        sleeve = 'upperarm'
        num = opt.exp_num
    elif opt.forearm and opt.exp_num:
        node_name = "PR2_forearm_collect_data"
        topic_subscriber = "/pr2_fabric_forearm_sensor/taxels/forces"
        topic_publisher = 'demo_forearm_taxelarray_with_cost'
        num_taxels = 22
        sleeve = 'forearm'
        num = opt.exp_num
    elif opt.gripper_left_link and opt.exp_num:
        node_name = "PR2_gripper_left_link_collect_data"
        topic_subscriber = "/pr2_fabric_gripper_left_link_sensor/taxels/forces"
        topic_publisher = 'demo_gripper_left_link_taxelarray_with_cost'
        num_taxels = 4
        sleeve = 'gripper_left_link'
        num = opt.exp_num
    elif opt.gripper_right_link and opt.exp_num:
        node_name = "PR2_gripper_right_link_collect_data"
        topic_subscriber = "/pr2_fabric_gripper_right_link_sensor/taxels/forces"
        topic_publisher = 'demo_gripper_right_link_taxelarray_with_cost'
        num_taxels = 4
        sleeve = 'gripper_right_link'
        num = opt.exp_num
    elif opt.gripper_palm and opt.exp_num:
        node_name = "PR2_gripper_palm_collect_data"
        topic_subscriber = "/pr2_fabric_gripper_palm_sensor/taxels/forces"
        topic_publisher = 'demo_gripper_palm_taxelarray_with_cost'
        num_taxels = 2
        sleeve = 'gripper_palm'
        num = opt.exp_num
    elif opt.fingertip_left and opt.exp_num:
        node_name = "PR2_fingertip_left_collect_data"
        topic_subscriber = "/pr2_pps_left_sensor/taxels/forces"
        topic_publisher = 'demo_fingertip_left_taxelarray_with_cost'
        num_taxels = 3
        sleeve = 'fingertip_left'
        num = opt.exp_num
    elif opt.fingertip_right and opt.exp_num:
        node_name = "PR2_fingertip_right_collect_data"
        topic_subscriber = "/pr2_pps_right_sensor/taxels/forces"
        topic_publisher = 'demo_fingertip_right_taxelarray_with_cost'
        num_taxels = 3
        sleeve = 'fingertip_right'
        num = opt.exp_num    
    else:
        rospy.logerr('Specify --exp_num and --upperarm or --forearm or --gripper_left_link or --gripper_right_link or --gripper_palm or --fingertip_left or --fingertip_right')
        sys.exit()

    # Global Params
    #obj_class = 'human/'
    obj_class = 'furniture/'
    #obj_class = 'mannequin/'
    trial_index = 1
    directory = '/home/nidrr/svn/robot1_data/usr/tapo/data/rapid_categorization/Taxel_Based/NIDRR_demo_data/model_estimation/' + obj_class + np.str(int(num))

    # Global Data dicts
    fmags = {}
    for i in range(num_taxels):
        fmags[i] = []
    global_contact_vector = {}
    for i in range(num_taxels):
        global_contact_vector[i] = []
    global_deformation = {}
    for i in range(num_taxels):
        global_deformation[i] = []

    # Function Calls
    getdata(node_name, topic_subscriber, sleeve)
    processdata()
    #plotdata()
    #pp.show()
