'''
Copyright (C) 2011 by Eike Welk.

Test the control.matlab toolbox.
'''

import unittest

import numpy as np
import scipy.signal
from numpy.testing import assert_array_almost_equal
from numpy import array, asarray, matrix, asmatrix, zeros, ones, linspace,\
                  all, hstack, vstack, c_, r_
from matplotlib.pylab import show, figure, plot, legend, subplot2grid
from control.matlab import ss, step, impulse, initial, lsim, dcgain, \
                           ss2tf
from control.statesp import _mimo2siso
from control.timeresp import _check_convert_array
from control.exception import slycot_check
import warnings

class TestControlMatlab(unittest.TestCase):

    def setUp(self):
        pass

    def plot_matrix(self):
        #Test: can matplotlib correctly plot matrices?
        #Yes, but slightly inconvenient
        figure()
        t = matrix([[ 1.],
                    [ 2.],
                    [ 3.],
                    [ 4.]])
        y = matrix([[ 1., 4.],
                    [ 4., 5.],
                    [ 9., 6.],
                    [16., 7.]])
        plot(t, y)
        #plot(asarray(t)[0], asarray(y)[0])


    def make_SISO_mats(self):
        """Return matrices for a SISO system"""
        A = matrix([[-81.82, -45.45],
                    [ 10.,    -1.  ]])
        B = matrix([[9.09],
                    [0.  ]])
        C = matrix([[0, 0.159]])
        D = zeros((1, 1))
        return A, B, C, D

    def make_MIMO_mats(self):
        """Return matrices for a MIMO system"""
        A = array([[-81.82, -45.45,   0,      0   ],
                   [ 10,     -1,      0,      0   ],
                   [  0,      0,    -81.82, -45.45],
                   [  0,      0,     10,     -1,  ]])
        B = array([[9.09, 0   ],
                   [0   , 0   ],
                   [0   , 9.09],
                   [0   , 0   ]])
        C = array([[0, 0.159, 0, 0    ],
                   [0, 0,     0, 0.159]])
        D = zeros((2, 2))
        return A, B, C, D

    def test_dcgain(self):
        """Test function dcgain with different systems"""
        if slycot_check():
            #Test MIMO systems
            A, B, C, D = self.make_MIMO_mats()

            gain1 = dcgain(ss(A, B, C, D))
            gain2 = dcgain(A, B, C, D)
            sys_tf = ss2tf(A, B, C, D)
            gain3 = dcgain(sys_tf)
            gain4 = dcgain(sys_tf.num, sys_tf.den)
            #print("gain1:", gain1)

            assert_array_almost_equal(gain1,
                                      array([[0.0269, 0.    ],
                                             [0.    , 0.0269]]),
                                      decimal=4)
            assert_array_almost_equal(gain1, gain2)
            assert_array_almost_equal(gain3, gain4)
            assert_array_almost_equal(gain1, gain4)

        #Test SISO systems
        A, B, C, D = self.make_SISO_mats()

        gain1 = dcgain(ss(A, B, C, D))
        assert_array_almost_equal(gain1,
                                  array([[0.0269]]),
                                  decimal=4)

    def test_dcgain_2(self):
        """Test function dcgain with different systems"""
        #Create different forms of a SISO system
        A, B, C, D = self.make_SISO_mats()
        num, den = scipy.signal.ss2tf(A, B, C, D)
        # numerator is only a constant here; pick it out to avoid numpy warning
        Z, P, k = scipy.signal.tf2zpk(num[0][-1], den)
        sys_ss = ss(A, B, C, D)

        #Compute the gain with ``dcgain``
        gain_abcd = dcgain(A, B, C, D)
        gain_zpk = dcgain(Z, P, k)
        gain_numden = dcgain(np.squeeze(num), den)
        gain_sys_ss = dcgain(sys_ss)
        # print('gain_abcd:', gain_abcd, 'gain_zpk:', gain_zpk)
        # print('gain_numden:', gain_numden, 'gain_sys_ss:', gain_sys_ss)

        #Compute the gain with a long simulation
        t = linspace(0, 1000, 1000)
        y, _t = step(sys_ss, t)
        gain_sim = y[-1]
        # print('gain_sim:', gain_sim)

        #All gain values must be approximately equal to the known gain
        assert_array_almost_equal([gain_abcd,   gain_zpk,
                                   gain_numden, gain_sys_ss, gain_sim],
                                  [0.026948, 0.026948, 0.026948, 0.026948,
                                   0.026948],