# Online haptic_map implementation

import pylab as pyl
import numpy as np
import matplotlib.pyplot as pp
#from enthought.mayavi import mlab

import scipy as scp
import scipy.ndimage as ni

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

import unittest
import ghmm
import ghmmwrapper
import random

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 visualization_msgs.msg import Marker
from visualization_msgs.msg import MarkerArray

import sys
sys.path.insert(0, '/home/tapo/svn/robot1_data/usr/tapo/data_code/rapid_categorization/')

from data_variable_length_force_sample import Fmat_original, temp_num_fol, temp_num_trunk

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
    fmags = np.row_stack([fmags,fmags_instant_tuned])

    # Calculating no. of contact regions with hand-tuned force threshold
    contact_regions = force_arr > threshold
    lb,ls = ni.label(contact_regions)
    total_contact = ni.sum(lb)
    
    # Calculating force data for contact with hand-tuned force threshold
    temp = force_arr*lb
    global max_force
    max_force = np.max(temp)

    # Gathering Contact Data for Haptic Mapping
    contact_vectors = np.row_stack([data.centers_x, data.centers_y, data.centers_z]).reshape((3,16,24))
    cx = ni.mean(contact_vectors[0], lb, 1) #update ls=1
    cy = ni.mean(contact_vectors[1], lb, 1) #update ls=1
    cz = ni.mean(contact_vectors[2], lb, 1) #update ls=1
    contact_vector = np.column_stack([cx, cy, cz])
    global global_contact_vector
    global_contact_vector = r1*(contact_vector.T) + t1
        
    global time_varying_data
    time_varying_data.append(max_force)
    test_data()
    pubdata()

def test_data():
    
    global FLAG_Trunk
    global FLAG_Foliage
    global FLAG_Unknown
    
    # Do Stuff For Testing which basically returns which FLAG is true
    global idx
    idx = 0 # 0 for Unknown (Red), 1 for Foliage (green), 2 for Trunk (brown)
    # For Testing
    global time_varying_data
    global max_force
    if (max_force > 0):      
        ts_obj = time_varying_data
        final_ts_obj = ghmm.EmissionSequence(F,ts_obj)

        # Find Viterbi Path

        global model_ff
        global model_tf
    
        path_ff_obj = model_ff.viterbi(final_ts_obj)
        path_tf_obj = model_tf.viterbi(final_ts_obj)
    
        print path_ff_obj[1], path_tf_obj[1]

        diff = abs(path_ff_obj[1]-path_tf_obj[1])
        obj = max(path_ff_obj[1],path_tf_obj[1])
        obj_min = min(abs(path_ff_obj[1]),abs(path_tf_obj[1]))
        if ((obj == path_ff_obj[1]) and (diff > 50)):
        #if ((obj == path_ff_obj[1]) and (obj_min > 1800)):
        #if ((obj == path_ff_obj[1])):
            print 'Foliage !'
            FLAG_Trunk = False
            FLAG_Foliage = True
            FLAG_Unknown = False
        #elif ((obj == path_tf_obj[1]) and (obj_min > 1800)):
        elif ((obj == path_tf_obj[1]) and (diff > 50)):
        #elif ((obj == path_tf_obj[1])):
            print 'Trunk !'
            FLAG_Trunk = True
            FLAG_Foliage = False
            FLAG_Unknown = False
        else:
            print 'Unknown'
            FLAG_Trunk = False
            FLAG_Leaf = False
            FLAG_Unknown = True
    else:
        print 'Unknown'
        FLAG_Trunk = False
        FLAG_Leaf = False
        FLAG_Unknown = True
        time_varying_data = [0]     

    if FLAG_Trunk:
        idx = 2
    elif FLAG_Foliage:
        idx = 1
    else: 
        idx = 0

def getdata():
    rospy.loginfo('Initializing the Node !')
    rospy.init_node('Online_Haptic_Map_Builder', 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()

def pubdata():
    rospy.loginfo('Publishing data')    
    marker = Marker()
    marker.ns = 'Haptic_Map_Markers'
    marker.header.frame_id = '/torso_lift_link'
    marker.type = marker.SPHERE
    marker.action = marker.ADD
    marker.scale.x = 0.02
    marker.scale.y = 0.02
    marker.scale.z = 0.02
    if idx == 1:
        # Green for Foliage
        marker.color.a = 1.0;
        marker.color.r = 0.0;
        marker.color.g = 1.0;
        marker.color.b = 0.0;
    elif idx == 2:
        # Brown for Trunk
        marker.color.a = 1.0;
        marker.color.r = 0.5;
        marker.color.g = 0.25;
        marker.color.b = 0.125;
    else:
        # Red for Unknown
        marker.color.a = 1.0;
        marker.color.r = 1.0;
        marker.color.g = 0.0;
        marker.color.b = 0.0;
    marker.pose.orientation.w = 1.0
    marker.pose.position.x = global_contact_vector[0]
    marker.pose.position.y = global_contact_vector[1]
    marker.pose.position.z = global_contact_vector[2]

    markerArray.markers.append(marker)

    # Renumber the marker IDs
    id = 0
    for m in markerArray.markers:
        m.id = id
        id += 1

    # Publish the MarkerArray
    publisher.publish(markerArray)
    #rospy.sleep(0.01)
    

if __name__ == '__main__':

    topic = 'visualization_marker_array'
    publisher = rospy.Publisher(topic, MarkerArray)
    markerArray = MarkerArray()

    print "Initializing the HMM Models"

# HMM Implementation
    Fmat = Fmat_original
    Foliage_Trials = temp_num_fol
    Trunk_Trials = temp_num_trunk
    
# Getting mean / covariance

    i = 0
    number_states = 10
    feature_1_final_data = [0.0]*number_states
    state_1 = [0.0]
    while (i < Foliage_Trials):
        data_length = len(Fmat[i])
        feature_length = data_length/1
        sample_length = feature_length/number_states
        Feature_1 = Fmat[i][0:feature_length]
        if i == 0:
            j = 0
            while (j < number_states):
                feature_1_final_data[j] = Feature_1[sample_length*j:sample_length*(j+1)]
                j=j+1  
        else:
            j = 0
            while (j < number_states):
                state_1 = Feature_1[sample_length*j:sample_length*(j+1)]
                #print np.shape(state_1)
                #print np.shape(feature_1_final_data[j])
                feature_1_final_data[j] = feature_1_final_data[j]+state_1
                j=j+1 
        i = i+1

    j = 0
    mu_ff_force = np.zeros((number_states,1))
    sigma_ff = np.zeros((number_states,1))
    
    while (j < number_states):
        mu_ff_force[j] = np.mean(feature_1_final_data[j])
        sigma_ff[j] = scp.std(feature_1_final_data[j])
        j = j+1

    i = Foliage_Trials
    feature_1_final_data = [0.0]*number_states
    state_1 = [0.0]
    while (i < (Foliage_Trials + Trunk_Trials)):
        data_length = len(Fmat[i])
        feature_length = data_length/1
        sample_length = feature_length/number_states
        Feature_1 = Fmat[i][0:feature_length]
        if i == Foliage_Trials:
            j = 0
            while (j < number_states):
                feature_1_final_data[j] = Feature_1[sample_length*j:sample_length*(j+1)]
                j=j+1  
        else:
            j = 0
            while (j < number_states):
                state_1 = Feature_1[sample_length*j:sample_length*(j+1)]
                feature_1_final_data[j] = feature_1_final_data[j]+state_1
                j=j+1 
        i = i+1
    
    j = 0
    mu_tf_force = np.zeros((number_states,1))
    sigma_tf = np.zeros((number_states,1))

    while (j < number_states):
        mu_tf_force[j] = np.mean(feature_1_final_data[j])
        sigma_tf[j] = scp.std(feature_1_final_data[j])
        j = j+1

# HMM - Implementation:

# 10 Hidden States
# Max. Force(For now), Contact Area(Not now), and Contact Motion(Not Now) as Continuous Gaussian Observations from each hidden state
# Four HMM-Models for Rigid-Fixed, Soft-Fixed, Rigid-Movable, Soft-Movable
# Transition probabilities obtained as upper diagonal matrix (to be trained using Baum_Welch)
# For new objects, it is classified according to which model it represenst the closest..

    F = ghmm.Float()  # emission domain of this model

    # A - Transition Matrix
    
    if number_states == 3:
        A  = [[0.2, 0.5, 0.3],
              [0.0, 0.5, 0.5],
              [0.0, 0.0, 1.0]]
    elif number_states == 5:
        A  = [[0.2, 0.35, 0.2, 0.15, 0.1],
              [0.0, 0.2, 0.45, 0.25, 0.1],
              [0.0, 0.0, 0.2, 0.55, 0.25],
              [0.0, 0.0, 0.0, 0.2, 0.8],
              [0.0, 0.0, 0.0, 0.0, 1.0]]
    elif number_states == 10:
        A  = [[0.1, 0.25, 0.15, 0.15, 0.1, 0.05, 0.05, 0.05, 0.05, 0.05],
              [0.0, 0.1, 0.25, 0.25, 0.2, 0.1, 0.05, 0.05, 0.05, 0.05],
              [0.0, 0.0, 0.1, 0.25, 0.25, 0.2, 0.05, 0.05, 0.05, 0.05],
              [0.0, 0.0, 0.0, 0.1, 0.3, 0.30, 0.20, 0.1, 0.05, 0.05],
              [0.0, 0.0, 0.0, 0.0, 0.1, 0.30, 0.30, 0.20, 0.05, 0.05],
              [0.0, 0.0, 0.0, 0.0, 0.00, 0.1, 0.35, 0.30, 0.20, 0.05],
              [0.0, 0.0, 0.0, 0.0, 0.00, 0.00, 0.2, 0.30, 0.30, 0.20],
              [0.0, 0.0, 0.0, 0.0, 0.00, 0.00, 0.00, 0.2, 0.50, 0.30],
              [0.0, 0.0, 0.0, 0.0, 0.00, 0.00, 0.00, 0.00, 0.4, 0.60],
              [0.0, 0.0, 0.0, 0.0, 0.00, 0.00, 0.00, 0.00, 0.00, 1.00]]
    elif number_states == 15:
        A  = [[0.1, 0.25, 0.15, 0.15, 0.1, 0.05, 0.05, 0.05, 0.04, 0.01, 0.01, 0.01, 0.01, 0.01, 0.01],
              [0.0, 0.1, 0.25, 0.25, 0.2, 0.1, 0.05, 0.05, 0.04, 0.01, 0.01, 0.01, 0.01, 0.01, 0.01],
              [0.0, 0.0, 0.1, 0.25, 0.25, 0.2, 0.05, 0.05, 0.04, 0.01, 0.01, 0.01, 0.01, 0.01, 0.01],
              [0.0, 0.0, 0.0, 0.1, 0.3, 0.30, 0.20, 0.1, 0.04, 0.01, 0.01, 0.01, 0.01, 0.01, 0.01],
              [0.0, 0.0, 0.0, 0.0, 0.1, 0.30, 0.30, 0.20, 0.04, 0.01, 0.01, 0.01, 0.01, 0.01, 0.01],
              [0.0, 0.0, 0.0, 0.0, 0.00, 0.1, 0.35, 0.30, 0.15, 0.05, 0.01, 0.01, 0.01, 0.01, 0.01],
              [0.0, 0.0, 0.0, 0.0, 0.00, 0.00, 0.1, 0.30, 0.30, 0.10, 0.05, 0.05, 0.05, 0.03, 0.02],
              [0.0, 0.0, 0.0, 0.0, 0.00, 0.00, 0.00, 0.1, 0.30, 0.30, 0.10, 0.05, 0.05, 0.05, 0.05],
              [0.0, 0.0, 0.0, 0.0, 0.00, 0.00, 0.00, 0.00, 0.1, 0.30, 0.20, 0.15, 0.10, 0.10, 0.05],
              [0.0, 0.0, 0.0, 0.0, 0.00, 0.00, 0.00, 0.00, 0.00, 0.1, 0.30, 0.20, 0.15, 0.15, 0.10],
              [0.0, 0.0, 0.0, 0.0, 0.00, 0.00, 0.00, 0.00, 0.00, 0.0, 0.1, 0.30, 0.30, 0.20, 0.10],
              [0.0, 0.0, 0.0, 0.0, 0.00, 0.00, 0.00, 0.00, 0.00, 0.0, 0.0, 0.1, 0.40, 0.30, 0.20],
              [0.0, 0.0, 0.0, 0.0, 0.00, 0.00, 0.00, 0.00, 0.00, 0.0, 0.0, 0.0, 0.20, 0.50, 0.30],
              [0.0, 0.0, 0.0, 0.0, 0.00, 0.00, 0.00, 0.00, 0.00, 0.0, 0.0, 0.0, 0.00, 0.40, 0.60],
              [0.0, 0.0, 0.0, 0.0, 0.00, 0.00, 0.00, 0.00, 0.00, 0.00, 0.0, 0.0, 0.0, 0.00, 1.00]]
    elif number_states == 20:
        A  = [[0.1, 0.25, 0.15, 0.15, 0.1, 0.05, 0.05, 0.03, 0.01, 0.01, 0.01, 0.01, 0.01, 0.01, 0.01, 0.01, 0.01, 0.01, 0.01, 0.01],
              [0.0, 0.1, 0.25, 0.25, 0.2, 0.1, 0.05, 0.03, 0.01, 0.01, 0.01, 0.01, 0.01, 0.01, 0.01, 0.01, 0.01, 0.01, 0.01, 0.01],
              [0.0, 0.0, 0.1, 0.25, 0.25, 0.2, 0.05, 0.03, 0.01, 0.01, 0.01, 0.01, 0.01, 0.01, 0.01, 0.01, 0.01, 0.01, 0.01, 0.01],
              [0.0, 0.0, 0.0, 0.1, 0.3, 0.30, 0.20, 0.09, 0.01, 0.01, 0.01, 0.01, 0.01, 0.01, 0.01, 0.01, 0.01, 0.01, 0.01, 0.01],
              [0.0, 0.0, 0.0, 0.0, 0.1, 0.30, 0.30, 0.15, 0.04, 0.01, 0.01, 0.01, 0.01, 0.01, 0.01, 0.01, 0.01, 0.01, 0.01, 0.01],
              [0.0, 0.0, 0.0, 0.0, 0.00, 0.1, 0.35, 0.30, 0.10, 0.05, 0.01, 0.01, 0.01, 0.01, 0.01, 0.01, 0.01, 0.01, 0.01, 0.01],
              [0.0, 0.0, 0.0, 0.0, 0.00, 0.00, 0.1, 0.30, 0.20, 0.10, 0.05, 0.05, 0.05, 0.03, 0.02, 0.02, 0.02, 0.02, 0.02, 0.02],
              [0.0, 0.0, 0.0, 0.0, 0.00, 0.00, 0.00, 0.1, 0.30, 0.20, 0.10, 0.05, 0.05, 0.05, 0.05, 0.02, 0.02, 0.02, 0.02, 0.02],
              [0.0, 0.0, 0.0, 0.0, 0.00, 0.00, 0.00, 0.00, 0.1, 0.30, 0.20, 0.15, 0.05, 0.05, 0.05, 0.02, 0.02, 0.02, 0.02, 0.02],
              [0.0, 0.0, 0.0, 0.0, 0.00, 0.00, 0.00, 0.00, 0.00, 0.1, 0.30, 0.20, 0.15, 0.10, 0.05, 0.02, 0.02, 0.02, 0.02, 0.02],
              [0.0, 0.0, 0.0, 0.0, 0.00, 0.00, 0.00, 0.00, 0.00, 0.0, 0.1, 0.30, 0.30, 0.10, 0.10, 0.02, 0.02, 0.02, 0.02, 0.02],
              [0.0, 0.0, 0.0, 0.0, 0.00, 0.00, 0.00, 0.00, 0.00, 0.0, 0.0, 0.1, 0.40, 0.30, 0.10, 0.02, 0.02, 0.02, 0.02, 0.02], 
              [0.0, 0.0, 0.0, 0.0, 0.00, 0.00, 0.00, 0.00, 0.00, 0.0, 0.0, 0.0, 0.20, 0.40, 0.20, 0.10, 0.04, 0.02, 0.02, 0.02],
              [0.0, 0.0, 0.0, 0.0, 0.00, 0.00, 0.00, 0.00, 0.00, 0.0, 0.0, 0.0, 0.00, 0.20, 0.40, 0.20, 0.10, 0.05, 0.03, 0.02],
              [0.0, 0.0, 0.0, 0.0, 0.00, 0.00, 0.00, 0.00, 0.00, 0.00, 0.0, 0.0, 0.0, 0.00, 0.20, 0.40, 0.20, 0.10, 0.05, 0.05],
              [0.0, 0.0, 0.0, 0.0, 0.00, 0.00, 0.00, 0.00, 0.00, 0.00, 0.0, 0.0, 0.0, 0.00, 0.00, 0.20, 0.40, 0.20, 0.10, 0.10],
              [0.0, 0.0, 0.0, 0.0, 0.00, 0.00, 0.00, 0.00, 0.00, 0.00, 0.0, 0.0, 0.0, 0.00, 0.00, 0.00, 0.20, 0.40, 0.20, 0.20],
              [0.0, 0.0, 0.0, 0.0, 0.00, 0.00, 0.00, 0.00, 0.00, 0.00, 0.0, 0.0, 0.0, 0.00, 0.00, 0.00, 0.00, 0.30, 0.50, 0.20],
              [0.0, 0.0, 0.0, 0.0, 0.00, 0.00, 0.00, 0.00, 0.00, 0.00, 0.0, 0.0, 0.0, 0.00, 0.00, 0.00, 0.00, 0.00, 0.40, 0.60],
              [0.0, 0.0, 0.0, 0.0, 0.00, 0.00, 0.00, 0.00, 0.00, 0.00, 0.0, 0.0, 0.0, 0.00, 0.00, 0.00, 0.00, 0.00, 0.00, 1.00]]

    # B - Emission Matrix, parameters of emission distributions in pairs of (mu, sigma)
    
    B_ff = [0.0]*number_states
    B_tf = [0.0]*number_states
          
    for num_states in range(number_states):
        B_ff[num_states] = [mu_ff_force[num_states][0],sigma_ff[num_states][0]]
        B_tf[num_states] = [mu_tf_force[num_states][0],sigma_tf[num_states][0]]
        
       
    # pi - initial probabilities per state
    if number_states == 3:
        pi = [1./3.] * 3
    elif number_states == 5:
        pi = [0.2] * 5
    elif number_states == 10:
        pi = [0.1] * 10
    elif number_states == 15:
        pi = [1./15.] * 15
    elif number_states == 20:
        pi = [0.05] * 20

    # generate FF, TF models from parameters
    model_ff = ghmm.HMMFromMatrices(F,ghmm.GaussianDistribution(F), A, B_ff, pi) # Will be Trained
    model_tf = ghmm.HMMFromMatrices(F,ghmm.GaussianDistribution(F), A, B_tf, pi) # Will be Trained

    total_seq = Fmat
    for i in range((Foliage_Trials + Trunk_Trials)):
        total_seq[i][:] = sum(total_seq[i][:],[])

    total_seq_ff = total_seq[0:Foliage_Trials]
    total_seq_tf = total_seq[Foliage_Trials:Foliage_Trials + Trunk_Trials]
    
    #print len(total_seq_ff)
    #print len(total_seq_tf)
    
    print "Training the HMM Models..."
       
    train_seq_ff = total_seq_ff
    train_seq_tf = total_seq_tf
       
    final_ts_ff = ghmm.SequenceSet(F,train_seq_ff)
    final_ts_tf = ghmm.SequenceSet(F,train_seq_tf)    

    model_ff.baumWelch(final_ts_ff)
    model_tf.baumWelch(final_ts_tf)

    print "Models Trained: Ready to Collect Data !"

    # Gather Data from Robot Online
    index = 0
    idx = 0
    max_force = 0
    fmags = np.zeros(384)
    global_contact_vector = [0,0,0]
    time_varying_data = [0]
    FLAG_Trunk = False
    FLAG_Foliage = False
    FLAG_Unknown = True
    getdata()
    
    
        

    
    
   

    
    
