import Polynomial
reload(Polynomial)
from Polynomial import *

import DifferenceEqIISM
reload(DifferenceEqIISM)
from DifferenceEqIISM import *

class SystemFunction:
    # Represent a system function as a ratio of polynomials in R
    
    # Input to the initializer is two instances of the Polynomial class
    def __init__(self, numeratorPoly, denominatorPoly):
        self.numerator = numeratorPoly
        self.denominator = denominatorPoly

    # Return a list of the poles and print out a description of the
    # overall stability of the system (ignoring input)
    def poles(self, verbose = True):
        # The poles of a system are the roots of the denominator
        # polynomial in z, which is 1/R.  To get that polynomial,
        # reverse the coefficients of our polynomial (which is in R).

        # Your code here
        pass

    # If this equation is second order and homogeneous (has no
    # dependence on input), then return a function that consumes n as
    # input and computes y[n] directly (in time independent of n) and
    # returns it.  Takes list of two initial conditions [y[0], [y1]]
    # as input.
    def closedForm(self, initialValues):
        pass

    # Let k be the order of the denominator polynomial (and, hence,
    # the system) and j be the order of the numerator polynomial.  We
    # need m = max(j+1,k) initial values. You need more than k initial
    # values if you have a deeper input dependence (say y[n] = x[n-5]).

    # Given initial values (the lists [x[1] ... x[j+1]] and [y[0]... y[k]]), 
    # return a state machine that transduces a stream of inputs X to
    # the stream of outputs Y.
    def stateMachine(self, initialXs, initialYs):
        return DifferenceEquationSMII(reverseCopy(self.denominator.coeffs),
                                      reverseCopy(self.numerator.coeffs),
                                      initialXs, initialYs)
    
    # Plot the transformed version of an input sequence in a window
    #  initialValues : list [y[0] ... y[k-1]]
    #  inputSource : n -> x[n]
    #  result is a plot
    # n is how many points to plot
    # set idle to True if you want the graph to play nice with idle;
    #  but you have to close the window before the program will
    #  continue
    def plotSequence(self, initialXs, initialYs, inputSource,
                     n = 30, idle = True):
        self.stateMachine(initialXs, initialYs).plot(inputSource, n, idle,
                                                     initialYs[:-1])

    def __str__(self):
        return 'SF(' + Polynomial.__str__(self.numerator, 'R') + \
               '/' + Polynomial.__str__(self.denominator, 'R') + ')'

    def __repr__(self):
        return self.__str__()

    # Given two system functions, return the system function of the
    # cascade
    def cascade(sf1, sf2):
        # Your code here
        pass

    # sf maps (X - W) into Y
    # return the sf for a machine that maps X into Y, where W = Y
    # Use Black's formula.  SF_new = sf / (1 + sf)
    # (n/d) / (1 + (n/d)) = (n/d) / ((d+n)/d) = n / (d+n)
    def feedback(sf):
        # your code here
        pass

    def scale(sf, v):
        # your code here
        pass

######################################################################
# Create a system function from a difference equation
######################################################################
    
# Assume that the difference equation has the form
# a_0 y[n] + a_1 y[n-1] + ... + a_k y[n-k] =
# b_0 y[n] + b_1 y[n-1] + ... + b_j y[n-j]
# It is crucial that the first coefficient of both polynomials be
# from the same time step;  but they can go back different depths
# in history.

# Now, the system function wants numerator and denominator polynomials
# in R.  Those polynomials have the form
#      a_0 Y + a_1 R Y + a_2 R^2 Y + ... + a_k R^k Y =             
#      b_0 X + b_1 R X + b_2 R^2 X + ... + b_j R^j X
# so we need to reverse the coefficient lists to make the necessary
# polynomials.           

def systemFunctionFromDifferenceEquation(outputCoeffs, inputCoeffs):
    return SystemFunction(Polynomial(reverseCopy(inputCoeffs)),
                          Polynomial(reverseCopy(outputCoeffs)))


######################################################################
## Utilities
######################################################################

# Reverse a list non-destructively
def reverseCopy(items):
    itemCopy = items[:]
    itemCopy.reverse()
    return itemCopy

# Functional push an item onto a fixed size queue
def fpush(item, queue):
    return ([item] + queue)[:-1]

# Return the dot product of two lists of numbers
def dotProd(a, b):
    return sum([ai*bi for (ai,bi) in zip(a,b)])

######################################################################
## Test
######################################################################

# Robot moving toward wall
# d[t] = d[t-1] - k*deltaT (d[t-2] - dDes)
# d[t] - d[t-1] + k*deltaT d[t-2] =  k*deltaT*dDes x[t-1]
#       where x[t] = 1 for all t
k = 1.0
deltaT = 0.2
dDes = 1.0
def robotWallSF(k):
    return systemFunctionFromDifferenceEquation([1, -1, k * deltaT],
                                                [k * deltaT * dDes])

def robotGainTest():
    return [(k, max([abs(r) for r in robotWallSF(k).poles()])) \
            for k in range(-10, 10)]

# Delay by 1
# y[t] = x[t-1]
d1sf = systemFunctionFromDifferenceEquation([1], [0, 1])

# Delay by 5
d5sf = systemFunctionFromDifferenceEquation([1], [0, 0, 0, 0, 0, 1])

# Fib
fib = systemFunctionFromDifferenceEquation([1,-1, -1],[])
fibsm = fib.stateMachine([],[1, 1])

# D as a function of V
# D[t] = D[t-1] - deltaT V[t-1]
wall = systemFunctionFromDifferenceEquation([1, -1], [0, -deltaT])

# V as a function of E:  V[t] = -k E[t-1]
def robot(k):
    return systemFunctionFromDifferenceEquation([1],[0, -k])

# E[t] = DDes[t] - D[t]

# We use feedback.  Input of the cascade machine is E, output is D,
# input of the feedback machine is -DDes
def robotWall(k):
    return wall.cascade(robot(k)).feedback()