# 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/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((15,1)))
    sigma = np.matrix(np.zeros((15,1)))
    DIVS = m/15

    while (index < 15):
        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

# Returns sequence given raw data
  
def create_seq(fvec): 
    
    m,n = np.shape(fvec)
    #print m,n
    seq = np.matrix(np.zeros((15,n)))
    DIVS = m/15
    for i in range(n):
        index = 0
        while (index < 15):
            m_init = index*DIVS
            temp_fvec = fvec[(m_init):(m_init+DIVS),i]
            #if index == 1:
                #print temp_fvec
            seq[index,i] = scp.mean(temp_fvec)
            index = index+1
        
    return seq


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])
    mu_obj1,sigma_obj1 = feature_to_mu_sigma(Fmat[0:121,140:141])
    mu_obj2,sigma_obj2 = feature_to_mu_sigma(Fmat[0:121,141:142])
    #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.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]]

    # B - Emission Matrix, parameters of emission distributions in pairs of (mu, sigma)
    
    B_rf = np.zeros((15,2))
    B_rm = np.zeros((15,2))
    B_sf = np.zeros((15,2))
    B_sm = np.zeros((15,2))
      
    for num_states in range(15):
        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 = [1.0/15.0] * 15

    # 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

    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)))

    while (trial_number < 6):

        # For Training
   
        total_seq = Fmat[0:121,:]
        m_total, n_total = np.shape(total_seq)
        #print 'Total_Sequence_Shape:', m_total, n_total
        if (trial_number == 1):
            j = 5
            total_seq_rf = total_seq[0:121,1:5]
            total_seq_rm = total_seq[0:121,36:40]
            total_seq_sf = total_seq[0:121,71:75]
            total_seq_sm = total_seq[0:121,106:110]
            while (j < 35):
                total_seq_rf = np.column_stack((total_seq_rf,total_seq[0:121,j+1:j+5]))
                total_seq_rm = np.column_stack((total_seq_rm,total_seq[0:121,j+36:j+40]))
                total_seq_sf = np.column_stack((total_seq_sf,total_seq[0:121,j+71:j+75]))
                total_seq_sm = np.column_stack((total_seq_sm,total_seq[0:121,j+106:j+110]))
                j = j+5

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

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

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

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

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

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

        if (trial_number == 5): 
            j = 5
            total_seq_rf = total_seq[0:121,4]
            total_seq_rm = total_seq[0:121,39]
            total_seq_sf = total_seq[0:121,74]
            total_seq_sm = total_seq[0:121,109]
            while (j < 35):
                total_seq_rf = np.column_stack((total_seq_rf,total_seq[0:121,j+4]))
                total_seq_rm = np.column_stack((total_seq_rm,total_seq[0:121,j+39]))
                total_seq_sf = np.column_stack((total_seq_sf,total_seq[0:121,j+74]))
                total_seq_sm = np.column_stack((total_seq_sm,total_seq[0:121,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[0:121,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()
    
    
        

    
    
   

    
    
