# Thomas Johnston
# Physics 250
# Final Project

# Filename: Walker.py


from numpy import array, arange, sin, cos, size, shape, transpose, max, \
        identity, real
from pylab import plot, xlabel, ylabel, title, figure, \
     subplot, show, legend, axis
import numpy.linalg as LA
from WalkerDefs import *
from WalkerPlot import animate
import time
from scipy.optimize import fsolve


# Model Parameters
gamma = 0.009               # radians
L = 1.0                     # dimensionless
Mhip = 1.0                  # dimensionless
g = 1.0
params = array([gamma, L, Mhip, g])


# Initial Conditions (state)
#theta = 0.2                 # radians
#phi = 0.4001                # radians
#thetadot = -0.2             # rad/non-dimensional time (NDT)
#phidot = 0                  # rad/NDT
#X = array([theta, phi, thetadot, phidot])
#Xto = X
#Xk = X


# Integrate the ODEs
dt = 0.01                   # time step size
tmin = 0.0                  # start time of integration
err = 1.0e6                 # dummy value
err_tol = 1.0e-12            # error tolerance
#err_tol = 1.0e-6            # error tolerance


# Initial Empty Lists to Store State Values
Time = [tmin]
X_LP =[]
X_SP = []
X_jacob = []
X_LP_eigvals = []
X_SP_eigvals = []
#eigvals_LP = []
#eigvals_SP = []
thetastar = []


# Range of ground slopes
#gamma = arange(0.009,0.0145,0.0005)
#X = array([ 0.20, 0.40, -0.20, -0.016 ])

gamma = arange(0.0005,0.045,0.0005)
X = array([ 0.07669, 0.1538, -0.07987, -0.0009435 ])
Xto = X         # state at "toe off" (beginning of step)
Xk = X          # kth state, before an integration step
print 'the fixed point of the limit cycle is:\n',
print '   [theta]       [phi]    [thetadot]    [phidot]   [gamma]  [max_eigval]'
for i in xrange(len(gamma)):
    params = array([gamma[i], L, Mhip, g])
    xstar = fsolve(StrideMapError,X, args=(Xto,Xk,params,dt,err,Time))
    #xstar,j = NRA(StrideMapError,xstar, (Xto,Xk,params,dt,err,Time))
    X = xstar
    Xto = X
    X_LP.append( X )            # collect fixed point solutions
    thetastar.append( X[0] )    # stance leg angle at the fixed point
    j,f0 = Jacob_StrideMap(StrideMap,X,(Xto,Xk,params,dt,err,Time))
    eigvals_LP = LA.eigvals(j)
    X_LP_eigvals.append( eigvals_LP )
    abs_eigvals_LP = abs( eigvals_LP )
    print xstar, ' ',gamma[i], ' ',max(abs_eigvals_LP)


print '\n\n'
tstar = []
X = array([ 0.07479, 0.14958, -0.08111, -0.00091 ])
print 'the fixed point of the limit cycle is:\n',
print '   [theta]       [phi]    [thetadot]    [phidot]   [gamma]  [max_eigval]'

for i in xrange(len(gamma)):
    params = array([gamma[i], L, Mhip, g])
    xstar = fsolve(StrideMapError,X,args=(Xto,Xk,params,dt,err,Time))
    X = xstar
    Xto = xstar
    X_SP.append( X )
    tstar.append( X[0] )
    j,f0 = Jacob_StrideMap(StrideMap,X,(Xto,Xk,params,dt,err,Time))
    eigvals_SP = LA.eigvals(j)
    X_SP_eigvals.append( eigvals_SP )
    abs_eigvals_SP = abs( eigvals_SP )
    print xstar, ' ',gamma[i], ' ',max(abs_eigvals_SP)


# Plot results
scale = []
for i in gamma:
    scale.append( i**(1./3.) )
figure(1)
plot(gamma,thetastar,'b*')
plot(gamma,scale)
plot(gamma,tstar,'ro')
xlabel('ground slope, $\\gamma$ (rad)')
ylabel('stance leg angle at fixed point, $\\theta^*$')
axis([0, 0.045, 0, 0.35])



"""
# Code to generate state trajectories for the swing phase
# an attempt to recreate Figure 2 in Garcia's paper
# Model Parameters
gamma = 0.009               # radians
L = 1.0                     # dimensionless
Mhip = 1.0                  # dimensionless
g = 1.0
params = array([gamma, L, Mhip, g])

# Initial Conditions (state)
theta = 0.2                 # radians
phi = 0.4001                # radians
thetadot = -0.2             # rad/non-dimensional time (NDT)
phidot = 0                  # rad/NDT
X = array([theta, phi, thetadot, phidot])
Xto = X
Xk = X

# Initial Empty Lists to Store State Values
TrajTheta = [theta]
TrajPhi = [phi]
TrajThetadot = [thetadot]
TrajPhidot = [phidot]
Time = [tmin]

Time, TrajTheta, TrajPhi, TrajThetadot, TrajPhidot = SwingPhase(X, \
        (Xto,Xk,params,dt,err,TrajTheta,TrajPhi, \
        TrajThetadot,TrajPhidot,Time))

print 'The time at foot strike is:',Time[-1]

figure(2)
plot(Time,TrajTheta,'b-',Time,TrajPhi,'r--')
xlabel('Non-dimensional Time')
ylabel('stance & swing leg angle, (rad)')
title('$\\gamma$ = 0.009 rad')
#show()
"""

X_LP_theta_real = []
X_LP_phi_real = []
X_SP_theta_real = []
X_SP_phi_real = []
for i in xrange( len(X_LP_eigvals) ):
    X_LP_theta_real.append( real(X_LP_eigvals[i][0]) )
    X_LP_phi_real.append( real(X_LP_eigvals[i][1]) )
    X_SP_theta_real.append( real(X_SP_eigvals[i][0]) )
    X_SP_phi_real.append( real(X_SP_eigvals[i][1]) )

#print X_LP_theta_real
#print type(X_LP_theta_real)
#print len(X_LP_theta_real)


figure(3)
plot(gamma,X_LP_eigvals)
#plot(gamma,X_LP_phi_real)
title('Long Period Gait')
xlabel('ground slope, $\\gamma$ (rad)')
ylabel('eigenvalues')
axis([0, 0.045, -6, 1])

figure(4)
plot(gamma,X_SP_theta_real,'bx')
plot(gamma,X_SP_phi_real,'rx')
title('Short Period Gait')
xlabel('ground slope, $\\gamma$ (rad)')
ylabel('eigenvalues')
axis([0, 0.045, -1, 8])
show()