import numpy as np
from scipy.integrate import odeint
import matplotlib.pyplot as plt
from math import pi, exp, fabs

a = 0
b = pi / 2
n = 100
alpha = 10
epsilon = 0.001
delta = 0.1

psi1_0 = 0
psi2_0 = 1  # произвольно
x1_0 = 1  # произвольно
x2_0 = 0

t = np.linspace(a, b, n)  # vector of time


# create function
def f(y, t):
    x1, x2, psi1, psi2 = y
    return [x2, psi2 / 2 - x1 * exp(-alpha * t), psi2 * exp(-alpha * t), -psi1]


def ode_solve(x1_init, psi2_init):
    y0 = [x1_init, x2_0, psi1_0, psi2_init]  # start value
    w = odeint(f, y0, t)  # solve eq.
    x1 = w[:, 0]
    x2 = w[:, 1]
    psi1 = w[:, 2]
    psi2 = w[:, 3]
    return x1, x2, psi1, psi2


def newton():
    global x1_0, psi2_0
    x1, x2, psi1, psi2 = ode_solve(x1_0 + delta, psi2_0)
    delta_11 = psi1[-1]
    delta_21 = psi2[-1]
    x1, x2, psi1, psi2 = ode_solve(x1_0, psi2_0 + delta)
    delta_12 = psi1[-1]
    delta_22 = psi2[-1]

    a11 = (delta_11 - delta_10) / delta
    a12 = (delta_12 - delta_10) / delta
    a21 = (delta_21 - delta_20) / delta
    a22 = (delta_22 - delta_20) / delta

    b1 = a11 * psi2_0 + a12 * x1_0 - delta_10
    b2 = a21 * psi2_0 + a22 * x1_0 - delta_20

    psi2_0 = (a11 * b2 - a21 * b1) / (a22 * a11 - a12 * a21)
    x1_0 = (b1 * a22 - b2 * a12) / (a22 * a11 - a12 * a21)

    return ode_solve(x1_0, psi2_0)


if __name__ == '__main__':
    x1, x2, psi1, psi2 = ode_solve(x1_0, psi2_0)
    delta_10 = psi1[-1]
    delta_20 = psi2[-1]
    iteration_count = 1
    while True:
        print('Итерация: ', iteration_count)
        x1, x2, psi1, psi2 = newton()
        u = [point / 2 for point in psi2]  # u = psi2 / 2
        delta_10 = psi1[-1]
        delta_20 = psi2[-1]
        if fabs(delta_10) < epsilon and fabs(delta_20) < epsilon:
            break
        iteration_count += 1

    plt.plot(t, u, label='u')  # plot u
    # plt.plot(t, x1, label='x1')
    # plt.plot(t, x2, label='x2')
    plt.plot(t, psi1, label='psi1')
    plt.plot(t, psi2, label='psi2')  # plot x1, x2, psi1, psi2
    plt.grid(True)
    plt.legend()
    plt.show()  # display