#!/usr/bin/env python
"""
Ported by SCL, 27 Oct 2012.
See original comments below.  Requires predprey.py
"""
# L4_1_statefbk.m - MATLAB source code for Lecture 4.1
# RMM, 26 Oct 03 (updated 22 Oct 06; 20 Oct 2012)

import numpy as np
import scipy.linalg as la
from scipy.integrate import odeint
from scipy.optimize import fsolve
import control
from control.matlab import ss
import matplotlib.pyplot as plt

from predprey import predprey, predprey_rh


##
## Cart Pendulum system
##
## This is the cart pendulum system that we have used in lectures and 
## homework.  Some of the parameters have been changed and I have derived
## the equations a bit differently, but the dynamics are basically the
## same.
##

M = 5  # mass of cart, kg
m = 0.5  # mass of pendulum, kg
b = 0  # damping of cart, N/m/sec
l = 0.3  # length of pendulum, m
J = 0.006  # inertia of pendulum, kg m^2
g = 9.8  # gravitational constant, kg m/sec^2

# Compute the inertia matrix for the system
InertiaMatrix = np.array([[M + m, -m*l],
                          [-m*l, J+m*l**2]])

# ...and its inverse
IM_inv = la.inv(InertiaMatrix)

# Build the state space representation
Acp = np.vstack([np.hstack([np.zeros((2,2)), np.eye(2)]),
                 np.hstack([np.dot(IM_inv, np.array([[0, 0],
                                                     [0, m*g*l]])),
                            np.dot(IM_inv, np.array([[-b, 0],
                                                     [0, 0]]))])])
Bcp = np.vstack([np.zeros((2,1)),
                 np.dot(IM_inv, np.array([[1.],[0]]))])
Ccp =np.array([1, -l, 0, 0])
Dcp = 0
cartpend=ss(Acp,Bcp,Ccp,Dcp)

# Build up controllability matrix
Wr = np.hstack([Bcp,
                np.dot(Acp, Bcp),
                np.dot(Acp,np.dot(Acp,Bcp)),
                np.dot(Acp, np.dot(Acp, np.dot(Acp, Bcp)))])

##
## Predator Prey example
##

# Set paramater values for the simulation
alpha = 12.5  # multiplier for hare population
beta = 25  # multipler for lynx population
gamma = 0.8  # multiplier for time scale

r = 2 * gamma  # birth rate of hares
d = 0.7 * gamma  # death rate for lynxes
b = 0.3 * beta / alpha  # birth coefficient for lynxes
k = 10 * alpha  # carrying capacity for hares
a = 8 * gamma * alpha / beta  # interaction coefficient
c = 4 * alpha  # saturation coefficient

# Compute the equilibrium point (Note: parameters must be set in predprey.m)
f = lambda x: predprey(x, 0)
xeq = fsolve(f, np.array([20, 30]))
He = xeq[0]
Le = xeq[1]

# Generate the linearization around the equilibrium point
# App, Bpp computed from symbolic expressions (using Mathematica)
App = np.array([[-((a*c*k*Le + (c + He)**2*(2*He - k)*r)/((c + He)**2*k)), -((a*He)/(c + He))],
                [(a*b*c*Le)/(c + He)**2, -d + (a*b*He)/(c + He)]])
Bpp = np.array([[He*(1 - He/k)],
                [0.]])
Cpp = np.array([0, 1.])

# Assign the desired eigenvalues for the system
Kpp = control.place(App, Bpp, [-0.1, -0.2])

# Compute the reference input and gain
yd = 30
ref = 0
krpp = -1 / np.dot(Cpp, np.dot(la.inv(App - np.dot(Bpp, Kpp)), Bpp))

# Now set the parameters for our feedback function
predprey_K = Kpp
predprey_xd = xeq
predprey_ud = krpp * ref

# Initial condition and simulation time
x0 = np.array([15., 20])
Tmax = 100

# Generate a simulation of the system
plt.figure(1)
plt.clf()
t = np.linspace(0, Tmax, 100)
x = odeint(predprey, x0, t,
           args=(-r, d, b, k, a, c, (predprey_K, predprey_xd, predprey_ud)))
plt.plot(t, x[:,0], '-', t, x[:,1], '--')
plt.xlabel('Time (years)')
plt.ylabel( 'Population')

plt.legend(('Hare', 'Lynx'))
plt.savefig('predprey-sim.png')

# Save the final point we reached
xfin = x[-1, :]

# Generate the phase portrait
plt.figure(2)
plt.clf()
plt.hold(True)
control.phase_plot(predprey, [0, 50, 5], [0, 50, 6], 0.2,
                   parms=(-r, d, b, k, a, c,
                           (predprey_K, predprey_xd, predprey_ud)))
plt.hold(True)
control.phase_plot(predprey, [0, 50, 5], [60, 100, 4], 0.13,
                   parms=(-r, d, b, k, a, c,
                           (predprey_K, predprey_xd, predprey_ud)))
plt.hold(True)
control.phase_plot(predprey, [60, 100, 4], [0, 50, 6], 0.15,
                   parms=(-r, d, b, k, a, c,
                           (predprey_K, predprey_xd, predprey_ud)))
plt.hold(True)
control.phase_plot(predprey, [60, 100, 4], [60, 100, 4], 0.1,
                   parms=(-r, d, b, k, a, c,
                           (predprey_K, predprey_xd, predprey_ud)))
plt.hold(True)
control.phase_plot(predprey, [0, 100, 10], [0, 100, 10], 0,
                   [[0.5, 10],
                    [100, 45],
                    [100, 70]],
                   Tmax/2,
                   parms=(-r, d, b, k, a, c,
                           (predprey_K, predprey_xd, predprey_ud)))
plt.xlabel('Hares')
plt.ylabel('Lynxes')
plt.axis([0, 100, 0, 100])

plt.plot(xfin[0], xfin[1], 'ro')
plt.savefig('predprey-pp.png')

# Check the linearization to make sure that computations are OK
eqchk = xeq - np.dot(la.inv(App - np.dot(Bpp, Kpp)), np.dot(Bpp, krpp) * r)
print eqchk

plt.show()