
## Generating simulated data for variations in the object parameters:

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
import random

### User-defined functions

def rigid_fixed(K_robot, K_rf, eqbm_point_2, eqbm_point_3):
    
    Robot_Home_Position = 0.00
    Robot_Current_Position = 0.60
    time = np.arange(0.00,1.21,0.01)

    eq1 = 0.61
    atemp = np.linspace(eq1,eqbm_point_2,num=15)
    step = atemp[1]-atemp[0]
    btemp = np.linspace(eqbm_point_2+step,eqbm_point_3,num=15)
    step2 = btemp[1]-btemp[0]

    dist_eqbm_pt2 = np.arange(Robot_Current_Position,eqbm_point_2,step)
    dist_eqbm_pt3 = np.arange(eqbm_point_2+step2,eqbm_point_3,step2)
    dist = np.concatenate((dist_eqbm_pt2,dist_eqbm_pt3),axis=0)
    #print len(dist)
    #print len(time)    
    applied_force_rf = np.zeros((len(dist),1))
    deform_rf = np.zeros((len(dist),1))
    sensed_force_rf = np.zeros((len(time),1))
    robot_pos_rf = np.zeros((len(time),1))
    for i in range(len(robot_pos_rf)):
        robot_pos_rf[i] = Robot_Current_Position
    for i in range(1,len(dist)):
    
        applied_force_rf[i] = K_robot*(dist[i] - Robot_Current_Position)
        deform_rf[i] = applied_force_rf[i]/K_rf
        if i == 1:
            if (Robot_Current_Position + deform_rf[i]) <= dist[i]:
                Robot_Current_Position = Robot_Current_Position + deform_rf[i]
            else:
                Robot_Current_Position = Robot_Current_Position + (dist[i] - Robot_Current_Position)
        else:
            if (Robot_Current_Position + deform_rf[i]) <= dist[i]:
                Robot_Current_Position = Robot_Current_Position + deform_rf[i] - deform_rf[i-1]
            else:
                Robot_Current_Position = Robot_Current_Position + (dist[i] - dist[i-1])
         
        sensed_force_rf[i] = K_rf*(Robot_Current_Position - 0.60)
        robot_pos_rf[i] = Robot_Current_Position

    for i in range(len(dist),len(time)):
        #print i
        sensed_force_rf[i] = sensed_force_rf[i-1]
        robot_pos_rf[i] = Robot_Current_Position

    force = sum(sensed_force_rf.tolist(),[])
    pos = sum(robot_pos_rf.tolist(),[])

    #print force
    #print pos
    #print np.shape(pos)
    #print np.shape(force)
    return pos,force

def soft_fixed(K_robot, K_sf, eqbm_point_2, eqbm_point_3):

    Robot_Home_Position = 0.00
    Robot_Current_Position = 0.60
    time = np.arange(0.00,1.21,0.01)

    eq1 = 0.61
    atemp = np.linspace(eq1,eqbm_point_2,num=15)
    step = atemp[1]-atemp[0]
    btemp = np.linspace(eqbm_point_2+step,eqbm_point_3,num=15)
    step2 = btemp[1]-btemp[0]

    dist_eqbm_pt2 = np.arange(Robot_Current_Position,eqbm_point_2,step)
    dist_eqbm_pt3 = np.arange(eqbm_point_2+step2,eqbm_point_3,step2)
    dist = np.concatenate((dist_eqbm_pt2,dist_eqbm_pt3),axis=0)

    applied_force_sf = np.zeros((len(dist),1))
    deform_sf = np.zeros((len(dist),1))
    sensed_force_sf = np.zeros((len(time),1))
    robot_pos_sf = np.zeros((len(time),1))
    for i in range(len(robot_pos_sf)):
        robot_pos_sf[i] = Robot_Current_Position

    for i in range(1,len(dist)):
    
        applied_force_sf[i] = K_robot*(dist[i] - Robot_Current_Position)
        deform_sf[i] = applied_force_sf[i]/K_sf
        if i == 1:
            if (Robot_Current_Position + deform_sf[i]) <= dist[i]:
                Robot_Current_Position = Robot_Current_Position + deform_sf[i]
            else:
                Robot_Current_Position = Robot_Current_Position + (dist[i] - Robot_Current_Position)
        else:
            if (Robot_Current_Position + deform_sf[i]) <= dist[i]:
                Robot_Current_Position = Robot_Current_Position + deform_sf[i] - deform_sf[i-1]
            else:
                Robot_Current_Position = Robot_Current_Position + (dist[i] - dist[i-1])
         
        sensed_force_sf[i] = K_sf*(Robot_Current_Position - 0.60)
        robot_pos_sf[i] = Robot_Current_Position

    for i in range(len(dist),len(time)):
        sensed_force_sf[i] = sensed_force_sf[i-1]
        robot_pos_sf[i] = Robot_Current_Position

    force = sum(sensed_force_sf.tolist(),[])
    pos = sum(robot_pos_sf.tolist(),[])

    return pos,force

def rigid_movable(K_robot, K_rm, Mass_rm, mu_static_rigid, mu_dynamic_rigid, eqbm_point_2, eqbm_point_3):

    Robot_Home_Position = 0.00
    Robot_Current_Position = 0.60
    time = np.arange(0.00,1.21,0.01)

    g = 9.81
    eq1 = 0.61
    atemp = np.linspace(eq1,eqbm_point_2,num=15)
    step = atemp[1]-atemp[0]
    btemp = np.linspace(eqbm_point_2+step,eqbm_point_3,num=15)
    step2 = btemp[1]-btemp[0]

    dist_eqbm_pt2 = np.arange(Robot_Current_Position,eqbm_point_2,step)
    dist_eqbm_pt3 = np.arange(eqbm_point_2+step2,eqbm_point_3,step2)
    dist = np.concatenate((dist_eqbm_pt2,dist_eqbm_pt3),axis=0)

    applied_force_rm = np.zeros((len(dist),1))
    deform_rm = np.zeros((len(dist),1))
    acc_rm = np.zeros((len(dist),1))
    vel_rm = np.zeros((len(dist),1))
    
    pos_rm = np.zeros((len(time),1))
    sensed_force_rm = np.zeros((len(time),1))
    robot_pos_rm = np.zeros((len(time),1))
    for i in range(len(robot_pos_rm)):
        robot_pos_rm[i] = Robot_Current_Position

    stat_force = Mass_rm*g*mu_static_rigid
    index = 1
    for i in range(1,len(dist)):
        if index == 1:
            applied_force_rm[i] = K_robot*(dist[i] - Robot_Current_Position)
            deform_rm[i] = applied_force_rm[i]/K_rm
            if i == 1:
                if (Robot_Current_Position + deform_rm[i]) <= dist[i]:
                    Robot_Current_Position = Robot_Current_Position + deform_rm[i]
                else:
                    Robot_Current_Position = Robot_Current_Position + (dist[i] - Robot_Current_Position)
            else:
                if (Robot_Current_Position + deform_rm[i]) <= dist[i]:
                    Robot_Current_Position = Robot_Current_Position + deform_rm[i] - deform_rm[i-1]
                else:
                    Robot_Current_Position = Robot_Current_Position + (dist[i] - dist[i-1])

            sensed_force_rm[i] = K_rm*(Robot_Current_Position - 0.60)
        else:
            applied_force_rm[i] = K_rm*(dist[i] - Robot_Current_Position)

        if (applied_force_rm[i] <= stat_force) and (index == 1):
            sensed_force_rm[i] = sensed_force_rm[i]
            acc_rm[i] = 0
            vel_rm[i] = 0
            pos_rm[i] = 0
            Robot_Current_Position = Robot_Current_Position
        else:
            net_force_rm = applied_force_rm[i] - Mass_rm*g*mu_dynamic_rigid
            sensed_force_rm[i] = Mass_rm*g*mu_dynamic_rigid
            if net_force_rm < 0:
                net_force_rm = 0
                acc_rm[i] = 0
                vel_rm[i] = 0
                pos_rm[i] = 0
                Robot_Current_Position = Robot_Current_Position
            else:
                acc_rm[i] = net_force_rm/Mass_rm
                vel_rm[i] = vel_rm[i-1]+acc_rm[i]*0.01
                pos_rm[i] = pos_rm[i-1]+vel_rm[i]*0.01
                if (Robot_Current_Position + pos_rm[i] - pos_rm[i-1]) <= dist[i]:
                    Robot_Current_Position = Robot_Current_Position + pos_rm[i] - pos_rm[i-1]
                else:
                    Robot_Current_Position = Robot_Current_Position + dist[i] - dist[i-1]
            
            index = index+1

        robot_pos_rm[i] = Robot_Current_Position

    for i in range(len(dist),len(time)):
        sensed_force_rm[i] = sensed_force_rm[i-1]
        pos_rm[i] = pos_rm[i-1]
        robot_pos_rm[i] = Robot_Current_Position

    force = sum(sensed_force_rm.tolist(),[])
    pos = sum(robot_pos_rm.tolist(),[])

    return pos,force

def soft_movable(K_robot, K_sm, Mass_sm, mu_static_soft, mu_dynamic_soft, eqbm_point_2, eqbm_point_3):

    Robot_Home_Position = 0.00
    Robot_Current_Position = 0.60
    time = np.arange(0.00,1.21,0.01)

    g = 9.81
    eq1 = 0.61
    atemp = np.linspace(eq1,eqbm_point_2,num=15)
    step = atemp[1]-atemp[0]
    btemp = np.linspace(eqbm_point_2+step,eqbm_point_3,num=15)
    step2 = btemp[1]-btemp[0]

    dist_eqbm_pt2 = np.arange(Robot_Current_Position,eqbm_point_2,step)
    dist_eqbm_pt3 = np.arange(eqbm_point_2+step2,eqbm_point_3,step2)
    dist = np.concatenate((dist_eqbm_pt2,dist_eqbm_pt3),axis=0)

    applied_force_sm = np.zeros((len(dist),1))
    deform_sm = np.zeros((len(dist),1))
    acc_sm = np.zeros((len(dist),1))
    vel_sm = np.zeros((len(dist),1))
    
    pos_sm = np.zeros((len(time),1))
    sensed_force_sm = np.zeros((len(time),1))
    robot_pos_sm = np.zeros((len(time),1))
    for i in range(len(robot_pos_sm)):
        robot_pos_sm[i] = Robot_Current_Position

    stat_force = Mass_sm*g*mu_static_soft
    index = 1
    for i in range(1,len(dist)):
        if index == 1:
            applied_force_sm[i] = K_robot*(dist[i] - Robot_Current_Position)
            deform_sm[i] = applied_force_sm[i]/K_sm
            if i == 1:
                if (Robot_Current_Position + deform_sm[i]) <= dist[i]:
                    Robot_Current_Position = Robot_Current_Position + deform_sm[i]
                else:
                    Robot_Current_Position = Robot_Current_Position + (dist[i] - Robot_Current_Position)
            else:
                if (Robot_Current_Position + deform_sm[i]) <= dist[i]:
                    Robot_Current_Position = Robot_Current_Position + deform_sm[i] - deform_sm[i-1]
                else:
                    Robot_Current_Position = Robot_Current_Position + (dist[i] - dist[i-1])

            sensed_force_sm[i] = K_sm*(Robot_Current_Position - 0.60)
        else:
            applied_force_sm[i] = K_sm*(dist[i] - Robot_Current_Position)

        if (applied_force_sm[i] <= stat_force) and (index == 1):
            sensed_force_sm[i] = sensed_force_sm[i]
            acc_sm[i] = 0
            vel_sm[i] = 0
            pos_sm[i] = 0
            Robot_Current_Position = Robot_Current_Position
        else:
            net_force_sm = applied_force_sm[i] - Mass_sm*g*mu_dynamic_soft
            sensed_force_sm[i] = Mass_sm*g*mu_dynamic_soft
            if net_force_sm < 0:
                net_force_sm = 0
                acc_sm[i] = 0
                vel_sm[i] = 0
                pos_sm[i] = 0
                Robot_Current_Position = Robot_Current_Position
            else:
                acc_sm[i] = net_force_sm/Mass_sm
                vel_sm[i] = vel_sm[i-1]+acc_sm[i]*0.01
                pos_sm[i] = pos_sm[i-1]+vel_sm[i]*0.01
                if (Robot_Current_Position + pos_sm[i] - pos_sm[i-1]) <= dist[i]:
                    Robot_Current_Position = Robot_Current_Position + pos_sm[i] - pos_sm[i-1]
                else:
                    Robot_Current_Position = Robot_Current_Position + dist[i] - dist[i-1]
            
            index = index+1

        robot_pos_sm[i] = Robot_Current_Position

    for i in range(len(dist),len(time)):
        sensed_force_sm[i] = sensed_force_sm[i-1]
        pos_sm[i] = pos_sm[i-1]
        robot_pos_sm[i] = Robot_Current_Position

    force = sum(sensed_force_sm.tolist(),[])
    pos = sum(robot_pos_sm.tolist(),[])

    return pos,force


### Main Program

if __name__ == '__main__':
    
    time = np.arange(0.00,1.21,0.01)
    eq2 = [0.77, 0.80]
    eq3 = [0.97, 1.00]

    # For Rigid-Fixed

    K_robot = np.zeros((5,1))
    K_rf = np.zeros((100,1))

    K_robot = np.array([55,56,57,58,59])

    for i in range(100):
        K_rf[i] = 1000*(i+1)
    
    row_rf = np.size(K_rf,0)
    row_robot = np.size(K_robot,0)
    row_pt2 = np.size(eq2,0)
    row_pt3 = np.size(eq3,0)
    samples = len(time)
    #print samples
    trials_rf = row_rf*row_robot*row_pt2*row_pt3
    robot_pos_rf = np.zeros((trials_rf,samples))
    #print np.shape(robot_pos_rf)
    sensed_force_rf = np.zeros((trials_rf,samples))
    #print np.shape(robot_pos_rf)
    p=0
    for i in range(100):
        for j in range(5):
            for k in range(2):
                for l in range(2):
                    robot_pos_rf[p,:],sensed_force_rf[p,:] = rigid_fixed(K_robot[j], K_rf[i], eq2[k], eq3[l])
                    p = p+1


    # For Soft-Fixed

    K_robot = np.zeros((5,1))
    K_sf = np.zeros((100,1))

    for i in range(100):
        K_sf[i] = 0.5*(i+1)

    K_robot = np.array([1,3,5,7,9])
    
    row_sf = np.size(K_sf,0)
    row_robot = np.size(K_robot,0)
    row_pt2 = np.size(eq2,0)
    row_pt3 = np.size(eq3,0)
    samples = len(time)
    trials_sf = row_sf*row_robot*row_pt2*row_pt3
    robot_pos_sf = np.zeros((trials_sf,samples))
    sensed_force_sf = np.zeros((trials_sf,samples))
    p=0
    for i in range(100):
        for j in range(5):
            for k in range(2):
                for l in range(2):
                    robot_pos_sf[p,:],sensed_force_sf[p,:] = soft_fixed(K_robot[j], K_sf[i], eq2[k], eq3[l])
                    p = p+1


    # For Rigid_Movable

    K_robot = np.zeros((5,1))
    K_rm = np.zeros((5,1))
    Mass_rm = [2.0, 2.1, 2.2, 2.3, 2.4]
    mu_static_rigid = [0.45, 0.55]
    mu_dynamic_rigid = [0.15, 0.2]

    for i in range(5):
        K_rm[i] = 500*(i+1)

    for i in range(5):
        K_robot[i] = 50*(i+1)
    
    row_K_rm = np.size(K_rm,0)
    row_Mass_rm = np.size(Mass_rm,0)
    row_mu_static_rm = len(mu_static_rigid)
    row_mu_dynamic_rm = len(mu_dynamic_rigid)
    samples = len(time)
    row_robot = np.size(K_robot,0)
    row_pt2 = np.size(eq2,0)
    row_pt3 = np.size(eq3,0)
    trials_rm = row_robot*row_pt2*row_pt3*row_K_rm*row_Mass_rm*row_mu_static_rm*row_mu_dynamic_rm
    robot_pos_rm = np.zeros((trials_rm,samples))
    sensed_force_rm = np.zeros((trials_rm,samples))
    p=0
    for j in range(5):
        for k in range(5):
            for m in range(2):
                for n in range(2):
                    for q in range(5):
                        for r in range(2):
                            for s in range(2):
                                #print p
                                robot_pos_rm[p,:], sensed_force_rm[p,:] = rigid_movable(K_robot[q], K_rm[j], Mass_rm[k], mu_static_rigid[m], mu_dynamic_rigid[n], eq2[r], eq2[s])
                                p=p+1


   # For Soft-Movable

    K_robot = np.zeros((5,1))
    K_sm = np.zeros((5,1))
    Mass_sm = [0.3, 0.35, 0.4, 0.45, 0.5]
    mu_static_soft = [0.15, 0.35]
    mu_dynamic_soft = [0.05, 0.1]

    for i in range(5):
        K_sm[i] = 100*(i+1)

    K_robot = np.array([20,22,24,26,28])
    
    row_K_sm = np.size(K_sm,0)
    row_Mass_sm = np.size(Mass_sm,0)
    row_mu_static_sm = len(mu_static_soft)
    row_mu_dynamic_sm = len(mu_dynamic_soft)
    samples = len(time)
    row_robot = np.size(K_robot,0)
    row_pt2 = np.size(eq2,0)
    row_pt3 = np.size(eq3,0)
    trials_sm = row_robot*row_pt2*row_pt3*row_K_sm*row_Mass_sm*row_mu_static_sm*row_mu_dynamic_sm
    robot_pos_sm = np.zeros((trials_sm,samples))
    sensed_force_sm = np.zeros((trials_sm,samples))
    p=0
    for j in range(5):
        for k in range(5):
            for m in range(2):
                for n in range(2):
                    for q in range(5):
                        for r in range(2):
                            for s in range(2):
                                #print p
                                robot_pos_sm[p,:], sensed_force_sm[p,:] = rigid_movable(K_robot[q], K_sm[j], Mass_sm[k], mu_static_soft[m], mu_dynamic_soft[n], eq2[r], eq2[s])
                                p=p+1


    # Store data

    rf_data = {}
    rm_data = {}
    sf_data = {}
    sm_data = {}

    rf_data['sensed_force_rf'] = sensed_force_rf
    rm_data['sensed_force_rm'] = sensed_force_rm
    sf_data['sensed_force_sf'] = sensed_force_sf
    sm_data['sensed_force_sm'] = sensed_force_sm

    rf_data['robot_pos_rf'] = robot_pos_rf
    rm_data['robot_pos_rm'] = robot_pos_rm
    sf_data['robot_pos_sf'] = robot_pos_sf
    sm_data['robot_pos_sm'] = robot_pos_sm

    scipy.io.savemat('rigid_fixed.mat',rf_data)
    scipy.io.savemat('rigid_movable.mat',rm_data)
    scipy.io.savemat('soft_fixed.mat',sf_data)
    scipy.io.savemat('soft_movable.mat',sm_data)

    # Load data

   
    data_rf = scipy.io.loadmat('rigid_fixed.mat')
    data_sf = scipy.io.loadmat('soft_fixed.mat')
    data_rm = scipy.io.loadmat('rigid_movable.mat')
    data_sm = scipy.io.loadmat('soft_movable.mat')
       
    dataforce_rf = np.transpose(data_rf['sensed_force_rf'])
    dataforce_sf = np.transpose(data_sf['sensed_force_sf'])
    dataforce_rm = np.transpose(data_rm['sensed_force_rm'])
    dataforce_sm = np.transpose(data_sm['sensed_force_sm'])
    
    datamotion_rf = np.transpose(data_rf['robot_pos_rf'])
    datamotion_sf = np.transpose(data_sf['robot_pos_sf'])
    datamotion_rm = np.transpose(data_rm['robot_pos_rm'])
    datamotion_sm = np.transpose(data_sm['robot_pos_sm'])
   

    # Plot data
   
    # Force
   
    mpu.figure(1)
    pp.subplot(221)
    pp.title('Rigid Fixed',fontsize='24')
    pp.xlabel('Time (s)',fontsize='24')
    pp.ylabel('Force (N)',fontsize='24')
    pp.plot(time, dataforce_rf, linewidth=3.0)
    pp.xlim((0.0, 1.3))
    pp.grid('True')

    pp.subplot(222)
    pp.title('Soft Fixed',fontsize='24')
    pp.xlabel('Time (s)',fontsize='24')
    pp.ylabel('Force (N)',fontsize='24')
    pp.plot(time, dataforce_sf, linewidth=3.0)
    pp.xlim((0.0, 1.3))
    pp.grid('True')

    pp.subplot(223)
    pp.title('Rigid Movable',fontsize='24')
    pp.xlabel('Time (s)',fontsize='24')
    pp.ylabel('Force (N)',fontsize='24')
    pp.plot(time, dataforce_rm, linewidth=3.0)
    pp.xlim((0.0, 1.3))
    pp.grid('True')

    pp.subplot(224)
    pp.title('Soft Movable',fontsize='24')
    pp.xlabel('Time (s)',fontsize='24')
    pp.ylabel('Force (N)',fontsize='24')
    pp.plot(time, dataforce_sm, linewidth=3.0)
    pp.xlim((0.0, 1.3))
    pp.grid('True')

    # Position

    mpu.figure(2)
    pp.subplot(221)
    pp.title('Rigid Fixed',fontsize='24')
    pp.xlabel('Position (s)',fontsize='24')
    pp.ylabel('Force (N)',fontsize='24')
    pp.plot(datamotion_rf, dataforce_rf, linewidth=3.0)
    pp.xlim((0.5, 1.0))
    pp.grid('True')

    pp.subplot(222)
    pp.title('Soft Fixed',fontsize='24')
    pp.xlabel('Position (m)',fontsize='24')
    pp.ylabel('Force (N)',fontsize='24')
    pp.plot(datamotion_sf, dataforce_sf, linewidth=3.0)
    pp.xlim((0.5, 1.0))
    pp.grid('True')

    pp.subplot(223)
    pp.title('Rigid Movable',fontsize='24')
    pp.xlabel('Position (m)',fontsize='24')
    pp.ylabel('Force (N)',fontsize='24')
    pp.plot(datamotion_rm, dataforce_rm, linewidth=3.0)
    pp.xlim((0.5, 1.0))
    pp.grid('True')

    pp.subplot(224)
    pp.title('Soft-Movable',fontsize='24')
    pp.xlabel('Position (m)',fontsize='24')
    pp.ylabel('Force (N)',fontsize='24')
    pp.plot(datamotion_sm, dataforce_sm, linewidth=3.0)
    pp.xlim((0.5, 1.0))
    pp.grid('True')
    
    pp.show() 

    




    
