# 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
import time

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/taxel_based/')

from data_variable_length_force_sample import Fmat_original, temp_num_fol, temp_num_trunk

def callback(data, callback_args):
    #rospy.loginfo('Getting data!')
    global start_time
    start_time = time.time()
    # 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
    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)):
        if fmags_instant_tuned[i] > 0.0:
            fmags[i].append(fmags_instant_tuned[i])
        else:
            fmags[i] = []

    # Calculating no. of contact regions with hand-tuned force threshold
    global total_contact
    total_contact = sum(i > 0 for i in fmags_instant_tuned)

    # Gathering Contact Data for Haptic Mapping
    global global_contact_vector
    for i in range(len(fmags_instant_tuned)):
        global_contact_vector[i] = r1*((np.column_stack([data.centers_x[i], data.centers_y[i], data.centers_z[i]])).T) + t1
        
    test_data()
    global taxel_FLAG
    for i in range(len(fmags_instant_tuned)):
        if taxel_FLAG[i] > -1:
            idx = taxel_FLAG[i]
            contact_info = global_contact_vector[i]
            pubdata(idx, contact_info)

def test_data():
    global exp_time
    # Do Stuff For Testing which basically returns which FLAG is true
    global taxel_FLAG # -1 for not in Contact, 0 for Unknown (Red), 1 for Foliage (green), 2 for Trunk (brown)
    global trunk_contact
    trunk_contact = 0
    global foliage_contact
    foliage_contact = 0
    global total_contact
    global other_contact
    other_contact = 0
    num_samples = []
    # For Testing
    global start_time
    global fmags
    for i in range(384):
        if (len(fmags[i]) > 0):      
            ts_obj = fmags[i]
            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 > 5)):
            #if ((obj == path_ff_obj[1])):
                #print 'Taxel', i, 'is Foliage !'
                taxel_FLAG[i] = 1
                foliage_contact = foliage_contact+1
                num_samples.append(len(fmags[i]))
            elif ((obj == path_tf_obj[1]) and (diff > 20)):
            #elif ((obj == path_tf_obj[1])):
                #print 'Taxel', i, 'is Trunk !'
                taxel_FLAG[i] = 2
                trunk_contact = trunk_contact+1
                num_samples.append(len(fmags[i]))
            #elif ((obj == path_tf_obj[1]):
                #print 'Taxel', i, 'is Uncertain'
                #taxel_FLAG[i] = 0
            else:
                taxel_FLAG[i] = 0
                #print 'Taxel', i, 'is Unknown'
                other_contact = other_contact+1
                                
        else:
            #print 'Taxel', i, 'is not in Contact'
            taxel_FLAG[i] = -1
   
    time_taken = time.time()-start_time
    global quantitative_data
    exp_time = exp_time+0.01
    if len(num_samples) > 0:
        quant_instant_data = [exp_time, time_taken, min(num_samples), total_contact, trunk_contact, foliage_contact, other_contact]
    else:
        quant_instant_data = [exp_time, time_taken, 0, total_contact, trunk_contact, foliage_contact, other_contact]
    quantitative_data = np.row_stack([quantitative_data, quant_instant_data])
            

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, callback, callback_args = (tf_lstnr))
    rospy.spin()

def pubdata(idx, contact_info):
    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 = 0.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 = contact_info[0]
    marker.pose.position.y = contact_info[1]
    marker.pose.position.z = contact_info[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
    taxel_FLAG = {}
    for i in range(384):
        taxel_FLAG[i] = -1 # -1 for not in Contact, 0 for Unknown (Red), 1 for Foliage (green), 2 for Trunk (brown) 
    fmags = {}
    for i in range(384):
        fmags[i] = []
    global_contact_vector = {}
    for i in range(384):
        global_contact_vector[i] = []
    FLAG_Trunk = False
    FLAG_Foliage = False
    FLAG_Unknown = True

    total_contact = 0
    trunk_contact = 0
    foliage_contact = 0
    other_contact = 0
    start_time = 0.0
    exp_time = 0.0
    quantitative_data = [0,0,0,0,0,0,0]
    getdata()
    ut.save_pickle(quantitative_data, '/home/tapo/svn/robot1_data/usr/tapo/data/rapid_categorization/Taxel_Based/Tests/foliage_reach_6.pkl')
    #ut.save_pickle(quantitative_data, '/home/tapo/svn/robot1_data/usr/tapo/data/rapid_categorization/Taxel_Based/Tests/trunk_reach_2.pkl')
    
    
        

    
    
   

    
    
