
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 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
from m3skin_ros.msg import TaxelArray as TaxelArray_Meka


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

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

    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
    force_arr = fmags_instant.reshape((16,24))
    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 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)

def processdata():
    rospy.loginfo('Processing data!')
    global fmags
    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)
                else:
                    temp_contact_motion.append(abs(lin.norm(global_contact_vector[key][i] - init_contact_store)))
            else:
                if len(temp_force_store) > 0:
                    print "Broke Contact !"
                    savedata(temp_force_store, temp_contact_motion)
                    temp_force_store = []
                    temp_contact_motion = []
                    init_contact = 0
                    init_contact_store = 0.0
                    

def savedata(force, motion):
    global trial_index
    global directory 
    time = []
    contact_area = []
    if len(force) > 10:
        rospy.loginfo('Saving data!')
        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], directory + '/trial_' + 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
    for trial_num in range(1, trial_index):
        ta = ut.load_pickle(directory + '/trial_' + np.str(trial_num) +'.pkl')
        
        mpu.figure(3*trial_num-2)
        pp.title('Time-Varying Force')
        pp.xlabel('Time (s)')
        pp.ylabel('Max Force')
        pp.plot(ta[0], ta[1])
        pp.grid('on') 

        mpu.figure(3*trial_num-1)
        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(3*trial_num)
        pp.title('Point Tracker')
        pp.xlabel('Time (s)')
        pp.ylabel('Contact Point Distance')
        pp.plot(ta[0], ta[3])
        pp.grid('on')
    
    
def getdata():
    rospy.loginfo('Initializing the Node !')
    rospy.init_node('Taxel_Based_Segmentation', anonymous=True)
    tf_lstnr = tf.TransformListener()
    rospy.loginfo('Waiting to Subscribe to the Skin Message...')
    rospy.Subscriber("/skin_patch_forearm_right/taxels/forces", TaxelArray_Meka, callback, callback_args = (tf_lstnr))
    rospy.spin()


if __name__ == '__main__':

    # Global Params
    trial_index = 1
    num = 54
    directory = '/home/tapo/svn/robot1_data/usr/tapo/data/rapid_categorization/Taxel_Based/Trunk/'+np.str(num)

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

    # Function Calls
    getdata()
    processdata()
    #plotdata()
    #pp.show()
