import matplotlib.pyplot as plt
import numpy as np

def y(u, alpha, beta):
  return alpha * (u + beta * u ** 3)


t = np.arange(0, 5, 0.01)
ut = np.sin(t) + np.sin(np.pi * t) + np.sin((np.pi**2) * t)

fig, [ax1, ax2] = plt.subplots(1,2, figsize=(6, 3))

for a in np.arange(0.1, 0.6, 0.2):
  for b in np.arange(0, 2.5, 1):
    yt = y(ut, a, b)
    ax1.plot(ut, yt, 'r');
    ax2.plot(t, yt, 'r')

yt = y(ut, 0.2, 1)
ax1.plot(ut, yt, 'b--');
ax2.plot(t, yt, 'b--')

ax1.set_title("I/O relationship")
ax1.set_xlabel("Input u")
ax1.set_ylabel("Output y")
ax1.set_ylim(-3, 3)
ax1.set_xlim(-2.5, 2.5)


ax2.set_title("Output signals")
ax2.set_xlabel("Time t")
ax2.set_ylabel("Output y")
ax2.set_ylim(-1, 5)
ax2.set_xlim(0, 2)
plt.tight_layout()
plt.show()




# closed loop system

dt = 0.0001
t = np.arange(0, 5, dt)
rt = np.sin(t) + np.sin(np.pi * t) + np.sin((np.pi**2) * t)

yt = np.zeros(len(t))
ut = np.zeros(len(t))
et = np.zeros(len(t))

K = 1000
tc = 0.3 # time constant of the amplifier
alpha = 0.2
beta = 1.0

for i in range(1, len(t)):
  et[i] = rt[i-1] - yt[i-1]
  ut[i] = ut[i-1] + (K * et[i] - ut[i-1]) * (dt/tc)
  yt[i] = y(ut[i], alpha, beta)

fig, [ax1, ax2, ax3, ax4] = plt.subplots(1,4, figsize=(12, 3))

ax1.set_title("I/O relationship")
ax1.plot(rt, yt, "red")
ax1.set_xlabel("Input r")
ax1.set_ylabel("Output y")

ax2.plot(ut, yt, "r")
ax2.set_xlabel("u")
ax2.set_ylabel("Output y")

ax3.plot(t, yt, "r")
ax3.set_xlabel("Time t")
ax3.set_ylabel("Output y")
ax3.set_xlim(0, 2)

ax4.plot(t, et, "r")
ax4.set_xlabel("Time t")
ax4.set_ylabel("Error e")
ax4.set_xlim(0, 2)

plt.tight_layout()
plt.show()