# 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 scipy.io

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

# Returns mu,sigma for 20 hidden-states from feature-vectors(123,35) for Smooth, Moderate, and Rough Surface 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__':

### Simulation Data
    tSamples = 400
    datasmooth = scipy.io.loadmat('smooth.mat')
    datamoderate = scipy.io.loadmat('medium.mat')
    datarough = scipy.io.loadmat('rough.mat')
    simulforce = np.zeros((tSamples,150))
    datatime = np.arange(0,7,0.01)
    dataforceSmooth = np.transpose(datasmooth['force'])
    dataforceModerate = np.transpose(datamoderate['force'])
    dataforceRough = np.transpose(datarough['force'])
    j = 0
    for i in dataforceSmooth:
        simulforce[:,j] = i
        j = j+1
    j = 50
    for i in dataforceModerate:
        simulforce[:,j] = i
        j = j+1
    j = 100
    for i in dataforceRough:
        simulforce[:,j] = i
        j = j+1

    Fmat = np.matrix(simulforce)
          
# Checking the Data-Matrix

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

    mu_smooth,sigma_smooth = feature_to_mu_sigma(Fmat[0:tSamples,0:50])
    mu_moderate,sigma_moderate = feature_to_mu_sigma(Fmat[0:tSamples,50:100])
    mu_rough,sigma_rough = feature_to_mu_sigma(Fmat[0:tSamples,100:150])
    
    #print [mu_smooth, sigma_smooth]

# HMM - Implementation:

# 10 Hidden States
# Force as Continuous Gaussian Observations from each hidden state
# Three HMM-Models for Smooth, Moderate, Rough Surfaces
# 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_smooth = np.zeros((15,2))
    B_moderate = np.zeros((15,2))
    B_rough = np.zeros((15,2))
          
    for num_states in range(15):
        B_smooth[num_states,0] = mu_smooth[num_states]
        B_smooth[num_states,1] = sigma_smooth[num_states]
        B_moderate[num_states,0] = mu_moderate[num_states]
        B_moderate[num_states,1] = sigma_moderate[num_states]
        B_rough[num_states,0] = mu_rough[num_states]
        B_rough[num_states,1] = sigma_rough[num_states]
        
    B_smooth = B_smooth.tolist()
    B_moderate = B_moderate.tolist()
    B_rough = B_rough.tolist()
        
    # pi - initial probabilities per state
    pi = [1.0/15.0] * 15

    # generate Smooth, Moderate, Rough Surface models from parameters
    model_smooth = ghmm.HMMFromMatrices(F,ghmm.GaussianDistribution(F), A, B_smooth, pi) # Will be Trained
    model_moderate = ghmm.HMMFromMatrices(F,ghmm.GaussianDistribution(F), A, B_moderate, pi) # Will be Trained
    model_rough = ghmm.HMMFromMatrices(F,ghmm.GaussianDistribution(F), A, B_rough, pi) # Will be Trained
    
    trial_number = 1
    smooth_final = np.matrix(np.zeros((30,1)))
    moderate_final = np.matrix(np.zeros((30,1)))
    rough_final = np.matrix(np.zeros((30,1)))
    
    while (trial_number < 6):

        # For Training
   
        total_seq = Fmat[0:tSamples,:]
        m_total, n_total = np.shape(total_seq)
        #print 'Total_Sequence_Shape:', m_total, n_total
        if (trial_number == 1):
            j = 5
            total_seq_smooth = total_seq[0:tSamples,1:5]
            total_seq_moderate = total_seq[0:tSamples,51:55]
            total_seq_rough = total_seq[0:tSamples,101:105]
            while (j < 50):
                total_seq_smooth = np.column_stack((total_seq_smooth,total_seq[0:tSamples,j+1:j+5]))
                total_seq_moderate = np.column_stack((total_seq_moderate,total_seq[0:tSamples,j+51:j+55]))
                total_seq_rough = np.column_stack((total_seq_rough,total_seq[0:tSamples,j+101:j+105]))
                j = j+5

        if (trial_number == 2):
            j = 5
            total_seq_smooth = np.column_stack((total_seq[0:tSamples,0],total_seq[0:tSamples,2:5]))
            total_seq_moderate = np.column_stack((total_seq[0:tSamples,50],total_seq[0:tSamples,52:55]))
            total_seq_rough = np.column_stack((total_seq[0:tSamples,100],total_seq[0:tSamples,102:105]))
            while (j < 50):
                total_seq_smooth = np.column_stack((total_seq_smooth,total_seq[0:tSamples,j+0],total_seq[0:tSamples,j+2:j+5]))
                total_seq_moderate = np.column_stack((total_seq_moderate,total_seq[0:tSamples,j+50],total_seq[0:tSamples,j+52:j+55]))
                total_seq_rough = np.column_stack((total_seq_rough,total_seq[0:tSamples,j+100],total_seq[0:tSamples,j+102:j+105]))
                j = j+5

        if (trial_number == 3):
            j = 5
            total_seq_smooth = np.column_stack((total_seq[0:tSamples,0:2],total_seq[0:tSamples,3:5]))
            total_seq_moderate = np.column_stack((total_seq[0:tSamples,50:52],total_seq[0:tSamples,53:55]))
            total_seq_rough = np.column_stack((total_seq[0:tSamples,100:102],total_seq[0:tSamples,103:105]))
            while (j < 50):
                total_seq_smooth = np.column_stack((total_seq_smooth,total_seq[0:tSamples,j+0:j+2],total_seq[0:tSamples,j+3:j+5]))
                total_seq_moderate = np.column_stack((total_seq_moderate,total_seq[0:tSamples,j+50:j+52],total_seq[0:tSamples,j+53:j+55]))
                total_seq_rough = np.column_stack((total_seq_rough,total_seq[0:tSamples,j+100:j+102],total_seq[0:tSamples,j+103:j+105]))
                j = j+5

        if (trial_number == 4):
            j = 5
            total_seq_smooth = np.column_stack((total_seq[0:tSamples,0:3],total_seq[0:tSamples,4:5]))
            total_seq_moderate = np.column_stack((total_seq[0:tSamples,50:53],total_seq[0:tSamples,54:55]))
            total_seq_rough = np.column_stack((total_seq[0:tSamples,100:103],total_seq[0:tSamples,104:105]))
            while (j < 50):
                total_seq_smooth = np.column_stack((total_seq_smooth,total_seq[0:tSamples,j+0:j+3],total_seq[0:tSamples,j+4:j+5]))
                total_seq_moderate = np.column_stack((total_seq_moderate,total_seq[0:tSamples,j+50:j+53],total_seq[0:tSamples,j+54:j+55]))
                total_seq_rough = np.column_stack((total_seq_rough,total_seq[0:tSamples,j+100:j+103],total_seq[0:tSamples,j+104:j+105]))
                j = j+5

        if (trial_number == 5):
            j = 5
            total_seq_smooth = total_seq[0:tSamples,0:4]
            total_seq_moderate = total_seq[0:tSamples,50:54]
            total_seq_rough = total_seq[0:tSamples,100:104]
            
            while (j < 50):
                total_seq_smooth = np.column_stack((total_seq_smooth,total_seq[0:tSamples,j+0:j+4]))
                total_seq_moderate = np.column_stack((total_seq_moderate,total_seq[0:tSamples,j+50:j+54]))
                total_seq_rough = np.column_stack((total_seq_rough,total_seq[0:tSamples,j+100:j+104]))
                j = j+5
      
        train_seq_smooth = (np.array(total_seq_smooth).T).tolist()
        train_seq_moderate = (np.array(total_seq_moderate).T).tolist()
        train_seq_rough = (np.array(total_seq_rough).T).tolist()
        #m,n = np.shape(train_seq_smooth)
        #print m,n
        #print train_seq_smooth
    
        final_ts_smooth = ghmm.SequenceSet(F,train_seq_smooth)
        final_ts_moderate = ghmm.SequenceSet(F,train_seq_moderate)
        final_ts_rough = ghmm.SequenceSet(F,train_seq_rough)
        
        model_smooth.baumWelch(final_ts_smooth)
        model_moderate.baumWelch(final_ts_moderate)
        model_rough.baumWelch(final_ts_rough)
            
        # For Testing
        
        if (trial_number == 1): 
            j = 5
            total_seq_smooth = total_seq[0:tSamples,0]
            total_seq_moderate = total_seq[0:tSamples,50]
            total_seq_rough = total_seq[0:tSamples,100]
            while (j < 50):
                total_seq_smooth = np.column_stack((total_seq_smooth,total_seq[0:tSamples,j]))
                total_seq_moderate = np.column_stack((total_seq_moderate,total_seq[0:tSamples,j+50]))
                total_seq_rough = np.column_stack((total_seq_rough,total_seq[0:tSamples,j+100]))
                j = j+5
        
        if (trial_number == 2): 
            j = 5
            total_seq_smooth = total_seq[0:tSamples,1]
            total_seq_moderate = total_seq[0:tSamples,51]
            total_seq_rough = total_seq[0:tSamples,101]
            while (j < 50):
                total_seq_smooth = np.column_stack((total_seq_smooth,total_seq[0:tSamples,j+1]))
                total_seq_moderate = np.column_stack((total_seq_moderate,total_seq[0:tSamples,j+51]))
                total_seq_rough = np.column_stack((total_seq_rough,total_seq[0:tSamples,j+101]))
                j = j+5 

        if (trial_number == 3): 
            j = 5
            total_seq_smooth = total_seq[0:tSamples,2]
            total_seq_moderate = total_seq[0:tSamples,52]
            total_seq_rough = total_seq[0:tSamples,102]
            while (j < 50):
                total_seq_smooth = np.column_stack((total_seq_smooth,total_seq[0:tSamples,j+2]))
                total_seq_moderate = np.column_stack((total_seq_moderate,total_seq[0:tSamples,j+52]))
                total_seq_rough = np.column_stack((total_seq_rough,total_seq[0:tSamples,j+102]))
                j = j+5 

        if (trial_number == 4): 
            j = 5
            total_seq_smooth = total_seq[0:tSamples,3]
            total_seq_moderate = total_seq[0:tSamples,53]
            total_seq_rough = total_seq[0:tSamples,103]
            while (j < 50):
                total_seq_smooth = np.column_stack((total_seq_smooth,total_seq[0:tSamples,j+3]))
                total_seq_moderate = np.column_stack((total_seq_moderate,total_seq[0:tSamples,j+53]))
                total_seq_rough = np.column_stack((total_seq_rough,total_seq[0:tSamples,j+103]))
                j = j+5

        if (trial_number == 5): 
            j = 5
            total_seq_smooth = total_seq[0:tSamples,4]
            total_seq_moderate = total_seq[0:tSamples,54]
            total_seq_rough = total_seq[0:tSamples,104]
            while (j < 50):
                total_seq_smooth = np.column_stack((total_seq_smooth,total_seq[0:tSamples,j+4]))
                total_seq_moderate = np.column_stack((total_seq_moderate,total_seq[0:tSamples,j+54]))
                total_seq_rough = np.column_stack((total_seq_rough,total_seq[0:tSamples,j+104]))
                j = j+5
    
        total_seq_obj = np.matrix(np.column_stack((total_seq_smooth,total_seq_moderate,total_seq_rough)))
       
        smooth = np.matrix(np.zeros(np.size(total_seq_obj,1)))
        moderate = np.matrix(np.zeros(np.size(total_seq_obj,1)))
        rough = np.matrix(np.zeros(np.size(total_seq_obj,1)))
        m,n = np.shape(smooth)
        print m,n
        k = 0
        while (k < np.size(total_seq_obj,1)):
            test_seq_obj = (np.array(total_seq_obj[0:tSamples,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_smooth_obj = model_smooth.viterbi(final_ts_obj)
            path_moderate_obj = model_moderate.viterbi(final_ts_obj)
            path_rough_obj = model_rough.viterbi(final_ts_obj)
            obj = max(path_smooth_obj[1],path_moderate_obj[1],path_rough_obj[1])
            if obj == path_smooth_obj[1]:
                smooth[0,k] = 1
            elif obj == path_moderate_obj[1]:
                moderate[0,k] = 1
            else:
                rough[0,k] = 1
            k = k+1
        #print smooth.T
        smooth_final = smooth_final + smooth.T
        moderate_final = moderate_final + moderate.T
        rough_final = rough_final + rough.T
        trial_number = trial_number + 1

    #print smooth_final
    #print moderate_final
    #print rough_final
     
    # Confusion Matrix
    cmat = np.zeros((3,3))
    arrsum_smooth = np.zeros((3,1))
    arrsum_moderate = np.zeros((3,1))
    arrsum_rough= np.zeros((3,1))
    k = 10
    i = 0
    while (k < 31):
        arrsum_smooth[i] = np.sum(smooth_final[k-10:k,0])
        arrsum_moderate[i] = np.sum(moderate_final[k-10:k,0])
        arrsum_rough[i] = np.sum(rough_final[k-10:k,0])
        i = i+1
        k = k+10
        
    
    i=0
    while (i < 3):
        j=0
        while (j < 3):
            if (i == 0):
                cmat[i][j] = arrsum_smooth[j]
            elif (i == 1):
                cmat[i][j] = arrsum_moderate[j]
            else:
                cmat[i][j] = arrsum_rough[j]
            j = j+1
        i = i+1

    #print cmat
    
    # Plot Confusion Matrix
    Nlabels = 3
    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])
    ax.set_xticklabels(['Smooth', 'Moderate', 'Rough'])
    ax.set_yticks([2.5,1.5,0.5])
    ax.set_yticklabels(['Smooth', 'Moderate', 'Rough'])
    figbar = fig.colorbar(figplot)

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

    
    
   

    
    
