from controlboros import StateSpaceBuilder, RateWrapper
import matplotlib.pyplot as plt
import numpy as np
from scipy import signal

%matplotlib inline
%config InlineBackend.figure_format = 'retina'

"""
Explanation: Simple multi-rate simulation
Mikhail Pak, 2017
End of explanation
"""


t_begin, t_end = 0.0, 1.0

"""
Explanation: In this notebook, we shall simulate the step response of a simple system using scipy.signal and controlboros.
The system is defined by the following transfer function:
$$H(s) = \frac{1}{s + 1}$$
We want to simulate from 0 to 1 s:
End of explanation
"""


dt = 1.0e-3

"""
Explanation: We use 1 ms sample time for the reference simulation performed with scipy.signal:
End of explanation
"""


multiplier = 100

"""
Explanation: As for controlboros, we shall discretise $H(s)$ with 0.1 ms sample rate but simulate it with 1 ms sample rate. We define the corresponding rate multiplier:
End of explanation
"""


t_ref, y_ref = signal.step(
    ([1.0], [1.0, 1.0]),
    T=np.arange(t_begin, t_end, dt),
    )

"""
Explanation: Ok. Now we create a reference system using scipy.signal and compute its step response:
End of explanation
"""


sys_cb = StateSpaceBuilder().from_tf([1.0], [1.0, 1.0])\
                            .discretise(dt*multiplier)\
                            .build()
sys_cb_wrapped = RateWrapper(sys_cb, multiplier)

"""
Explanation: Now we create a controlboros.StateSpace model using the builder pattern and wrap it using controlboros.RateWrapper:
End of explanation
"""


excitation = np.array([1.0])  # Constant step excitation
t_cb = np.arange(t_begin, t_end, dt)
y_cb = np.zeros((len(t_cb),))

"""
Explanation: And we initialise arrays for the step response:
End of explanation
"""


# Reset the inital state of the systems,
# helpful if you run this cell multiple times!
sys_cb_wrapped.set_state_to_zero()
sys_cb_wrapped.reset_wrapper()

for i in range(len(t_cb)):
    y_cb[i] = sys_cb_wrapped.push_stateful(excitation)

"""
Explanation: We're ready to run the main loop:
End of explanation
"""


plt.figure(figsize=(7, 5))
plt.plot(t_ref, y_ref)
plt.step(t_cb, y_cb, where="post")
plt.xlabel("Time (s)")
plt.ylabel("Output")
plt.legend(["scipy.signal", "controlboros"])
plt.grid()
plt.show()

"""
Explanation: Plot and compare results:
End of explanation
"""


