
# 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

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

# Define features

def feature_vector_diff_2_objs(Zt1,Zt2,i): 

    data_matrix = np.array([0,0,0])
    n = i+121
    while (i < n):
        data_instant = np.array([Zt1[i,3],Zt1[i,4],Zt2[i,1]])
        data_matrix = np.row_stack([data_matrix, data_instant])
        i = i+1
    Fvec_a = np.matrix(data_matrix[1:,0]).T
    # With Scaling
    max_a = np.max(abs(Fvec_a))
    min_a = np.min(abs(Fvec_a))
    mean_a = np.mean(Fvec_a)
    std_a = np.std(Fvec_a)
    #Fvec_a = (Fvec_a)/max_a
    #Fvec_a = (Fvec_a-mean_a)
    #Fvec_a = (Fvec_a-mean_a)/max_a
    #Fvec_a = (Fvec_a-mean_a)/std_a
    
    Fvec_b = np.matrix(data_matrix[1:,1]).T
    # With Scaling
    max_b = np.max(abs(Fvec_b))
    min_b = np.min(abs(Fvec_b))
    mean_b = np.mean(Fvec_b)
    std_b = np.std(Fvec_b)
    #Fvec_b = (Fvec_b)/max_b
    #Fvec_b = (Fvec_b-mean_b)
    #Fvec_b = (Fvec_b-mean_b)/max_b
    #Fvec_b = (Fvec_b-mean_b)/std_b
    
    Fvec_c = np.matrix(data_matrix[1:,2]).T
    # With Scaling
    max_c = np.max(abs(Fvec_c))
    min_c = np.min(abs(Fvec_c))
    mean_c = np.mean(Fvec_c)
    std_c = np.std(Fvec_c)
    #Fvec_c = (Fvec_c)/max_c
    #Fvec_c = (Fvec_c-mean_c)
    #Fvec_c = (Fvec_c-mean_c)/max_c
    #Fvec_c = (Fvec_c-mean_c)/std_c
    #Fvec_c = Fvec_c*np.max((max_a,max_b))/max_c

    Fvec = np.row_stack([Fvec_a,Fvec_b,Fvec_c])
    n_Fvec, m_Fvec = np.shape(Fvec) 
    #print 'Feature_Vector_Shape:',n_Fvec, m_Fvec
    return Fvec

# Returns mu,sigma for 20 hidden-states from feature-vectors(123,35) for RF,SF,RM,SM models
  
def feature_to_mu_cov(fvec1,fvec2): 
    
    index = 0
    m,n = np.shape(fvec1)
    #print m,n
    mu_1 = np.zeros((10,1))
    mu_2 = np.zeros((10,1))
    cov = np.zeros((10,2,2))
    DIVS = m/10

    while (index < 10):
        m_init = index*DIVS
        temp_fvec1 = fvec1[(m_init):(m_init+DIVS),0:]
        temp_fvec2 = fvec2[(m_init):(m_init+DIVS),0:]
        temp_fvec1 = np.reshape(temp_fvec1,DIVS*n)
        temp_fvec2 = np.reshape(temp_fvec2,DIVS*n)
        mu_1[index] = np.mean(temp_fvec1)
        mu_2[index] = np.mean(temp_fvec2)
        cov[index,:,:] = np.cov(np.concatenate((temp_fvec1,temp_fvec2),axis=0))
        if index == 0:
            print 'mean = ', mu_2[index]
            print 'mean = ', scp.mean(fvec2[(m_init):(m_init+DIVS),0:])
            print np.shape(np.concatenate((temp_fvec1,temp_fvec2),axis=0))
            print cov[index,:,:]
            print scp.std(fvec2[(m_init):(m_init+DIVS),0:])
            print scp.std(temp_fvec2)
        index = index+1
        
    return mu_1,mu_2,cov


if __name__ == '__main__':

    Fmat = Fmat_original

    # First_Object

    ta_no_fo_t1 = ut.load_pickle('/home/tapo/svn/robot1_data/usr/tapo/data/New_Experiments_Qi/Two_Objects/black_foam_blue_cup/reverse/new4/time_varying_data_black_foam_blue_cup_fixed_first_object_trial_5.pkl')
    fa_no_fo_t1 = ut.load_pickle('/home/tapo/svn/robot1_data/usr/tapo/data/New_Experiments_Qi/Two_Objects/black_foam_blue_cup/reverse/new4/time_varying_tracking_data_black_foam_blue_cup_fixed_first_object_trial_5.pkl')

    # Second_Object

    ta_no_so_t1 = ut.load_pickle('/home/tapo/svn/robot1_data/usr/tapo/data/New_Experiments_Qi/Two_Objects/black_foam_blue_cup/reverse/new4/time_varying_data_black_foam_blue_cup_fixed_second_object_trial_5.pkl')
    fa_no_so_t1 = ut.load_pickle('/home/tapo/svn/robot1_data/usr/tapo/data/New_Experiments_Qi/Two_Objects/black_foam_blue_cup/reverse/new4/time_varying_tracking_data_black_foam_blue_cup_fixed_second_object_trial_5.pkl')
    
    # Creating Feature Vector
    Fmat2 = np.matrix(np.zeros((363,2)))
    Fmat2[:,0] = feature_vector_diff_2_objs(ta_no_fo_t1,fa_no_fo_t1,0)
    Fmat2[:,1] = feature_vector_diff_2_objs(ta_no_so_t1,fa_no_so_t1,0)

# Checking the Data-Matrix

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

    mu_rf_force,mu_rf_motion,cov_rf = feature_to_mu_cov(Fmat[0:121,0:35],Fmat[242:363,0:35])
    mu_rm_force,mu_rm_motion,cov_rm = feature_to_mu_cov(Fmat[0:121,35:70],Fmat[242:363,35:70])
    mu_sf_force,mu_sf_motion,cov_sf = feature_to_mu_cov(Fmat[0:121,70:105],Fmat[242:363,70:105])
    mu_sm_force,mu_sm_motion,cov_sm = feature_to_mu_cov(Fmat[0:121,105:140],Fmat[242:363,105:140])

# HMM - Implementation:

# 20 Hidden States
# Max. Force, Contact Area(Not now), and Contact Motion 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 = [0.0]*10
    B_rm = [0.0]*10
    B_sf = [0.0]*10
    B_sm = [0.0]*10
      
    for num_states in range(10):
        B_rf[num_states] = [[mu_rf_force[num_states][0],mu_rf_motion[num_states][0]],[cov_rf[num_states][0][0],cov_rf[num_states][0][1],cov_rf[num_states][1][0],cov_rf[num_states][1][1]]]
        B_rm[num_states] = [[mu_rm_force[num_states][0],mu_rm_motion[num_states][0]],[cov_rm[num_states][0][0],cov_rm[num_states][0][1],cov_rm[num_states][1][0],cov_rm[num_states][1][1]]]        
        B_sf[num_states] = [[mu_sf_force[num_states][0],mu_sf_motion[num_states][0]],[cov_sf[num_states][0][0],cov_sf[num_states][0][1],cov_sf[num_states][1][0],cov_sf[num_states][1][1]]]
        B_sm[num_states] = [[mu_sm_force[num_states][0],mu_sm_motion[num_states][0]],[cov_sm[num_states][0][0],cov_sm[num_states][0][1],cov_sm[num_states][1][0],cov_sm[num_states][1][1]]]
    
    #print B_sm
    #print mu_sm_motion
    
    # pi - initial probabilities per state
    pi = [0.1] * 10

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

    # For Training
   
    total_seq = np.matrix(np.concatenate((np.array(Fmat[0:121,:]),np.array(Fmat[242:363,:])),axis=0))
    total_seq2 = np.matrix(np.concatenate((np.array(Fmat2[0:121,:]),np.array(Fmat2[242:363,:])),axis=0))
    m_total, n_total = np.shape(total_seq)
    #print 'Total_Sequence_Shape:', m_total, n_total
    
    train_seq_rf = (np.array(total_seq[:,0:35]).T).tolist()
    train_seq_rm = (np.array(total_seq[:,35:70]).T).tolist()
    train_seq_sf = (np.array(total_seq[:,70:105]).T).tolist()
    train_seq_sm = (np.array(total_seq[:,105:140]).T).tolist()
    #print train_seq_rf
    
    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)

    #print " "
    #print "TRAINED RIGID-FIXED MODEL"
    #print model_rf
    #print " "
    #print " "
    #print "TRAINED RIGID-MOVABLE MODEL"
    #print model_rm
    #print " "
    #print " "
    #print "TRAINED SOFT-FIXED MODEL"
    #print model_sf
    #print " "
    #print " "
    #print "TRAINED SOFT-MOVABLE MODEL"
    #print model_sm
    #print " "
    #print " "

    # Test New Objects
    
    test_seq_obj1 = (np.array(total_seq2[:,0]).T).tolist()
    #print test_seq_obj1
    test_seq_obj2 = (np.array(total_seq2[:,1]).T).tolist()

    new_test_seq_obj1 = np.array(sum(test_seq_obj1,[]))
    #print new_test_seq_obj1
    new_test_seq_obj2 = np.array(sum(test_seq_obj2,[]))

    ts_obj1 = new_test_seq_obj1
    ts_obj2 = new_test_seq_obj2
    
    final_ts_obj1 = ghmm.EmissionSequence(F,ts_obj1.tolist())
    final_ts_obj2 = ghmm.EmissionSequence(F,ts_obj2.tolist())
    #print final_ts_obj2

    # Find Viterbi Path
    
    path_rf_obj1 = model_rf.viterbi(final_ts_obj1)
    path_rm_obj1 = model_rm.viterbi(final_ts_obj1)
    path_sf_obj1 = model_sf.viterbi(final_ts_obj1)
    path_sm_obj1 = model_sm.viterbi(final_ts_obj1)

    path_rf_obj2 = model_rf.viterbi(final_ts_obj2)
    path_rm_obj2 = model_rm.viterbi(final_ts_obj2)
    path_sf_obj2 = model_sf.viterbi(final_ts_obj2)
    path_sm_obj2 = model_sm.viterbi(final_ts_obj2)

    obj1 = max(path_rf_obj1[1],path_rm_obj1[1],path_sf_obj1[1],path_sm_obj1[1])
    obj2 = max(path_rf_obj2[1],path_rm_obj2[1],path_sf_obj2[1],path_sm_obj2[1])
    print " "
    if obj1 == path_rf_obj1[1]:
        print "ONE OBJECT IS RIGID-FIXED"
    elif obj1 == path_rm_obj1[1]:
        print "ONE OBJECT IS RIGID-MOVABLE"
    elif obj1 == path_sf_obj1[1]:
        print "ONE OBJECT IS SOFT-FIXED"
    else:
        print "ONE OBJECT IS SOFT-MOVABLE"
    print " "
    if obj2 == path_rf_obj2[1]:
        print "THE OTHER OBJECT IS RIGID-FIXED"
    elif obj2 == path_rm_obj2[1]:
        print "THE OTHER OBJECT IS RIGID-MOVABLE"
    elif obj2 == path_sf_obj2[1]:
        print "THE OTHER OBJECT IS SOFT-FIXED"
    else:
        print "THE OTHER OBJECT IS SOFT-MOVABLE"
    print " "

    
    

    







 
