# Hidden Markov Model 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 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 unittest
import ghmm
import ghmmwrapper
import random

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

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

from data_variable_length_force_sample import Fmat_original

def compute_contact_regions(arr_2d, threshold):
    mask = arr_2d > threshold
    label_im, nb_labels = ni.label(mask)
    return label_im, nb_labels

def compute_obj_labels(label_im, nb_labels):
    i=1
    local_nb_label_first_obj = -1
    local_nb_label_second_obj = -1
    index = np.zeros(nb_labels)
    row,col = np.shape(label_im)
    while i <= nb_labels:
        j=0
        while j < row:
            k=0
            while k < col:
                if label_im[j][k] == i:
                    index[i-1] = index[i-1]+1
                k=k+1
            j=j+1
        i=i+1
    temp=0
    max_index_1 = np.max(index)
    #print max_index_1
    while temp < len(index):
        if index[temp] == max_index_1 and max_index_1 > 0:
            local_nb_label_first_obj = temp+1
        temp = temp+1
    index[local_nb_label_first_obj-1]=-1
    max_index_2 = np.max(index)
    #print max_index_2
    temp=0
    while temp < len(index):
        if index[temp] == max_index_2 and max_index_2 > 0:
            local_nb_label_second_obj = temp+1
        temp = temp+1
    return local_nb_label_first_obj, local_nb_label_second_obj, max_index_1, max_index_2

def compute_resultant_force_magnitudes(force_arr, label_im, nb_label):
    total_force = ni.sum(force_arr, label_im, nb_label)
    return total_force

def compute_max_force(force_arr, label_im, nb_label):
    max_force = ni.maximum(force_arr, label_im, nb_label)
    return max_force

def compute_center_of_pressure(cx_arr, cy_arr, cz_arr, label_im,
                                nb_label):
    cx = ni.mean(cx_arr, label_im, nb_label)
    cy = ni.mean(cy_arr, label_im, nb_label)
    cz = ni.mean(cz_arr, label_im, nb_label)
    contact_vector = np.column_stack([cx, cy, cz])
    return contact_vector

def track_object_connected_component(cx_arr, cy_arr, cz_arr, r1, t1, label_im, nb_label_first_obj1, nb_label_second_obj1, total_contact_first_obj1, total_contact_second_obj1):
    
    global cop_global_first_obj_prev
    global cop_global_second_obj_prev
    global iterindex
    cop_local_first_obj1 = compute_center_of_pressure(cx_arr,cy_arr,cz_arr,label_im,nb_label_first_obj1)
    cop_global_first_obj1 = r1*(cop_local_first_obj1.T) + t1
    cop_local_second_obj1 = compute_center_of_pressure(cx_arr,cy_arr,cz_arr,label_im,nb_label_second_obj1)
    cop_global_second_obj1 = r1*(cop_local_second_obj1.T) + t1
    #print lin.norm(cop_global_first_obj1-cop_global_first_obj_prev)
    if (lin.norm(cop_global_first_obj1-cop_global_first_obj_prev) >= lin.norm(cop_global_first_obj1 - cop_global_second_obj_prev)) and (iterindex > 0):
    #if (lin.norm(cop_global_second_obj1-cop_global_second_obj_prev) >= lin.norm(cop_global_second_obj1 - cop_global_first_obj_prev)) and (iterindex > 0):
    #if ((lin.norm(cop_global_first_obj1-cop_global_first_obj_prev) >= lin.norm(cop_global_first_obj1 - cop_global_second_obj_prev)) or (lin.norm(cop_global_second_obj1-cop_global_second_obj_prev) >= lin.norm(cop_global_second_obj1 - cop_global_first_obj_prev))) and (iterindex > 0):
    #if (lin.norm(cop_global_first_obj1-cop_global_first_obj_prev) > 0.005) and (iterindex > 0):
        #print "Need to be Exchanged: ", total_contact_first_obj1, " compared to ", total_contact_second_obj1
        #print "Need to be Exchanged: ", nb_label_first_obj1, " compared to ", nb_label_second_obj1
        temp1 = nb_label_first_obj1
        temp2 = total_contact_first_obj1
        nb_label_first_obj1 = nb_label_second_obj1
        total_contact_first_obj1 = total_contact_second_obj1
        nb_label_second_obj1 = temp1
        total_contact_second_obj1 = temp2
        cop_global_first_obj_prev = cop_global_second_obj1
        cop_global_second_obj_prev = cop_global_first_obj1
        print "The connected component corresponding to the first object is not the largest"
    else:
        cop_global_first_obj_prev = cop_global_first_obj1
        cop_global_second_obj_prev = cop_global_second_obj1
        #print "The largest connected component is the first object"

    iterindex = iterindex + 1

    return nb_label_first_obj1, nb_label_second_obj1, total_contact_first_obj1, total_contact_second_obj1

def track_object_connected_component_single(cx_arr, cy_arr, cz_arr, r1, t1, label_im, nb_label_first_obj1, nb_label_second_obj1, total_contact_first_obj1, total_contact_second_obj1):
    
    global cop_global_first_obj_prev
    global cop_global_second_obj_prev
    global iterindex
    cop_local_first_obj1 = compute_center_of_pressure(cx_arr,cy_arr,cz_arr,label_im,nb_label_first_obj1)
    cop_global_first_obj1 = r1*(cop_local_first_obj1.T) + t1
    
    #print lin.norm(cop_global_first_obj1-cop_global_first_obj_prev)
    if (lin.norm(cop_global_first_obj1-cop_global_first_obj_prev) >= lin.norm(cop_global_first_obj1 - cop_global_second_obj_prev)) and (iterindex > 0):
    #if (lin.norm(cop_global_second_obj1-cop_global_second_obj_prev) >= lin.norm(cop_global_second_obj1 - cop_global_first_obj_prev)) and (iterindex > 0):
    #if ((lin.norm(cop_global_first_obj1-cop_global_first_obj_prev) >= lin.norm(cop_global_first_obj1 - cop_global_second_obj_prev)) or (lin.norm(cop_global_second_obj1-cop_global_second_obj_prev) >= lin.norm(cop_global_second_obj1 - cop_global_first_obj_prev))) and (iterindex > 0):
    #if (lin.norm(cop_global_first_obj1-cop_global_first_obj_prev) > 0.005) and (iterindex > 0):
        #print "Need to be Exchanged: ", total_contact_first_obj1, " compared to ", total_contact_second_obj1
        #print "Need to be Exchanged: ", nb_label_first_obj1, " compared to ", nb_label_second_obj1
        temp1 = nb_label_first_obj1
        temp2 = total_contact_first_obj1
        nb_label_first_obj1 = nb_label_second_obj1
        total_contact_first_obj1 = total_contact_second_obj1
        nb_label_second_obj1 = temp1
        total_contact_second_obj1 = temp2
        cop_global_first_obj_prev = cop_global_second_obj_prev
        cop_global_second_obj_prev = cop_global_first_obj1
        print "The connected component corresponding to the first object is not the largest"
    else:
        cop_global_first_obj_prev = cop_global_first_obj1
        cop_global_second_obj_prev = cop_global_second_obj_prev
        #print "The largest connected component is the first object"

    iterindex = iterindex + 1

    return nb_label_first_obj1, nb_label_second_obj1, total_contact_first_obj1, total_contact_second_obj1

def callback(data, callback_args):
    rospy.loginfo('Getting data and Saving pics!')
    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)
    force_vectors = np.row_stack([data.values_x, data.values_y, data.values_z])
    contact_vectors = np.row_stack([data.centers_x, data.centers_y, data.centers_z]).reshape((3,16,24))
    fmags = ut.norm(force_vectors)
    force_arr = fmags.reshape((16,24))
    cx_arr = contact_vectors[0]
    cy_arr = contact_vectors[1]
    cz_arr = contact_vectors[2]

    global max_force_first_obj, max_force_second_obj
         
    label_im, nb_labels = compute_contact_regions(force_arr, 0.1)
    nb_label_first_object, nb_label_second_object, total_contact_first_object, total_contact_second_object = compute_obj_labels(label_im,nb_labels)

    if (total_contact_first_object > 1) and (total_contact_second_object > 1):

        nb_label_first_obj, nb_label_second_obj, total_contact_first_obj, total_contact_second_obj = track_object_connected_component(cx_arr, cy_arr, cz_arr, r1, t1, label_im, nb_label_first_object, nb_label_second_object, total_contact_first_object, total_contact_second_object)

        global cop_global_first_obj, cop_global_second_obj

        cop_local_first_obj = compute_center_of_pressure(cx_arr,cy_arr,cz_arr,label_im,nb_label_first_obj)
        cop_global_first_obj = r1*(cop_local_first_obj.T) + t1
        cop_local_second_obj = compute_center_of_pressure(cx_arr,cy_arr,cz_arr,label_im,nb_label_second_obj)
        cop_global_second_obj = r1*(cop_local_second_obj.T) + t1
        
        max_force_first_obj = compute_max_force(force_arr,label_im,nb_label_first_obj)
        max_force_second_obj = compute_max_force(force_arr,label_im,nb_label_second_obj)
        #if (nb_label_first_obj == nb_label_second_object) and (nb_label_second_obj == nb_label_first_object):
        #print "New Max. Forces, 1st object: ", max_force_first_obj, " compared to 2nd object: ", max_force_second_obj
        #print "Old Max. Forces, 1st object: ", max_force_first_obj_temp, " compared to 2nd object: ", max_force_second_obj_temp, "\n"

    elif (total_contact_first_object > 1) and (total_contact_second_object < 2):

        nb_label_first_obj, nb_label_second_obj, total_contact_first_obj, total_contact_second_obj = track_object_connected_component_single(cx_arr, cy_arr, cz_arr, r1, t1, label_im, nb_label_first_object, nb_label_second_object, total_contact_first_object, total_contact_second_object)
        global cop_global_first_obj, cop_global_second_obj

        cop_local_first_obj = compute_center_of_pressure(cx_arr,cy_arr,cz_arr,label_im,nb_label_first_obj)
        cop_global_first_obj = r1*(cop_local_first_obj.T) + t1
        cop_local_second_obj = compute_center_of_pressure(cx_arr,cy_arr,cz_arr,label_im,nb_label_second_obj)
        cop_global_second_obj = r1*(cop_local_second_obj.T) + t1
        
        max_force_first_obj = compute_max_force(force_arr,label_im,nb_label_first_obj)
        max_force_second_obj = compute_max_force(force_arr,label_im,nb_label_second_obj)
        #if (nb_label_first_obj == nb_label_second_object) and (nb_label_second_obj == nb_label_first_object):
        #print "New Max. Forces, 1st object: ", max_force_first_obj, " compared to 2nd object: ", max_force_second_obj
        #print "Old Max. Forces, 1st object: ", max_force_first_obj_temp, " compared to 2nd object: ", max_force_second_obj_temp, "\n"

    elif (total_contact_first_object > 1) and (total_contact_second_object < 2):

        nb_label_first_obj, nb_label_second_obj, total_contact_first_obj, total_contact_second_obj = track_object_connected_component_single(cx_arr, cy_arr, cz_arr, r1, t1, label_im, nb_label_first_object, nb_label_second_object, total_contact_first_object, total_contact_second_object)

        global cop_global_first_obj, cop_global_second_obj
        if nb_label_first_obj > 0:
            cop_local_first_obj = compute_center_of_pressure(cx_arr,cy_arr,cz_arr,label_im,nb_label_first_obj)
            cop_global_first_obj = r1*(cop_local_first_obj.T) + t1
            max_force_first_obj = compute_max_force(force_arr,label_im,nb_label_first_obj)
        if nb_label_second_obj > 0:
            cop_local_second_obj = compute_center_of_pressure(cx_arr,cy_arr,cz_arr,label_im,nb_label_second_obj)
            cop_global_second_obj = r1*(cop_local_second_obj.T) + t1
            max_force_second_obj = compute_max_force(force_arr,label_im,nb_label_second_obj)
    
    global time_varying_data_first_obj
    time_varying_data_first_obj.append(max_force_first_obj)
    global time_varying_data_second_obj
    time_varying_data_second_obj.append(max_force_second_obj)
    test_data_first_obj()
    test_data_second_obj()

def test_data_first_obj():
    
    global FLAG_Trunk
    global FLAG_Trunk_List
    global FLAG_Leaf
    global FLAG_Leaf_List

    # For Testing
    global time_varying_data_first_obj
    global max_force_first_obj
    if (max_force_first_obj > 0):      
        ts_obj = time_varying_data_first_obj
        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_ff = abs(path_ff_obj[1]-path_tf_obj[1])
        diff_tf = abs(path_tf_obj[1]-path_ff_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_ff > 500)):
        if ((obj == path_ff_obj[1]) and (obj_min > 1800)):
            print 'Foliage :'
            FLAG_Trunk = False
            FLAG_Leaf = True
        elif ((obj == path_tf_obj[1]) and (obj_min > 1800)):
            print 'Trunk :'
            FLAG_Trunk = True
            FLAG_Leaf = False
        else:
            print 'Unknown'
            FLAG_Trunk = False
            FLAG_Leaf = False
        FLAG_Trunk_List.append(FLAG_Trunk)
        FLAG_Leaf_List.append(FLAG_Leaf)
    else:
        print 'Unknown'
        FLAG_Trunk = False
        FLAG_Leaf = False
        time_varying_data_first_obj = [0]  

def test_data_second_obj():
    
    global FLAG_Trunk
    global FLAG_Trunk_List
    global FLAG_Leaf
    global FLAG_Leaf_List

    FLAG_Trunk = False
    FLAG_Leaf = False
    FLAG_Trunk_List = [False]
    FLAG_Leaf_List = [False]

    # For Testing
    global time_varying_data_second_obj
    global max_force_second_obj
    if (max_force_second_obj > 0):      
        ts_obj = time_varying_data_second_obj
        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_ff = abs(path_ff_obj[1]-path_tf_obj[1])
        diff_tf = abs(path_tf_obj[1]-path_ff_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_ff > 500)):
        if ((obj == path_ff_obj[1]) and (obj_min > 1800)):
            print 'Foliage :'
            FLAG_Trunk = False
            FLAG_Leaf = True
        elif ((obj == path_tf_obj[1]) and (obj_min > 1800)):
            print 'Trunk :'
            FLAG_Trunk = True
            FLAG_Leaf = False
        else:
            print 'Unknown'
            FLAG_Trunk = False
            FLAG_Leaf = False
        FLAG_Trunk_List.append(FLAG_Trunk)
        FLAG_Leaf_List.append(FLAG_Leaf)
    else:
        print 'Unknown'
        FLAG_Trunk = False
        FLAG_Leaf = False
        time_varying_data_second_obj = [0]     
        
def getdata():
    rospy.init_node('time_varying_data_two_objects', anonymous=True)
    tf_lstnr = tf.TransformListener()
    rospy.Subscriber("/skin_patch_forearm_right/taxels/forces", TaxelArray_Meka, callback, callback_args = (tf_lstnr))
    rospy.spin()

def getpics():
    
    global FLAG_Trunk_List
    global FLAG_Leaf_List

    ## Need to figure out a way to visualize multiple object classification ####


  
if __name__ == '__main__':

    Fmat = Fmat_original
    Foliage_Trials = 5
    Trunk_Trials = 5
    
# 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][:],[])

   # For Training
   
    total_seq_ff = total_seq[0:Foliage_Trials]
    total_seq_tf = total_seq[Foliage_Trials:(Foliage_Trials + Trunk_Trials)]
       
    train_seq_ff = (np.array(total_seq_ff).T).tolist()
    train_seq_tf = (np.array(total_seq_tf).T).tolist()
       
    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)
    
    # Gather Data from Robot Online
    index = 0
    frame = 0
    max_force_first_obj = 0
    max_force_second_obj = 0
    cop_global_first_obj = [0,0,0]
    cop_global_second_obj = [0,0,0]
    cop_global_first_obj_prev = [0,0,0]
    cop_global_second_obj_prev = [0,0,0]
    format = 'png'
    fmags = np.zeros(384)
    iterindex = 0
    time_varying_data_first_obj = [0]
    time_varying_data_second_obj = [0]
    FLAG_Trunk = False
    FLAG_Leaf = False
    FLAG_Trunk_List = [False]
    FLAG_Leaf_List = [False]
    getdata()
    getpics()  
    
    
        

    
    
   

    
    
