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

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

from data_Scaled_method_III 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_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__' or __name__ != '__main__':

    Fmat = Fmat_original
    
# Checking the Data-Matrix
    Fmat_import = Fmat
    #print np.shape(Fmat[0])
    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:

# 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 = [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 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 "----"
    #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

    trial_number = 1
    rf_final = np.matrix(np.zeros((28,1)))
    rm_final = np.matrix(np.zeros((28,1)))
    sf_final = np.matrix(np.zeros((28,1)))
    sm_final = np.matrix(np.zeros((28,1)))
    total_seq = np.matrix(np.concatenate((np.array(Fmat[0:121,:]),np.array(Fmat[242:363,:])),axis=0))
    m_total, n_total = np.shape(total_seq)
    #print 'Total_Sequence_Shape:', m_total, n_total

    while (trial_number < 6):

        # For Training
          
        if (trial_number == 1):
            j = 5
            total_seq_rf = total_seq[:,1:5]
            total_seq_rm = total_seq[:,36:40]
            total_seq_sf = total_seq[:,71:75]
            total_seq_sm = total_seq[:,106:110]
            while (j < 35):
                total_seq_rf = np.column_stack((total_seq_rf,total_seq[:,j+1:j+5]))
                total_seq_rm = np.column_stack((total_seq_rm,total_seq[:,j+36:j+40]))
                total_seq_sf = np.column_stack((total_seq_sf,total_seq[:,j+71:j+75]))
                total_seq_sm = np.column_stack((total_seq_sm,total_seq[:,j+106:j+110]))
                j = j+5

        if (trial_number == 2):
            j = 5
            total_seq_rf = np.column_stack((total_seq[:,0],total_seq[:,2:5]))
            total_seq_rm = np.column_stack((total_seq[:,35],total_seq[:,37:40]))
            total_seq_sf = np.column_stack((total_seq[:,70],total_seq[:,72:75]))
            total_seq_sm = np.column_stack((total_seq[:,105],total_seq[:,107:110]))
            while (j < 35):
                total_seq_rf = np.column_stack((total_seq_rf,total_seq[:,j+0],total_seq[:,j+2:j+5]))
                total_seq_rm = np.column_stack((total_seq_rm,total_seq[:,j+35],total_seq[:,j+37:j+40]))
                total_seq_sf = np.column_stack((total_seq_sf,total_seq[:,j+70],total_seq[:,j+72:j+75]))
                total_seq_sm = np.column_stack((total_seq_sm,total_seq[:,j+105],total_seq[:,j+107:j+110]))
                j = j+5

        if (trial_number == 3):
            j = 5
            total_seq_rf = np.column_stack((total_seq[:,0:2],total_seq[:,3:5]))
            total_seq_rm = np.column_stack((total_seq[:,35:37],total_seq[:,38:40]))
            total_seq_sf = np.column_stack((total_seq[:,70:72],total_seq[:,73:75]))
            total_seq_sm = np.column_stack((total_seq[:,105:107],total_seq[:,108:110]))
            while (j < 35):
                total_seq_rf = np.column_stack((total_seq_rf,total_seq[:,j+0:j+2],total_seq[:,j+3:j+5]))
                total_seq_rm = np.column_stack((total_seq_rm,total_seq[:,j+35:j+37],total_seq[:,j+38:j+40]))
                total_seq_sf = np.column_stack((total_seq_sf,total_seq[:,j+70:j+72],total_seq[:,j+73:j+75]))
                total_seq_sm = np.column_stack((total_seq_sm,total_seq[:,j+105:j+107],total_seq[:,j+108:j+110]))
                j = j+5

        if (trial_number == 4):
            j = 5
            total_seq_rf = np.column_stack((total_seq[:,0:3],total_seq[:,4:5]))
            total_seq_rm = np.column_stack((total_seq[:,35:38],total_seq[:,39:40]))
            total_seq_sf = np.column_stack((total_seq[:,70:73],total_seq[:,74:75]))
            total_seq_sm = np.column_stack((total_seq[:,105:108],total_seq[:,109:110]))
            while (j < 35):
                total_seq_rf = np.column_stack((total_seq_rf,total_seq[:,j+0:j+3],total_seq[:,j+4:j+5]))
                total_seq_rm = np.column_stack((total_seq_rm,total_seq[:,j+35:j+38],total_seq[:,j+39:j+40]))
                total_seq_sf = np.column_stack((total_seq_sf,total_seq[:,j+70:j+73],total_seq[:,j+74:j+75]))
                total_seq_sm = np.column_stack((total_seq_sm,total_seq[:,j+105:j+108],total_seq[:,j+109:j+110]))
                j = j+5

        if (trial_number == 5):
            j = 5
            total_seq_rf = total_seq[:,0:4]
            total_seq_rm = total_seq[:,35:39]
            total_seq_sf = total_seq[:,70:74]
            total_seq_sm = total_seq[:,105:109]
            while (j < 35):
                total_seq_rf = np.column_stack((total_seq_rf,total_seq[:,j+0:j+4]))
                total_seq_rm = np.column_stack((total_seq_rm,total_seq[:,j+35:j+39]))
                total_seq_sf = np.column_stack((total_seq_sf,total_seq[:,j+70:j+74]))
                total_seq_sm = np.column_stack((total_seq_sm,total_seq[:,j+105:j+109]))
                j = j+5
      
        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()
        #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)
    
        # For Testing
        
        if (trial_number == 1): 
            j = 5
            total_seq_rf = total_seq[:,0]
            total_seq_rm = total_seq[:,35]
            total_seq_sf = total_seq[:,70]
            total_seq_sm = total_seq[:,105]
            while (j < 35):
                total_seq_rf = np.column_stack((total_seq_rf,total_seq[:,j]))
                total_seq_rm = np.column_stack((total_seq_rm,total_seq[:,j+35]))
                total_seq_sf = np.column_stack((total_seq_sf,total_seq[:,j+70]))
                total_seq_sm = np.column_stack((total_seq_sm,total_seq[:,j+105]))
                j = j+5
        
        if (trial_number == 2): 
            j = 5
            total_seq_rf = total_seq[:,1]
            total_seq_rm = total_seq[:,36]
            total_seq_sf = total_seq[:,71]
            total_seq_sm = total_seq[:,106]
            while (j < 35):
                total_seq_rf = np.column_stack((total_seq_rf,total_seq[:,j+1]))
                total_seq_rm = np.column_stack((total_seq_rm,total_seq[:,j+36]))
                total_seq_sf = np.column_stack((total_seq_sf,total_seq[:,j+71]))
                total_seq_sm = np.column_stack((total_seq_sm,total_seq[:,j+106]))
                j = j+5 

        if (trial_number == 3): 
            j = 5
            total_seq_rf = total_seq[:,2]
            total_seq_rm = total_seq[:,37]
            total_seq_sf = total_seq[:,72]
            total_seq_sm = total_seq[:,107]
            while (j < 35):
                total_seq_rf = np.column_stack((total_seq_rf,total_seq[:,j+2]))
                total_seq_rm = np.column_stack((total_seq_rm,total_seq[:,j+37]))
                total_seq_sf = np.column_stack((total_seq_sf,total_seq[:,j+72]))
                total_seq_sm = np.column_stack((total_seq_sm,total_seq[:,j+107]))
                j = j+5 

        if (trial_number == 4): 
            j = 5
            total_seq_rf = total_seq[:,3]
            total_seq_rm = total_seq[:,38]
            total_seq_sf = total_seq[:,73]
            total_seq_sm = total_seq[:,108]
            while (j < 35):
                total_seq_rf = np.column_stack((total_seq_rf,total_seq[:,j+3]))
                total_seq_rm = np.column_stack((total_seq_rm,total_seq[:,j+38]))
                total_seq_sf = np.column_stack((total_seq_sf,total_seq[:,j+73]))
                total_seq_sm = np.column_stack((total_seq_sm,total_seq[:,j+108]))
                j = j+5

        if (trial_number == 5): 
            j = 5
            total_seq_rf = total_seq[:,4]
            total_seq_rm = total_seq[:,39]
            total_seq_sf = total_seq[:,74]
            total_seq_sm = total_seq[:,109]
            while (j < 35):
                total_seq_rf = np.column_stack((total_seq_rf,total_seq[:,j+4]))
                total_seq_rm = np.column_stack((total_seq_rm,total_seq[:,j+39]))
                total_seq_sf = np.column_stack((total_seq_sf,total_seq[:,j+74]))
                total_seq_sm = np.column_stack((total_seq_sm,total_seq[:,j+109]))
                j = j+5
    
        total_seq_obj = np.matrix(np.column_stack((total_seq_rf,total_seq_rm,total_seq_sf,total_seq_sm)))
       
        rf = np.matrix(np.zeros(np.size(total_seq_obj,1)))
        rm = np.matrix(np.zeros(np.size(total_seq_obj,1)))
        sf = np.matrix(np.zeros(np.size(total_seq_obj,1)))
        sm = np.matrix(np.zeros(np.size(total_seq_obj,1)))

        k = 0
        while (k < np.size(total_seq_obj,1)):
            test_seq_obj = (np.array(total_seq_obj[:,k]).T).tolist()
            new_test_seq_obj = np.array(sum(test_seq_obj,[]))
            ts_obj = new_test_seq_obj
            final_ts_obj = ghmm.EmissionSequence(F,ts_obj.tolist())
            # Find Viterbi Path
            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)
            obj = max(path_rf_obj[1],path_rm_obj[1],path_sf_obj[1],path_sm_obj[1])
            if obj == path_rf_obj[1]:
                rf[0,k] = 1
            elif obj == path_rm_obj[1]:
                rm[0,k] = 1
            elif obj == path_sf_obj[1]:
                sf[0,k] = 1
            else:
                sm[0,k] = 1
            k = k+1
        #print rf.T
        rf_final = rf_final + rf.T
        rm_final = rm_final + rm.T
        sf_final = sf_final + sf.T
        sm_final = sm_final + sm.T
        trial_number = trial_number + 1

    #print rf_final
    #print rm_final
    #print sf_final
    #print sm_final
 
    # Confusion Matrix
    cmat = np.zeros((4,4))
    arrsum_rf = np.zeros((4,1))
    arrsum_rm = np.zeros((4,1))
    arrsum_sf = np.zeros((4,1))
    arrsum_sm = np.zeros((4,1))
    k = 7
    i = 0
    while (k < 29):
        arrsum_rf[i] = np.sum(rf_final[k-7:k,0])
        arrsum_rm[i] = np.sum(rm_final[k-7:k,0])
        arrsum_sf[i] = np.sum(sf_final[k-7:k,0])
        arrsum_sm[i] = np.sum(sm_final[k-7:k,0])
        i = i+1
        k = k+7
        
    
    i=0
    while (i < 4):
        j=0
        while (j < 4):
            if (i == 0):
                cmat[i][j] = arrsum_rf[j]
            elif (i == 1):
                cmat[i][j] = arrsum_rm[j]
            elif (i == 2):
                cmat[i][j] = arrsum_sf[j]
            else:
                cmat[i][j] = arrsum_sm[j]
            j = j+1
        i = i+1

    #print cmat
    
    # Plot Confusion Matrix
    #Nlabels = 4
    #fig = pp.figure()
    #ax = fig.add_subplot(111)
    #figplot = ax.matshow(cmat, interpolation = 'nearest', origin = 'upper', extent=[0, Nlabels, 0, Nlabels])
    #ax.set_title('Performance of HMM Models')
    #pp.xlabel("Targets")
    #pp.ylabel("Predictions")
    #ax.set_xticks([0.5,1.5,2.5,3.5])
    #ax.set_xticklabels(['Rigid-Fixed', 'Rigid-Movable', 'Soft-Fixed', 'Soft-Movable'])
    #ax.set_yticks([3.5,2.5,1.5,0.5])
    #ax.set_yticklabels(['Rigid-Fixed', 'Rigid-Movable', 'Soft-Fixed', 'Soft-Movable'])
    #figbar = fig.colorbar(figplot)

    #i = 0    
    #while (i < 4):
        #j = 0
        #while (j < 4):
            #pp.text(j+0.5,3.5-i,cmat[i][j])
            #j = j+1
        #i = i+1
    #pp.show()
    
    
        

    
    
   

    
    
