#!/usr/bin/python
# 7.1
import os, sys
sys.path += [os.getcwdu() + os.sep + ".." + os.sep + "python"]

import pylab as pl
from scipy import signal

def rlocus(H, k_vec=None):
    """Plots the Root Locus given 
    an lti system and 
    a vector of gain values"""
    if k_vec is None:
        k_vec = linspace(0, 100, 10000)

    num = pl.poly1d(H.num)
    den = pl.poly1d(H.den)
    
    roots = []
    for k in k_vec:
        p = den + k*num
        roots.append(pl.roots(p))
    roots = pl.array(roots)

    # Tried to do it with array slicing. Looks BAD
    # n = H.num
    # d = H.den
    # l = d.shape[0]-n.shape[0]
    # n_vec = n[pl.newaxis,:]*k_vec[:,pl.newaxis]
    # A = pl.zeros((k_vec.shape[0], d.shape[0]))
    # A = A+d
    # A[:,l:] = n_vec 
    # r = map(pl.roots, A)
    # r = pl.array(r)

    x = pl.linspace(0, 2*pl.pi, 100)
    pl.plot(pl.cos(x), pl.sin(x), 'b.', markersize=2)

    for i, k_root in enumerate(roots[0]):
        pl.plot(pl.real(roots[:,i]), pl.imag(roots[:,i]))
  
    xmin, xmax = pl.xlim()
    x = pl.linspace(int(xmin), int(xmax)) 
    ymin, ymax = pl.ylim()
    y = pl.linspace(int(ymin), int(ymax))

    pl.plot(x, pl.zeros_like(x), 'b.', pl.zeros_like(y), y, 'b.', markersize=2)

    pl.xlim(xmin, xmax)
    pl.ylim(ymin, ymax)

    pl.xlabel('Real part')
    pl.ylabel('Imaginary part')

    pl.show()

n = [1]
d = [1, -1, 0]
k_vec = pl.linspace(0, 20, 10000)
H = signal.lti(n, d)

rlocus(H, k_vec)