import numpy as np
import math
import pylab as pl

h = 0.01
n = int(1 / h)
f = lambda x, y: -50 * y + 50 * x ** 2 + 2 * x


def E(f, x0, y0, h, N):
    x = np.zeros(N + 1)
    y = np.zeros(N + 1)
    x[0] = x0
    y[0] = y0
    for i in range(N):
        x[i + 1] = x[i] + h
        y[i + 1] = y[i] + h * f(x[i], y[i])
    return x, y


def E_r(f, x0, y0, h, N):
    x = np.zeros(N + 1)
    y = np.zeros(N + 1)
    x[0] = x0
    y[0] = y0
    for i in range(N):
        x[i + 1] = x[i] + h
        y_b = y[i] + h * f(x[i], y[i])
        y[i + 1] = y[i] + h / 2 * (f(x[i], y[i]) + f(x[i + 1], y_b))
    return x, y


def R_k(f, x0, y0, h, N):
    x = np.zeros(N + 1)
    y = np.zeros(N + 1)
    x[0] = x0
    y[0] = y0
    for i in range(N):
        x[i + 1] = x[i] + h
        k1 = f(x[i], y[i])
        k2 = f(x[i] + h / 2, y[i] + h / 2 * k1)
        k3 = f(x[i] + h / 2, y[i] + h / 2 * k2)
        k4 = f(x[i] + h, y[i] + h * k3)
        y[i + 1] = y[i] + h / 6 * (k1 + k2 * 2 + k3 * 2 + k4)
    return x, y


[xx, yy] = E(f, 0, 1 / 3, h, n)
print(xx)
print(yy)
[xx_r, yy_r] = E_r(f, 0, 1 / 3, h, n)
print(yy_r)
[xx_R_k, yy_R_k] = R_k(f, 0, 1 / 3, h, n)

xxx = np.arange(0, 1 + h, h)
yyy = []
for i in xxx:
    yyy.append((math.e ** (-50 * i) / 3) + i ** 2)
    # yyy.append(xxx)
pl.plot(xx, yy, label="Euler", linestyle=":")
pl.plot(xx_r, yy_r, label="E_r", color="green", linestyle="-")
pl.plot(xx_R_k, yy_R_k, label="R_k", linestyle="--", color="purple")
pl.plot(xxx, yyy, color="red", label="right", alpha=0.3)
pl.legend()
pl.show()