# coding:utf-8
'''
readme!!!

我们需要
1、 x 的区间 [a, b]                 ----   line 17
2、 步长 h (定义为 x_i+1 - x_i)      ----   line 21
3、 y(a) = k                       ----   line 24
4、 函数 f(x, y) 的解析式            ----   line 28
5、 解得的 y          相关的行 解不到的时候 需要 注释掉  ----   line 34  60  103
'''


# 导包
import matplotlib.pyplot as plt

# 输入区间 a、b
a = 0
b = 1

# 定义 h
h = 0.1

# 输入 y(a) = k
k = 1
y_e = [k]

# 定义函数 f
def f(x, y):
    res = 3 * y / (1 + x)
    return res


# 真实的 f: 也即 y = f(x), 一般用不上 这里写着只是为了某些人验证 代码的正确性
def f_true(x):
    y = (1 + x)**3
    return y


# (b-a)/h -> 得到总共有多少个 数对
num = int((b - a) / h)

# 得到 x ，但需要注意 此时仍少了 b ，如果有些傻宝没有注意整除怎么办？ 没事！反正只差了一点点 我们晚点直接append！
x = [(a + i * h) for i in range(num)]

# 得到最后一个 x
x.append(b)

''' 欧拉法 '''

# 得到 欧拉法 得到的 y
for x_i in x:
    # res 即 y_i+1、此即 y_i+1 = y_i + h * f(x_i, y_i)
    res = y_e[-1] + h * f(x_i, y_e[-1])
    y_e.append(res)

# 就像视频里面说的一样 这里多了个元素 我们pop一下
y_e.pop()

# 获得真实的 y
y_t = [f_true(i) for i in x]


''' 改进的欧拉法 '''

# 创建一个 改进的欧拉法 列表
y_r = [k]

# 获得 改进的欧拉法 的 y 值, 因为总共有num + 1个元素
for i in range(num):
    res = y_r[-1] + h * 0.5 * (f(x[i], y_r[-1]) + f(x[i + 1], y_e[i + 1]))
    y_r.append(res)


''' 四阶龙格-库塔方法 '''

# 创建 龙格-库塔方法 列表
y_RK = [k]

# 创建 x_half 列表 存储 x_i+1/2 值
x_half = []

# 利用切片 获得 x_i+1/2 的值
for x_i in x[1:]:
    x_half.append(x_i-h/2)

# 创建 k1、k2、k3、k4 列表
k1 = []
k2 = []
k3 = []
k4 = []
for i in range(len(x_half)):
    k1.append(f(x[i], y_RK[-1]))
    k2.append(f(x_half[i], y_RK[-1] + h * k1[-1]/2))
    k3.append(f(x_half[i], y_RK[-1] + h * k2[-1]/2))
    k4.append(f(x[i + 1], y_RK[-1] + h * k3[-1]))
    res = y_RK[-1] + h * (k1[-1] + 2*k2[-1] + 2*k3[-1] + k4[-1]) / 6
    y_RK.append(res)


# 绘制图像
plt.plot(x, y_e, label = 'Euler', color = 'green')
plt.plot(x, y_t, label = 'true', color = 'red')
plt.plot(x, y_r, label = 'Improved Euler', color = 'orange')
plt.plot(x, y_RK, label = 'Runge-Kutta', color = 'blue')

plt.legend()
plt.title('Experimental Image')
plt.show()