#!/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)