import matplotlib.pyplot as plt
import numpy as np
from scipy.optimize import least_squares, minimize


def calculate(args):
    T = 30
    h = 0.01
    a=0.7
    b=0.8
    c=3
    z=-0.4
    iterations = int(T / h)
    t_arr = np.arange(0, T, h)
    x_arr = np.array(t_arr)
    y_arr = np.array(t_arr)
    z_arr=np.array(t_arr)
    e_arr = np.array(t_arr)
    u_arr = np.array(t_arr)
    e_sum = 0
    x_arr[0] = 1
    y_arr[0] = 1
    y_need = 0
    e_arr[0] = 0
    u_arr[0] = 0
    set_up=2
    set_down=-1

    for i in range(iterations-1):
        e_arr[i + 1] = y_need - x_arr[i]
        e_sum += + e_arr[i] * h

        u_arr[i + 1] = args[0] * e_arr[i + 1] + args[1] * e_sum + args[2] * ((e_arr[i + 1] - e_arr[i]) / h)
        if u_arr[i+1]>set_up:
            u_arr[i+1]=set_up
        elif u_arr[i+1]<set_down:
            u_arr[i+1]=set_down

        x_arr[i + 1] = x_arr[i] + h * (9*u_arr[i+1]-(z_arr[i]*(a+9))-(x_arr[i]*(a+1)))
        y_arr[i+1] = y_arr[i]+ h * (z_arr[i])
        z_arr[i + 1] = z_arr[i] + h * (x_arr[i])

    return e_sum, t_arr, x_arr, y_arr,z_arr, e_arr, u_arr


def optimize(args):
    e_sum, a, b, c,d, e_arr, e = calculate(args)
    lam=1.0
    return np.sum(abs(e_arr))+lam*(args[0]**2+args[1]**2+args[2]**2)


args = minimize(optimize, np.array([0.0, 0.0, 0.0], dtype=float), method='BFGS')
print(args)
e_sum, t_arr, x_arr, y_arr,z_arr, e_arr, u_arr = calculate(args.x)

plt.subplot(311)
plt.grid(True)
plt.plot(t_arr, x_arr, label="x")
plt.plot(t_arr, y_arr, label="y")
plt.legend(loc="upper right")
plt.ylim()
plt.ylabel("stan")

plt.subplot(312)
plt.grid(True)
plt.plot(t_arr, e_arr)

plt.ylabel("blad")

plt.subplot(313)
plt.grid(True)
plt.plot(t_arr, u_arr, label="wejscie")
plt.ylabel("u")

plt.show()