# 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/Variable_length')

from data_variable_length_force import Fmat_original

  
if __name__ == '__main__' or __name__ != '__main__':

    print "Inside outlier HMM model training file"

    Fmat = Fmat_original

    
# Getting mean / covariance

    i = 0
    number_states = 10
    feature_1_final_data = [0.0]*number_states
    state_1 = [0.0]
    while (i < 35):
        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_rf_force = np.zeros((number_states,1))
    sigma_rf = np.zeros((number_states,1))
    
    while (j < number_states):
        mu_rf_force[j] = np.mean(feature_1_final_data[j])
        sigma_rf[j] = scp.std(feature_1_final_data[j])
        j = j+1

    i = 35
    feature_1_final_data = [0.0]*number_states
    state_1 = [0.0]
    while (i < 70):
        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 == 35:
            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_rm_force = np.zeros((number_states,1))
    sigma_rm = np.zeros((number_states,1))

    while (j < number_states):
        mu_rm_force[j] = np.mean(feature_1_final_data[j])
        sigma_rm[j] = scp.std(feature_1_final_data[j])
        j = j+1

    i = 70
    feature_1_final_data = [0.0]*number_states
    state_1 = [0.0]
    while (i < 105):
        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 == 70:
            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_sf_force = np.zeros((number_states,1))
    sigma_sf = np.zeros((number_states,1))

    while (j < number_states):
        mu_sf_force[j] = np.mean(feature_1_final_data[j])
        sigma_sf[j] = scp.std(feature_1_final_data[j])  
        j = j+1

    i = 105
    feature_1_final_data = [0.0]*number_states
    state_1 = [0.0]
    while (i < 140):
        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 == 105:
            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_sm_force = np.zeros((number_states,1))
    sigma_sm = np.zeros((number_states,1))

    while (j < number_states):
        mu_sm_force[j] = np.mean(feature_1_final_data[j])
        sigma_sm[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_rf = [0.0]*number_states
    B_rm = [0.0]*number_states
    B_sf = [0.0]*number_states
    B_sm = [0.0]*number_states
      
    for num_states in range(number_states):
        B_rf[num_states] = [mu_rf_force[num_states][0],sigma_rf[num_states][0]]
        B_rm[num_states] = [mu_rm_force[num_states][0],sigma_rm[num_states][0]]        
        B_sf[num_states] = [mu_sf_force[num_states][0],sigma_sf[num_states][0]]
        B_sm[num_states] = [mu_sm_force[num_states][0],sigma_sm[num_states][0]]

    #print B_sm
    #print mu_sm_motion
    
    # 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 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)))

    total_seq = Fmat
    for i in range(140):
        total_seq[i][:] = sum(total_seq[i][:],[])

    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]
            #print total_seq_rf
            while (j < 35):
                total_seq_rf = total_seq_rf+total_seq[j+1:j+5]
                total_seq_rm = total_seq_rm+total_seq[j+36:j+40]
                total_seq_sf = total_seq_sf+total_seq[j+71:j+75]
                total_seq_sm = total_seq_sm+total_seq[j+106:j+110]
                j = j+5

        if (trial_number == 2):
            j = 5
            total_seq_rf = [total_seq[0]]+total_seq[2:5]
            total_seq_rm = [total_seq[35]]+total_seq[37:40]
            total_seq_sf = [total_seq[70]]+total_seq[72:75]
            total_seq_sm = [total_seq[105]]+total_seq[107:110]
            #print total_seq_rf
            while (j < 35):
                total_seq_rf = total_seq_rf+[total_seq[j+0]]+total_seq[j+2:j+5]
                total_seq_rm = total_seq_rm+[total_seq[j+35]]+total_seq[j+37:j+40]
                total_seq_sf = total_seq_sf+[total_seq[j+70]]+total_seq[j+72:j+75]
                total_seq_sm = 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 = total_seq[0:2]+total_seq[3:5]
            total_seq_rm = total_seq[35:37]+total_seq[38:40]
            total_seq_sf = total_seq[70:72]+total_seq[73:75]
            total_seq_sm = total_seq[105:107]+total_seq[108:110]
            while (j < 35):
                total_seq_rf = total_seq_rf+total_seq[j+0:j+2]+total_seq[j+3:j+5]
                total_seq_rm = total_seq_rm+total_seq[j+35:j+37]+total_seq[j+38:j+40]
                total_seq_sf = total_seq_sf+total_seq[j+70:j+72]+total_seq[j+73:j+75]
                total_seq_sm = 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 = total_seq[0:3]+total_seq[4:5]
            total_seq_rm = total_seq[35:38]+total_seq[39:40]
            total_seq_sf = total_seq[70:73]+total_seq[74:75]
            total_seq_sm = total_seq[105:108]+total_seq[109:110]
            while (j < 35):
                total_seq_rf = total_seq_rf+total_seq[j+0:j+3]+total_seq[j+4:j+5]
                total_seq_rm = total_seq_rm+total_seq[j+35:j+38]+total_seq[j+39:j+40]
                total_seq_sf = total_seq_sf+total_seq[j+70:j+73]+total_seq[j+74:j+75]
                total_seq_sm = 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 = total_seq_rf+total_seq[j+0:j+4]
                total_seq_rm = total_seq_rm+total_seq[j+35:j+39]
                total_seq_sf = total_seq_sf+total_seq[j+70:j+74]
                total_seq_sm = total_seq_sm+total_seq[j+105:j+109]
                j = j+5
      
        train_seq_rf = total_seq_rf
        train_seq_rm = total_seq_rm
        train_seq_sf = total_seq_sf
        train_seq_sm = total_seq_sm
        #print train_seq_rf[27]
    
        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]]
            #print np.shape(total_seq_rf)
            while (j < 35):
                total_seq_rf = total_seq_rf+[total_seq[j]]
                total_seq_rm = total_seq_rm+[total_seq[j+35]]
                total_seq_sf = total_seq_sf+[total_seq[j+70]]
                total_seq_sm = 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 = total_seq_rf+[total_seq[j+1]]
                total_seq_rm = total_seq_rm+[total_seq[j+36]]
                total_seq_sf = total_seq_sf+[total_seq[j+71]]
                total_seq_sm = 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 = total_seq_rf+[total_seq[j+2]]
                total_seq_rm = total_seq_rm+[total_seq[j+37]]
                total_seq_sf = total_seq_sf+[total_seq[j+72]]
                total_seq_sm = 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 = total_seq_rf+[total_seq[j+3]]
                total_seq_rm = total_seq_rm+[total_seq[j+38]]
                total_seq_sf = total_seq_sf+[total_seq[j+73]]
                total_seq_sm = 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 = total_seq_rf+[total_seq[j+4]]
                total_seq_rm = total_seq_rm+[total_seq[j+39]]
                total_seq_sf = total_seq_sf+[total_seq[j+74]]
                total_seq_sm = total_seq_sm+[total_seq[j+109]]
                j = j+5
    
        trial_number = trial_number + 1

    print "Outlier HMM model trained"
    
    
    
        

    
    
   

    
    
