# stochsys_test.py - test stochastic system operations
# RMM, 16 Mar 2022

import numpy as np
import pytest
from control.tests.conftest import asmatarrayout

import control as ct
import control.optimal as opt
from control import lqe, dlqe, rss, drss, tf, ss, ControlArgument, slycot_check
from math import log, pi

# Utility function to check LQE answer
def check_LQE(L, P, poles, G, QN, RN):
    P_expected = asmatarrayout(np.sqrt(G @ QN @ G @ RN))
    L_expected = asmatarrayout(P_expected / RN)
    poles_expected = -np.squeeze(np.asarray(L_expected))
    np.testing.assert_almost_equal(P, P_expected)
    np.testing.assert_almost_equal(L, L_expected)
    np.testing.assert_almost_equal(poles, poles_expected)

# Utility function to check discrete LQE solutions
def check_DLQE(L, P, poles, G, QN, RN):
    P_expected = asmatarrayout(G.dot(QN).dot(G))
    L_expected = asmatarrayout(0)
    poles_expected = -np.squeeze(np.asarray(L_expected))
    np.testing.assert_almost_equal(P, P_expected)
    np.testing.assert_almost_equal(L, L_expected)
    np.testing.assert_almost_equal(poles, poles_expected)

@pytest.mark.parametrize("method", [None, 'slycot', 'scipy'])
def test_LQE(matarrayin, method):
    if method == 'slycot' and not slycot_check():
        return

    A, G, C, QN, RN = (matarrayin([[X]]) for X in [0., .1, 1., 10., 2.])
    L, P, poles = lqe(A, G, C, QN, RN, method=method)
    check_LQE(L, P, poles, G, QN, RN)

@pytest.mark.parametrize("cdlqe", [lqe, dlqe])
def test_lqe_call_format(cdlqe):
    # Create a random state space system for testing
    sys = rss(4, 3, 2)
    sys.dt = None           # treat as either continuous or discrete time

    # Covariance matrices
    Q = np.eye(sys.ninputs)
    R = np.eye(sys.noutputs)
    N = np.zeros((sys.ninputs, sys.noutputs))

    # Standard calling format
    Lref, Pref, Eref = cdlqe(sys.A, sys.B, sys.C, Q, R)

    # Call with system instead of matricees
    L, P, E = cdlqe(sys, Q, R)
    np.testing.assert_almost_equal(Lref, L)
    np.testing.assert_almost_equal(Pref, P)
    np.testing.assert_almost_equal(Eref, E)

    # Make sure we get an error if we specify N
    with pytest.raises(ct.ControlNotImplemented):
        L, P, E = cdlqe(sys, Q, R, N)

    # Inconsistent system dimensions
    with pytest.raises(ct.ControlDimension, match="Incompatible"):
        L, P, E = cdlqe(sys.A, sys.C, sys.B, Q, R)

    # Incorrect covariance matrix dimensions
    with pytest.raises(ct.ControlDimension, match="Incompatible"):
        L, P, E = cdlqe(sys.A, sys.B, sys.C, R, Q)

    # Too few input arguments
    with pytest.raises(ct.ControlArgument, match="not enough input"):
        L, P, E = cdlqe(sys.A, sys.C)

    # First argument is the wrong type (use SISO for non-slycot tests)
    sys_tf = tf(rss(3, 1, 1))
    sys_tf.dt = None        # treat as either continuous or discrete time
    with pytest.raises(ct.ControlArgument, match="LTI system must be"):
        L, P, E = cdlqe(sys_tf, Q, R)

@pytest.mark.parametrize("method", [None, 'slycot', 'scipy'])
def test_DLQE(matarrayin, method):
    if method == 'slycot' and not slycot_check():
        return

    A, G, C, QN, RN = (matarrayin([[X]]) for X in [0., .1, 1., 10., 2.])
    L, P, poles = dlqe(A, G, C, QN, RN, method=method)
    check_DLQE(L, P, poles, G, QN, RN)

def test_lqe_discrete():
    """Test overloading of lqe operator for discrete time systems"""
    csys = ct.rss(2, 1, 1)
    dsys = ct.drss(2, 1, 1)
    Q = np.eye(1)
    R = np.eye(1)

    # Calling with a system versus explicit A, B should be the sam
    K_csys, S_csys, E_csys = ct.lqe(csys, Q, R)
    K_expl, S_expl, E_expl = ct.lqe(csys.A, csys.B, csys.C, Q, R)
    np.testing.assert_almost_equal(K_csys, K_expl)
    np.testing.assert_almost_equal(S_csys, S_expl)
    np.testing.assert_almost_equal(E_csys, E_expl)

    # Calling lqe() with a discrete time system should call dlqe()
    K_lqe, S_lqe, E_lqe = ct.lqe(dsys, Q, R)
    K_dlqe, S_dlqe, E_dlqe = ct.dlqe(dsys, Q, R)
    np.testing.assert_almost_equal(K_lqe, K_dlqe)
    np.testing.assert_almost_equal(S_lqe, S_dlqe)
    np.testing.assert_almost_equal(E_lqe, E_dlqe)

    # Calling lqe() with no timebase should call lqe()
    asys = ct.ss(csys.A, csys.B, csys.C, csys.D, dt=None)
    K_asys, S_asys, E_asys = ct.lqe(asys, Q, R)
    K_expl, S_expl, E_expl = ct.lqe(csys.A, csys.B, csys.C, Q, R)
    np.testing.assert_almost_equal(K_asys, K_expl)
    np.testing.assert_almost_equal(S_asys, S_expl)
    np.testing.assert_almost_equal(E_asys, E_expl)

    # Calling dlqe() with a continuous time system should raise an error
    with pytest.raises(ControlArgument, match="called with a continuous"):
        K, S, E = ct.dlqe(csys, Q, R)

def test_estimator_iosys():
    sys = ct.drss(4, 2, 2, strictly_proper=True)

    Q, R = np.eye(sys.nstates), np.eye(sys.ninputs)
    K, _, _ = ct.dlqr(sys, Q, R)

    P0 = np.eye(sys.nstates)
    QN = np.eye(sys.ninputs)
    RN = np.eye(sys.noutputs)
    estim = ct.create_estimator_iosystem(sys, QN, RN, P0)

    ctrl, clsys = ct.create_statefbk_iosystem(sys, K, estimator=estim)

    # Extract the elements of the estimator
    est = estim.linearize(0, 0)
    Be1 = est.B[:sys.nstates, :sys.noutputs]
    Be2 = est.B[:sys.nstates, sys.noutputs:]
    A_clchk = np.block([
        [sys.A, -sys.B @ K],
        [Be1 @ sys.C, est.A[:sys.nstates, :sys.nstates] - Be2 @ K]
    ])
    B_clchk = np.block([
        [sys.B @ K, sys.B],
        [Be2 @ K, Be2]
    ])
    C_clchk = np.block([
        [sys.C, np.zeros((sys.noutputs, sys.nstates))],
        [np.zeros_like(K), -K]
    ])
    D_clchk = np.block([
        [np.zeros((sys.noutputs, sys.nstates + sys.ninputs))],
        [K, np.eye(sys.ninputs)]
    ])

    # Check to make sure everything matches
    cls = clsys.linearize(0, 0)
    nstates = sys.nstates
    np.testing.assert_almost_equal(cls.A[:2*nstates, :2*nstates], A_clchk)
    np.testing.assert_almost_equal(cls.B[:2*nstates, :], B_clchk)
    np.testing.assert_almost_equal(cls.C[:, :2*nstates], C_clchk)
    np.testing.assert_almost_equal(cls.D, D_clchk)


@pytest.mark.parametrize("sys_args", [
    ([[-1]], [[1]], [[1]], 0),                          # scalar system
    ([[-1, 0.1], [0, -2]], [[0], [1]], [[1, 0]], 0),    # SISO, 2 state
    ([[-1, 0.1], [0, -2]], [[1, 0], [0, 1]], [[1, 0]], 0),    # 2i, 1o, 2s
    ([[-1, 0.1, 0.1], [0, -2, 0], [0.1, 0, -3]],        # 2i, 2o, 3s
     [[1, 0], [0, 0.1], [0, 1]],
     [[1, 0, 0.1], [0, 1, 0.1]], 0),
])
def test_estimator_iosys_ctime(sys_args):
    # Define the system we want to test
    sys = ct.ss(*sys_args)
    T = 10 * log(1e-2) / np.max(sys.poles().real)
    assert T > 0

    # Create nonlinear version of the system to match integration methods
    nl_sys = ct.NonlinearIOSystem(
        lambda t, x, u, params : sys.A @ x + sys.B @ u,
        lambda t, x, u, params : sys.C @ x + sys.D @ u,
        inputs=sys.ninputs, outputs=sys.noutputs, states=sys.nstates)

    # Define an initial condition, inputs (small, to avoid integration errors)
    timepts = np.linspace(0, T, 500)
    U = 2e-2 * np.array([np.sin(timepts + i*pi/3) for i in range(sys.ninputs)])
    X0 = np.ones(sys.nstates)

    # Set up the parameters for the filter
    P0 = np.eye(sys.nstates)
    QN = np.eye(sys.ninputs)
    RN = np.eye(sys.noutputs)

    # Construct the estimator
    estim = ct.create_estimator_iosystem(sys, QN, RN)

    # Compute the system response and the optimal covariance
    sys_resp = ct.input_output_response(nl_sys, timepts, U, X0)
    _, Pf, _ = ct.lqe(sys, QN, RN)
    Pf = np.array(Pf)           # convert from matrix, if needed

    # Make sure that we converge to the optimal estimate
    estim_resp = ct.input_output_response(
        estim, timepts, [sys_resp.outputs, U], [0*X0, P0])
    np.testing.assert_allclose(
        estim_resp.states[0:sys.nstates, -1], sys_resp.states[:, -1],
        atol=1e-6, rtol=1e-3)
    np.testing.assert_allclose(
        estim_resp.states[sys.nstates:, -1], Pf.reshape(-1),
        atol=1e-6, rtol=1e-3)

    # Make sure that optimal estimate is an eq pt
    ss_resp = ct.input_output_response(
        estim, timepts, [sys_resp.outputs, U], [X0, Pf])
    np.testing.assert_allclose(
        ss_resp.states[sys.nstates:],
        np.outer(Pf.reshape(-1), np.ones_like(timepts)),
        atol=1e-4, rtol=1e-2)
    np.testing.assert_allclose(
        ss_resp.states[0:sys.nstates], sys_resp.states,
        atol=1e-4, rtol=1e-2)


def test_estimator_errors():
    sys = ct.drss(4, 2, 2, strictly_proper=True)
    P0 = np.eye(sys.nstates)
    QN = np.eye(sys.ninputs)
    RN = np.eye(sys.noutputs)

    with pytest.raises(TypeError, match="unrecognized keyword"):
        estim = ct.create_estimator_iosystem(sys, QN, RN, unknown=True)

    with pytest.raises(ct.ControlArgument, match=".* system must be a linear"):
        sys_tf = ct.tf([1], [1, 1], dt=True)
        estim = ct.create_estimator_iosystem(sys_tf, QN, RN)

    with pytest.raises(ValueError, match="output must be full state"):
        C = np.eye(2, 4)
        estim = ct.create_estimator_iosystem(sys, QN, RN, C=C)

    with pytest.raises(ValueError, match="output is the wrong size"):
        sys_fs = ct.drss(4, 4, 2, strictly_proper=True)
        sys_fs.C = np.eye(4)
        C = np.eye(1, 4)
        estim = ct.create_estimator_iosystem(sys_fs, QN, RN, C=C)


def test_white_noise():
    # Scalar white noise signal
    T = np.linspace(0, 1000, 1000)
    R = 0.5
    V = ct.white_noise(T, R)
    assert abs(np.mean(V)) < 0.1                # can occassionally fail
    assert abs(np.cov(V) - 0.5) < 0.1           # can occassionally fail

    # Vector white noise signal
    R = [[0.5, 0], [0, 0.1]]
    V = ct.white_noise(T, R)
    assert abs(np.mean(V)) < 0.1                # can occassionally fail
    assert np.all(abs(np.cov(V) - R) < 0.1)     # can occassionally fail

    # Make sure time scaling works properly
    T = T / 10
    V = ct.white_noise(T, R)
    assert abs(np.mean(V)) < np.sqrt(10)        # can occassionally fail
    assert np.all(abs(np.cov(V) - R) < 10)      # can occassionally fail

    # Make sure discrete time works properly
    V = ct.white_noise(T, R, dt=T[1] - T[0])
    assert abs(np.mean(V)) < 0.1                # can occassionally fail
    assert np.all(abs(np.cov(V) - R) < 0.1)     # can occassionally fail

    # Test error conditions
    with pytest.raises(ValueError, match="T must be 1D"):
        V = ct.white_noise(R, R)

    with pytest.raises(ValueError, match="Q must be square"):
        R = np.outer(np.eye(2, 3), np.ones_like(T))
        V = ct.white_noise(T, R)

    with pytest.raises(ValueError, match="Time values must be equally"):
        T = np.logspace(0, 2, 100)
        R = [[0.5, 0], [0, 0.1]]
        V = ct.white_noise(T, R)


def test_correlation():
    # Create an uncorrelated random sigmal
    T = np.linspace(0, 1000, 1000)
    R = 0.5
    V = ct.white_noise(T, R)

    # Compute the correlation
    tau, Rtau = ct.correlation(T, V)

    # Make sure the correlation makes sense
    zero_index = np.where(tau == 0)
    np.testing.assert_almost_equal(Rtau[zero_index], np.cov(V), decimal=2)
    for i, t in enumerate(tau):
        if i == zero_index:
            continue
    assert abs(Rtau[i]) < 0.01

    # Try passing a second argument
    tau, Rneg = ct.correlation(T, V, -V)
    np.testing.assert_equal(Rtau, -Rneg)

    # Test error conditions
    with pytest.raises(ValueError, match="Time vector T must be 1D"):
        tau, Rtau = ct.correlation(V, V)

    with pytest.raises(ValueError, match="X and Y must be 2D"):
        tau, Rtau = ct.correlation(T, np.zeros((3, T.size, 2)))

    with pytest.raises(ValueError, match="X and Y must have same length as T"):
        tau, Rtau = ct.correlation(T, V[:, 0:-1])

    with pytest.raises(ValueError, match="Time values must be equally"):
        T = np.logspace(0, 2, T.size)
        tau, Rtau = ct.correlation(T, V)

@pytest.mark.slow
@pytest.mark.parametrize('dt', [0, 0.2])
def test_oep(dt):
    # Define the system to test, with additional input
    # Use fixed system to avoid random errors (was csys = ct.rss(4, 2, 5))
    csys = ct.ss(
        [[-0.5, 1, 0, 0], [0, -1, 1, 0], [0, 0, -2, 1], [0, 0, 0, -3]], # A
        [[0, 0.1], [0, 0.1], [0, 0.1], [1, 0.1]],                       # B
        [[1, 0, 0, 0], [0, 0, 1, 0]],                                   # C
        0, dt=0)
    dsys = ct.c2d(csys, dt)
    sys = csys if dt == 0 else dsys

    # Create disturbances and noise (fixed, to avoid random errors)
    dist_mag = 1e-1                     # disturbance magnitude
    meas_mag = 1e-3                     # measurement noise magnitude
    Rv = dist_mag**2 * np.eye(1)        # scalar disturbance
    Rw = meas_mag**2 * np.eye(sys.noutputs)
    timepts = np.arange(0, 1, 0.2)
    V = np.array(
        [0 if i % 2 == 1 else 1 if i % 4 == 0 else -1
         for i in range(timepts.size)]
    ).reshape(1, -1) * dist_mag / 10
    W = np.vstack([
        np.sin(10*timepts/timepts[-1]), np.cos(15*timepts)/timepts[-1]
    ]) * meas_mag / 10

    # Generate system data
    U = np.sin(timepts).reshape(1, -1)

    # With disturbances and noise
    res = ct.input_output_response(sys, timepts, [U, V])
    Y = res.outputs + W

    # Set up optimal estimation function using Gaussian likelihoods for cost
    traj_cost = opt.gaussian_likelihood_cost(sys, Rv, Rw)
    init_cost = lambda xhat, x: (xhat - x) @ (xhat - x)
    oep1 = opt.OptimalEstimationProblem(
        sys, timepts, traj_cost, terminal_cost=init_cost)

    # Compute the optimal estimate
    est1 = oep1.compute_estimate(Y, U)
    assert est1.success
    np.testing.assert_allclose(
        est1.states[:, -1], res.states[:, -1], atol=meas_mag, rtol=meas_mag)

    # Recompute using initial guess (should be pretty fast)
    est2 = oep1.compute_estimate(
        Y, U, initial_guess=(est1.states, est1.inputs))
    assert est2.success

    # Change around the inputs and disturbances
    sys2 = ct.ss(sys.A, sys.B[:, ::-1], sys.C, sys.D[::-1], sys.dt)
    oep2a = opt.OptimalEstimationProblem(
        sys2, timepts, traj_cost, terminal_cost=init_cost,
        control_indices=[1])
    est2a = oep2a.compute_estimate(
        Y, U, initial_guess=(est1.states, est1.inputs))
    np.testing.assert_allclose(est2a.states, est2.states)

    oep2b = opt.OptimalEstimationProblem(
        sys2, timepts, traj_cost, terminal_cost=init_cost,
        disturbance_indices=[0])
    est2b = oep2b.compute_estimate(
        Y, U, initial_guess=(est1.states, est1.inputs))
    np.testing.assert_allclose(est2b.states, est2.states)

    # Add disturbance constraints
    V3 = np.clip(V, 0.5, 1)
    traj_constraint = opt.disturbance_range_constraint(sys, 0.5, 1)
    oep3 = opt.OptimalEstimationProblem(
        sys, timepts, traj_cost, terminal_cost=init_cost,
        trajectory_constraints=traj_constraint)

    res3 = ct.input_output_response(sys, timepts, [U, V3])
    Y3 = res3.outputs + W

    # Make sure estimation is correct with constraint in place
    est3 = oep3.compute_estimate(Y3, U)
    assert est3.success
    np.testing.assert_allclose(
        est3.states[:, -1], res3.states[:, -1], atol=meas_mag, rtol=meas_mag)


@pytest.mark.slow
def test_mhe():
    # Define the system to test, with additional input
    csys = ct.ss(
        [[-0.5, 1, 0, 0], [0, -1, 1, 0], [0, 0, -2, 1], [0, 0, 0, -3]], # A
        [[0, 0.1], [0, 0.1], [0, 0.1], [1, 0.1]],                       # B
        [[1, 0, 0, 0], [0, 0, 1, 0]],                                   # C
        0, dt=0)
    dt = 0.1
    sys = ct.c2d(csys, dt)

    # Create disturbances and noise (fixed, to avoid random errors)
    Rv = 0.1 * np.eye(1)                # scalar disturbance
    Rw = 1e-6 * np.eye(sys.noutputs)
    P0 = 0.1 * np.eye(sys.nstates)

    timepts = np.arange(0, 10*dt, dt)
    mhe_timepts = np.arange(0, 5*dt, dt)
    V = np.array(
        [0 if i % 2 == 1 else 1 if i % 4 == 0 else -1
         for i, t in enumerate(timepts)]).reshape(1, -1) * 0.1
    W = np.sin(timepts / dt) * 1e-3

    # Create a moving horizon estimator
    traj_cost = opt.gaussian_likelihood_cost(sys, Rv, Rw)
    init_cost = lambda xhat, x: (xhat - x) @ P0 @ (xhat - x)
    oep = opt.OptimalEstimationProblem(
        sys, mhe_timepts, traj_cost, terminal_cost=init_cost,
        disturbance_indices=1)
    mhe = oep.create_mhe_iosystem()

    # Generate system data
    U = 10 * np.sin(timepts / (4*dt))
    inputs = np.vstack([U, V])
    resp = ct.input_output_response(sys, timepts, inputs)

    # Run the estimator
    estp = ct.input_output_response(
        mhe, timepts, [resp.outputs, resp.inputs[0:1]])

    # Make sure the estimated state is close to the actual state
    np.testing.assert_allclose(estp.outputs, resp.states, atol=1e-2, rtol=1e-4)

@pytest.mark.slow
@pytest.mark.parametrize("ctrl_indices, dist_indices", [
    (slice(0, 3), None),
    (3, None),
    (None, 2),
    ([0, 1, 4], None),
    (['u[0]', 'u[1]', 'u[4]'], None),
    (['u[0]', 'u[1]', 'u[4]'], ['u[1]', 'u[3]']),
    (slice(0, 3), slice(3, 5))
])
def test_indices(ctrl_indices, dist_indices):
    # Define a system with inputs (0:3), disturbances (3:5), and no noise
    sys = ct.ss(
        [[-1, 1, 0, 0], [0, -2, 1, 0], [0, 0, -3, 1], [0, 0, 0, -4]],
        [[0, 0, 0, 0, 0], [1, 0, 0, 0, 0], [0, 1, 0, .1, 0], [0, 0, 1, 0, .1]],
        [[1, 0, 0, 0], [0, 1, 0, 0]], 0)

    # Create a system whose state we want to estimate
    if ctrl_indices is not None:
        ctrl_idx = ct.namedio._process_indices(
            ctrl_indices, 'control', sys.input_labels, sys.ninputs)
        dist_idx = [i for i in range(sys.ninputs) if i not in ctrl_idx]
    else:
        arg = -dist_indices if isinstance(dist_indices, int) else dist_indices
        dist_idx = ct.namedio._process_indices(
            arg, 'disturbance', sys.input_labels, sys.ninputs)
        ctrl_idx = [i for i in range(sys.ninputs) if i not in dist_idx]
    sysm = ct.ss(sys.A, sys.B[:, ctrl_idx], sys.C, sys.D[:, ctrl_idx])

    # Set the simulation time based on the slowest system pole
    from math import log
    T = 10

    # Generate a system response with no disturbances
    timepts = np.linspace(0, T, 20)
    U = np.vstack([np.sin(timepts + i) for i in range(len(ctrl_idx))])
    resp = ct.input_output_response(
        sysm, timepts, U, np.zeros(sys.nstates),
        solve_ivp_kwargs={'method': 'RK45', 'max_step': 0.01,
                          'atol': 1, 'rtol': 1})
    Y = resp.outputs

    # Create an estimator
    QN = np.eye(len(dist_idx))
    RN = np.eye(sys.noutputs)
    P0 = np.eye(sys.nstates)
    estim = ct.create_estimator_iosystem(
        sys, QN, RN, control_indices=ctrl_indices,
        disturbance_indices=dist_indices)

    # Run estimator (no prediction + same solve_ivp params => should be exact)
    resp_estim = ct.input_output_response(
        estim, timepts, [Y, U], [np.zeros(sys.nstates), P0],
        solve_ivp_kwargs={'method': 'RK45', 'max_step': 0.01,
                          'atol': 1, 'rtol': 1},
            params={'correct': False})
    np.testing.assert_allclose(resp.states, resp_estim.outputs, rtol=1e-2)
