import numpy as np
import matplotlib.pyplot as plt
from scipy.integrate import solve_ivp

# Определим функцию для системы дифференциальных уравнений
def system_equations(gamma):
    def rhs(t, X):
        x, y = X
        dxdt = y
        dydt = +x**5 - 5*x**3 + 4*x
        return [dxdt, dydt]
    return rhs

# Создадим функцию для построения векторного поля
def plot_vector_field(rhs, limits, N=25, arrow_length=0.1):
    xlims, ylims = limits
    xs = np.linspace(xlims[0], xlims[1], N)
    ys = np.linspace(ylims[0], ylims[1], N)
    U = np.zeros((N, N))
    V = np.zeros((N, N))
    for i, y in enumerate(ys):
        for j, x in enumerate(xs):
            vfield = rhs(0.0, [x, y])
            u, v = vfield
            # Нормализуем вектор
            length = np.sqrt(u**2 + v**2)
            if length != 0:
                u /= length
                v /= length
            # Умножаем на фиксированную длину
            u *= arrow_length
            v *= arrow_length
            U[i][j] = u
            V[i][j] = v
    return xs, ys, U, V


# Функция для построения векторного поля на плоскости
def plot_on_plane(rhs, limits):
    plt.close()
    xlims, ylims = limits
    plt.xlim(xlims[0], xlims[1])
    plt.ylim(ylims[0], ylims[1])
    xs, ys, U, V = plot_vector_field(rhs, limits)
    plt.quiver(xs, ys, U, V, alpha=0.8)

# Определение параметра gamma и создание функции для системы с этим параметром
gamma = 1.0
rhs = system_equations(gamma)

# Построение векторного поля на плоскости
plot_on_plane(rhs, [(-3.0, 3.0), (-3.0, 3.0)])

# Решение системы дифференциальных уравнений и построение траекторий
sol1 = solve_ivp(rhs, [0.0, 100.0], (0.01, 0.02), method='RK45', rtol=1e-12)
x1, y1 = sol1.y
# plt.plot(x1, y1, 'b-')

plt.scatter(x1, y1, color='yellow', marker='.')

sol2 = solve_ivp(rhs, [0.0, 100.0], (2.0, 0.5), method='RK45', rtol=1e-12)
x2, y2 = sol2.y
plt.scatter(x2, y2, color='red', marker='.')

sol3 = solve_ivp(rhs, [0.0, 100.0], (-0.01, -0.02), method='RK45', rtol=1e-12)
x3, y3 = sol3.y
plt.scatter(x3, y3, color='yellow', marker='.')

sol4 = solve_ivp(rhs, [0.0, -100.0], (2.0, 0.5), method='RK45', rtol=1e-12)
x4, y4 = sol4.y
plt.scatter(x4, y4, color='red', marker='.')

sol5 = solve_ivp(rhs, [0.0, 100.0], (2.0, -0.5), method='RK45', rtol=1e-12)
x5, y5 = sol5.y
plt.scatter(x5, y5, color='red', marker='.')

sol6 = solve_ivp(rhs, [0.0, -100.0], (2.0, -0.5), method='RK45', rtol=1e-12)
x6, y6 = sol6.y
plt.scatter(x6, y6, color='red', marker='.')

sol7 = solve_ivp(rhs, [0.0, 100.0], (2.2, 0.0), method='RK45', rtol=1e-12)
x7, y7 = sol7.y
plt.scatter(x7, y7, color='red', marker='.')

sol8 = solve_ivp(rhs, [0.0, -100.0], (-2.2, 0.0), method='RK45', rtol=1e-12)
x8, y8 = sol8.y
plt.scatter(x8, y8, color='red', marker='.')

sol9 = solve_ivp(rhs, [0.0, -100.0], (2.2, 0.0), method='RK45', rtol=1e-12)
x9, y9 = sol9.y
plt.scatter(x9, y9, color='red', marker='.')

sol10 = solve_ivp(rhs, [0.0, 100.0], (-2.2, 0.0), method='RK45', rtol=1e-12)
x10, y10 = sol10.y
plt.scatter(x10, y10, color='red', marker='.')

sol11 = solve_ivp(rhs, [0.0, 100.0], (0.0, 1.0), method='RK45', rtol=1e-12)
x11, y11 = sol11.y
plt.scatter(x11, y11, color='green', marker='.')

sol12 = solve_ivp(rhs, [0.0, 100.0], (1.0, 0.0), method='RK45', rtol=1e-12)
x12, y12 = sol12.y
plt.scatter(x12, y12, color='blue', marker='0')

sol13 = solve_ivp(rhs, [0.0, 100.0], (2.01, np.sqrt(24)*0.01), method='RK45', rtol=1e-12)
x13, y13 = sol13.y
plt.scatter(x13, y13, color='yellow', marker='.')

sol13 = solve_ivp(rhs, [0.0, -100.0], (2.01, np.sqrt(24)*0.01), method='RK45', rtol=1e-12)
x13, y13 = sol13.y
plt.scatter(x13, y13, color='yellow', marker='.')

sol14 = solve_ivp(rhs, [0.0, -100.0], (1.99, np.sqrt(24)*0.01), method='RK45', rtol=1e-12)
x14,y14 = sol14.y
plt.scatter(x14, y14, color='yellow', marker='.')

sol15 = solve_ivp(rhs, [0.0, 100.0], (1.99, -np.sqrt(24)*0.01), method='RK45', rtol=1e-12)
x15,y15 = sol15.y
plt.scatter(x15, y15, color='yellow', marker='.')

sol16 = solve_ivp(rhs, [0.0, 100.0], (-1.0, 0.0), method='RK45', rtol=1e-12)
x16, y16 = sol16.y
plt.scatter(x16, y16, color='blue', marker='+')

sol17 = solve_ivp(rhs, [0.0, 100.0], (0.0, 0.0), method='RK45', rtol=1e-12)
x17, y17 = sol17.y
plt.scatter(x17, y17, color='red', marker='+')

sol17 = solve_ivp(rhs, [0.0, 100.0], (-2.0, 0.0), method='RK45', rtol=1e-12)
x17, y17 = sol17.y
plt.scatter(x17, y17, color='red', marker='+')

sol17 = solve_ivp(rhs, [0.0, 100.0], (2.0, 0.0), method='RK45', rtol=1e-12)
x17, y17 = sol17.y
plt.scatter(x17, y17, color='red', marker='+')

sol18 = solve_ivp(rhs, [0.0, 100.0], (0.5, 0.0), method='RK45', rtol=1e-12)
x18, y18 = sol18.y
plt.scatter(x18, y18, color='green', marker='.')

sol19 = solve_ivp(rhs, [0.0, 100.0], (-0.5, 0.0), method='RK45', rtol=1e-12)
x19, y19 = sol19.y
plt.scatter(x19, y19, color='green', marker='.')

# Построение графика траектории
# x1, y1 = sol1.y
# plt.plot(x1, y1, 'k-')

# Отображение результатов
plt.show()