# Hidden Markov Model Implementation

import math
import pylab as pyl
import numpy as np
import matplotlib.pyplot as pp
import matplotlib.cm as cm
#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 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

import ghmm

from visualization_msgs.msg import Marker
from visualization_msgs.msg import MarkerArray
import hrl_haptic_manipulation_in_clutter_msgs.msg as haptic_msgs
from hrl_haptic_manipulation_in_clutter_msgs.msg import SkinContact
from hrl_haptic_manipulation_in_clutter_msgs.msg import ObjectInfo
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

import std_msgs.msg

import sys
sys.path.insert(0, '/home/tapo/svn/robot1_data/usr/tapo/data_code/Classification/Data/Single_Contact_HMM/384')

from data_384 import Fmat_original

# Returns mu,sigma for 10 hidden-states from feature-vectors(123,35) for RF,SF,RM,SM models
  
def feature_to_mu_sigma(fvec): 
    
    index = 0
    m,n = np.shape(fvec)
    #print m,n
    mu = np.matrix(np.zeros((10,1)))
    sigma = np.matrix(np.zeros((10,1)))
    DIVS = m/10

    while (index < 10):
        m_init = index*DIVS
        temp_fvec = fvec[(m_init):(m_init+DIVS),0:]
        #if index == 1:
            #print temp_fvec
        mu[index] = scp.mean(temp_fvec)
        sigma[index] = scp.std(temp_fvec)
        index = index+1
        
    return mu,sigma

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

    # New Code for New Haptic State Message:
    global_values_x = []
    global_values_y = []
    global_values_z = [] 
    global_centers_x = []
    global_centers_y = []
    global_centers_z = []

    length = len(data.skins)
    for i in range(length):
        for element in data.skins[i].values_x:
	    global_values_x.append(element)
        for element in data.skins[i].values_y:
	    global_values_y.append(element)
        for element in data.skins[i].values_z:
	    global_values_z.append(element)

        for element in data.skins[i].centers_x:
	    global_centers_x.append(element)
        for element in data.skins[i].centers_y:
	    global_centers_y.append(element)
        for element in data.skins[i].centers_z:
	    global_centers_z.append(element)	
    #print global_values_x
        
    for i in range(np.size(global_values_x), 32):
        global_values_x.append(0.)
    for i in range(np.size(global_values_y), 32):
        global_values_y.append(0.)
    for i in range(np.size(global_values_z), 32):
        global_values_z.append(0.)

    for i in range(np.size(global_centers_x), 32):
        global_centers_x.append(0.)
    for i in range(np.size(global_centers_y), 32):
        global_centers_y.append(0.)
    for i in range(np.size(global_centers_z), 32):
        global_centers_z.append(0.)
    
        
    # Processing Force Data
    force_vectors = np.row_stack([global_values_x, global_values_y, global_values_z])
    fmags_instant = ut.norm(force_vectors)
    #print fmags_instant
    threshold = 0.01
    fmags_tuned = fmags_instant - threshold
    fmags_tuned[np.where(fmags_tuned<0)]=0
    #print fmags_tuned
    fmags_instant_tuned = fmags_tuned
    global fmags
    fmags = np.row_stack([fmags,fmags_instant_tuned])

    # Processing Contact Data
    contact_vectors = np.row_stack([global_centers_x, global_centers_y, global_centers_z])
    #print np.shape(contact_vectors)
    contact_regions = fmags_instant > threshold
    lb,ls = ni.label(contact_regions)
    total_contact = ni.sum(lb) # After thresholding, assuming one connected component (experiment designed that way)
    #print total_contact
    # Calculating force data for contact with hand-tuned force threshold
    #total_forces = ni.sum(fmags_instant,lb)
    #mean_forces = ni.mean(fmags_instant,lb)
    temp = fmags_instant*lb
    global max_force
    max_force = np.max(temp)

    global time_varying_data
    time_varying_data.append(max_force)

    if total_contact > 0:
            
        global contact_point_world
        global_x = ni.mean(contact_vectors[0,:],lb) 
        global_y = ni.mean(contact_vectors[1,:],lb) 
        global_z = ni.mean(contact_vectors[2,:],lb)
        print global_x, global_y, global_z 
        contact_point_world = [global_x,global_y,global_z]
        #print contact_point_world
        #print np.shape(contact_point_world)
        global time_instant_contact_data
        time_instant_contact_data = [contact_point_world[0],contact_point_world[1],contact_point_world[2]]
        #print time_instant_contact_data
        global time_varying_contact_data
        time_varying_contact_data = np.row_stack([time_varying_contact_data, time_instant_contact_data])

    test_data()

def tracking_point():
    rospy.loginfo('Tracking Distance!')
    ta = time_varying_contact_data
    k = 0
    for i in ta[:,0]:
        if i != ta[-1,0]:
            instant_dist = math.sqrt((ta[k+1,1]-ta[1,1])**2 + (ta[k+1,2]-ta[1,2])**2 + (ta[k+1,3]-ta[1,3])**2)
            time_instant_tracker = [ta[k+1,0], instant_dist]
            global time_varying_tracker
            time_varying_tracker = np.row_stack([time_varying_tracker, time_instant_tracker])
            k=k+1

def test_data():
    
    pub = rospy.Publisher("/hmm/object_data", ObjectInfo)
    obj_info = ObjectInfo()
    obj_info.obj_id = 0
    obj_info.contact_x = 0
    obj_info.contact_y = 0
    obj_info.contact_z = 0 

    global index   

    global FLAG_Trunk
    global FLAG_Trunk_List
    global FLAG_Leaf
    global FLAG_Leaf_List
    
    global time_instant_contact_data
    
    # 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_rf
        global model_rm
        global model_sf
        global model_sm

        path_rf_obj = model_rf.viterbi(final_ts_obj)
        path_rm_obj = model_rm.viterbi(final_ts_obj)
        path_sf_obj = model_sf.viterbi(final_ts_obj)
        path_sm_obj = model_sm.viterbi(final_ts_obj)
        #print path_rf_obj[1], path_rm_obj[1], path_sf_obj[1], path_sm_obj[1]

        diff_rf = min(abs(path_rf_obj[1]-path_rm_obj[1]),abs(path_rf_obj[1]-path_sf_obj[1]),abs(path_rf_obj[1]-path_sm_obj[1]))
        diff_rm = min(abs(path_rm_obj[1]-path_rf_obj[1]),abs(path_rm_obj[1]-path_sf_obj[1]),abs(path_rm_obj[1]-path_sm_obj[1]))
        diff_sf = min(abs(path_sf_obj[1]-path_rf_obj[1]),abs(path_sf_obj[1]-path_rm_obj[1]),abs(path_sf_obj[1]-path_sm_obj[1]))
        diff_sm = min(abs(path_sm_obj[1]-path_rf_obj[1]),abs(path_sm_obj[1]-path_rm_obj[1]),abs(path_sm_obj[1]-path_sf_obj[1]))
        obj = max(path_rf_obj[1],path_rm_obj[1],path_sf_obj[1],path_sm_obj[1])
        obj_min = min(abs(path_rf_obj[1]),abs(path_rm_obj[1]),abs(path_sf_obj[1]),abs(path_sm_obj[1]))
        #print ' '
        #print '################################################################################'
        if ((obj == path_rf_obj[1]) and (diff_rf > 500)):
        #if ((obj == path_rf_obj[1]) and (obj_min > 1800)):
            print 'Rigid_Fixed :'
            FLAG_Trunk = True
            FLAG_Leaf = False
            obj_info.obj_id = 1
            obj_info.contact_x = time_instant_contact_data[0]
            obj_info.contact_y = time_instant_contact_data[1]
            obj_info.contact_z = time_instant_contact_data[2]
            pub.publish(obj_info)
        elif ((obj == path_rm_obj[1]) and (diff_rm > 500)):
        #elif ((obj == path_rm_obj[1]) and (obj_min > 1800)):
            print 'Rigid_Movable :'
            FLAG_Trunk = False
            FLAG_Leaf = False
            obj_info.obj_id = 2
            obj_info.contact_x = time_instant_contact_data[0]
            obj_info.contact_y = time_instant_contact_data[1]
            obj_info.contact_z = time_instant_contact_data[2]
            pub.publish(obj_info)
        elif ((obj == path_sf_obj[1]) and (diff_sf > 500)):
        #elif ((obj == path_sf_obj[1]) and (obj_min > 1800)):
            print 'Soft_Fixed :'
            FLAG_Trunk = False
            FLAG_Leaf = True
            obj_info.obj_id = 3
            obj_info.contact_x = time_instant_contact_data[0]
            obj_info.contact_y = time_instant_contact_data[1]
            obj_info.contact_z = time_instant_contact_data[2]
            pub.publish(obj_info)
        elif ((obj == path_sm_obj[1]) and (diff_sm > 500)):
        #elif ((obj == path_sm_obj[1]) and (obj_min > 1800)):
            print 'Soft_Movable :'
            FLAG_Trunk = False
            FLAG_Leaf = False
            obj_info.obj_id = 4
            obj_info.contact_x = time_instant_contact_data[0]
            obj_info.contact_y = time_instant_contact_data[1]
            obj_info.contact_z = time_instant_contact_data[2]
            pub.publish(obj_info)
        else:
            print 'Unknown'
            FLAG_Trunk = False
            FLAG_Leaf = False
            object_ID = 0
            pub.publish(obj_info)

        FLAG_Trunk_List.append(FLAG_Trunk)
        FLAG_Leaf_List.append(FLAG_Leaf)
        #print '################################################################################'
        #print ' '
    else:
        print 'Unknown'
        FLAG_Trunk = False
        FLAG_Leaf = False
        object_ID = 0
        pub.publish(obj_info)
        time_varying_data = [0]    
       
        
def getdata():
    rospy.init_node('getdata', anonymous=True)
    rospy.Subscriber("/haptic_mpc/robot_state", haptic_msgs.RobotHapticState, callback)
    rospy.spin()

def getpics():
    
    global FLAG_Trunk_List
    global FLAG_Leaf_List

    j = 0
    #index = 0 # For Fixed Pillow
    index = 16 #For Fixed Styrofoam, Fixed Foliage, Movable Styrofoam, and Movable Pillow
    for i in fmags:
        force_arr_raw = fmags[j,:]
        force_arr_raw_reshaped = force_arr_raw.reshape((16,24))
        force_arr = np.column_stack((force_arr_raw_reshaped[:,index:],force_arr_raw_reshaped[:,:index]))
        global frame
        frame = frame + 1
        if ((not FLAG_Trunk_List[j]) and (not FLAG_Leaf_List[j])):
            #Keep it Black
            pp.imshow(force_arr, interpolation='nearest', cmap=cm.binary,
              origin='upper', vmin=0, vmax=1)
        elif ((FLAG_Trunk_List[j]) and (not FLAG_Leaf_List[j])):
            #Make it Red
            pp.imshow(force_arr, interpolation='nearest', cmap=cm.Reds,
              origin='upper', vmin=0, vmax=1)
        elif ((not FLAG_Trunk_List[j]) and (FLAG_Leaf_List[j])):
            #Make it Green
            pp.imshow(force_arr, interpolation='nearest', cmap=cm.Greens,
              origin='upper', vmin=0, vmax=1)
        pp.title('Unrolled Taxel Array')
        pp.xlabel('Along the circumference')
        pp.ylabel('Along the forearm')
        pp.xlim(-0.5,23.5)
        pp.ylim(15.5, -0.5)
        pyl.savefig('%05d.%s'% (frame, format))
        j = j+1


if __name__ == '__main__':

    Fmat = Fmat_original

# Checking the Data-Matrix

    m_tot, n_tot = np.shape(Fmat)
    #print " "
    #print 'Total_Matrix_Shape:',m_tot,n_tot

    mu_rf,sigma_rf = feature_to_mu_sigma(Fmat[0:121,0:35])
    mu_rm,sigma_rm = feature_to_mu_sigma(Fmat[0:121,35:70])
    mu_sf,sigma_sf = feature_to_mu_sigma(Fmat[0:121,70:105])
    mu_sm,sigma_sm = feature_to_mu_sigma(Fmat[0:121,105:140])
    #print [mu_rf, sigma_rf]

# 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
    
    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]]

    # B - Emission Matrix, parameters of emission distributions in pairs of (mu, sigma)
    
    B_rf = np.zeros((10,2))
    B_rm = np.zeros((10,2))
    B_sf = np.zeros((10,2))
    B_sm = np.zeros((10,2))
      
    for num_states in range(10):
        B_rf[num_states,0] = mu_rf[num_states]
        B_rf[num_states,1] = sigma_rf[num_states]
        B_rm[num_states,0] = mu_rm[num_states]
        B_rm[num_states,1] = sigma_rm[num_states]
        B_sf[num_states,0] = mu_sf[num_states]
        B_sf[num_states,1] = sigma_sf[num_states]
        B_sm[num_states,0] = mu_sm[num_states]
        B_sm[num_states,1] = sigma_sm[num_states]
    B_rf = B_rf.tolist()
    B_rm = B_rm.tolist()
    B_sf = B_sf.tolist()
    B_sm = B_sm.tolist()
    
    # pi - initial probabilities per state
    pi = [0.1] * 10

    # generate RF, RM, SF, SM models from parameters
    model_rf = ghmm.HMMFromMatrices(F,ghmm.GaussianDistribution(F), A, B_rf, pi) # Will be Trained
    model_rm = ghmm.HMMFromMatrices(F,ghmm.GaussianDistribution(F), A, B_rm, pi) # Will be Trained
    model_sf = ghmm.HMMFromMatrices(F,ghmm.GaussianDistribution(F), A, B_sf, pi) # Will be Trained
    model_sm = ghmm.HMMFromMatrices(F,ghmm.GaussianDistribution(F), A, B_sm, pi) # Will be Trained

    total_seq = Fmat[0:121,:]

    # For Training

    total_seq_rf = total_seq[0:121,1:35]
    total_seq_rm = total_seq[0:121,36:70]
    total_seq_sf = total_seq[0:121,71:105]
    total_seq_sm = total_seq[0:121,106:140]
   
    train_seq_rf = (np.array(total_seq_rf).T).tolist()
    train_seq_rm = (np.array(total_seq_rm).T).tolist()
    train_seq_sf = (np.array(total_seq_sf).T).tolist()
    train_seq_sm = (np.array(total_seq_sm).T).tolist()
    
    final_ts_rf = ghmm.SequenceSet(F,train_seq_rf)
    final_ts_rm = ghmm.SequenceSet(F,train_seq_rm)
    final_ts_sf = ghmm.SequenceSet(F,train_seq_sf)
    final_ts_sm = ghmm.SequenceSet(F,train_seq_sm)

    model_rf.baumWelch(final_ts_rf)
    model_rm.baumWelch(final_ts_rm)
    model_sf.baumWelch(final_ts_sf)
    model_sm.baumWelch(final_ts_sm)

    # Gather Data from Robot Online
    index = 0
    frame = 0
    max_force = 0
    format = 'png'
    fmags = np.zeros(32)
    contact_point_world = [0,0,0]
    time_varying_data = [0]
    time_instant_contact_data = [0,0,0]
    time_varying_contact_data = [0,0,0]
    time_varying_tracker = [0,0]
    FLAG_Trunk = False
    FLAG_Leaf = False
    FLAG_Trunk_List = [False]
    FLAG_Leaf_List = [False]
    getdata()
    #getpics()
    #tracking_point()
    
    
    
    
        

    
    
   

    
    
