#!/usr/bin/python

import pylab as pl
from scipy.linalg import schur
from dscr import dscr
from scipy import signal

def dlyap(a, b):
    n = len(a)
    x = pl.zeros_like(a)
    s, u = schur(a)
    b = pl.dot(u.T, pl.dot(b,u))
    j = n-1
    while j>=0:
        k = j
        ## Check for Schur block.
        if j==0:
            blksiz = 1
        elif s[j, j-1]!=0:
            blksiz = 2
            j = j - 1
        else:
            blksiz = 1
        Ajj = pl.kron(s[j:k+1,j:k+1], s) - pl.eye(blksiz*n)
        rhs = pl.reshape(b[:,j:k+1].T, (blksiz*n, 1))
        if (k < n-1):
            rhs2 = pl.dot(s, pl.dot(x[:,k+1:n], s[j:k+1, k+1:n].T))
            rhs = rhs + pl.reshape(rhs2, (blksiz*n, 1))
        v = -pl.solve(Ajj, rhs)
        x[:,j] = v.squeeze()[:n]
        if(blksiz == 2):
            x[:, k] = v[n:blksiz*n].squeeze()
        j = j - 1

    ## Back-transform to original coordinates.
    x = pl.dot(u, pl.dot(x, u.T))
    return x 


def covar_m(H, W):
    """
    User defined equivalent function to Matlab covar function
    For discrete time domain only
    Uses Lyapunov's equation for computation
    W: noise intensity (scalar)
    """
    a = pl.roots(H.den)
    if pl.any(abs(a) > 1):
#        print "Warning: System being unstable has infinite covariance"
        P = pl.inf
        return P
    else:
        A, B, C, D = H.A, H.B, H.C, H.D
        # Sylvester and Lyapunov solver
        Q1 = pl.dot(-B, pl.dot(W, B.T))
        Q = dlyap(A, -Q1)
        # Q = linmeq(2,A,Q1,[1, 0],1)
        # A*X*A' - X + B*W*B' = 0,                          (2b)
        # Discrete time Lyapunov equation; A is general form. Hessenberg-Schur method.
        # linmeq(2, A, C, [1,0], 1)
        # A*X*A' - X = C,                          (2b)
        #
        P = pl.dot(C, pl.dot(Q,C.T)) + pl.dot(D, pl.dot(W,D.T)) 

    return P