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