#!/usr/bin/env python

import math, numpy as np
#from enthought.mayavi import mlab
import matplotlib.pyplot as pp
import matplotlib.cm as cm

import scipy.ndimage as ni

import roslib; roslib.load_manifest('sandbox_tapo_darpa_m3')
import rospy
#import hrl_lib.mayavi2_util as mu
import hrl_lib.viz as hv
import hrl_lib.util as ut
import hrl_lib.matplotlib_util as mpu 
import pickle

from visualization_msgs.msg import Marker
from visualization_msgs.msg import MarkerArray
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
from hrl_msgs.msg import FloatArrayBare
from geometry_msgs.msg import Point
from geometry_msgs.msg import Vector3

def callback(data):
    rospy.loginfo('Getting data!')
    force_vectors = np.row_stack([data.forces_x, data.forces_y, data.forces_z])
    fmags = ut.norm(force_vectors)
    force_arr = fmags.reshape((16,24))
        
    # Calculating no. of contact regions with hand-tuned force threshold
    contact_regions = fmags > 0.01
    lb,ls = ni.label(contact_regions)
    total_contact = ni.sum(lb)

    # Calculating force data for contact with hand-tuned force threshold
    #total_forces = ni.sum(fmags,lb)
    #mean_forces = ni.mean(fmags,lb)
    temp = fmags*lb
    total_forces = ni.sum(temp)
    mean_forces = total_forces/total_contact
    max_forces = np.max(temp)
 
    global time
    time = time + 0.01
    time_instant_data = [time,total_forces,mean_forces,max_forces,total_contact]
    global time_varying_data
    time_varying_data = np.row_stack([time_varying_data, time_instant_data])
    
def savedata():
    rospy.loginfo('Saving data!')
    ut.save_pickle(time_varying_data, '/home/tapo/svn/robot1_data/usr/tapo/data/Variable_Stiffness_Variable_Velocity/low_vel_low_stiff/soft_movable/sponge_movable/trial_5.pkl')

def plotdata():
    rospy.loginfo('Plotting data!')
    ta = ut.load_pickle('/home/tapo/svn/robot1_data/usr/tapo/data/Variable_Stiffness_Variable_Velocity/low_vel_low_stiff/soft_movable/sponge_movable/trial_5.pkl')

    mpu.figure(1)
    pp.title('Time-Varying Force')
    pp.xlabel('Time (s)')
    pp.ylabel('Total Force')
    pp.plot(ta[:,0], ta[:,1])

    mpu.figure(2)
    pp.title('Time-Varying Force')
    pp.xlabel('Time (s)')
    pp.ylabel('Mean Force')
    pp.plot(ta[:,0], ta[:,2])

    mpu.figure(3)
    pp.title('Time-Varying Force')
    pp.xlabel('Time (s)')
    pp.ylabel('Max Force')
    pp.plot(ta[:,0], ta[:,3])

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

def getdata():
    rospy.init_node('getdata', anonymous=True)
    rospy.Subscriber("/skin_patch_forearm_right/taxels/forces", TaxelArray_Meka, callback)
    rospy.spin()

if __name__ == '__main__':
    time = 0
    time_varying_data = [0,0,0,0,0]
    getdata()
    savedata()
    plotdata()
    pp.show()


