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


class ODESolver:
    def euler_method(self, f, t_span, y0, h):
        t0, t_end = t_span
        N = int((t_end - t0) / h)
        t = np.linspace(t0, t_end, N + 1)
        y = np.zeros(N + 1)
        y[0] = y0

        for i in range(N):
            y[i + 1] = y[i] + h * f(t[i], y[i])

        return t, y

    def explicit_trapezoid(self, f, t_span, y0, h):
        t0, t_end = t_span
        N = int((t_end - t0) / h)
        t = np.linspace(t0, t_end, N + 1)
        y = np.zeros(N + 1)
        y[0] = y0

        for i in range(N):
            k1 = f(t[i], y[i])
            k2 = f(t[i] + h, y[i] + h * k1)
            y[i + 1] = y[i] + (h / 2) * (k1 + k2)

        return t, y

    def rk4_method(self, f, t_span, y0, h):
        t0, t_end = t_span
        N = int((t_end - t0) / h)
        t = np.linspace(t0, t_end, N + 1)
        y = np.zeros(N + 1)
        y[0] = y0

        for i in range(N):
            k1 = f(t[i], y[i])
            k2 = f(t[i] + h/2, y[i] + h/2 * k1)
            k3 = f(t[i] + h/2, y[i] + h/2 * k2)
            k4 = f(t[i] + h, y[i] + h * k3)
            y[i + 1] = y[i] + h/6 * (k1 + 2*k2 + 2*k3 + k4)

        return t, y

    def backward_euler(self, f, t_span, y0, h):
        t0, t_end = t_span
        N = int((t_end - t0) / h)
        t = np.linspace(t0, t_end, N + 1)
        y = np.zeros(N + 1)
        y[0] = y0

        for i in range(N):
            t_next = t[i + 1]

            def equation(z):
                return z - h * f(t_next, z) - y[i]

            y[i + 1] = fsolve(equation, y[i])[0]

        return t, y

    def implicit_trapezoid(self, f, t_span, y0, h):
        t0, t_end = t_span
        N = int((t_end - t0) / h)
        t = np.linspace(t0, t_end, N + 1)
        y = np.zeros(N + 1)
        y[0] = y0

        for i in range(N):
            t_n = t[i]
            t_np1 = t[i + 1]
            y_n = y[i]

            def equation(z):
                return z - y_n - (h/2) * (f(t_n, y_n) + f(t_np1, z))

            y[i + 1] = fsolve(equation, y_n)[0]

        return t, y

def problem_1_4():
    print("=== ЗАДАЧА 1.4: y' = y² - t, y(0) = 1 ===")

    def f(t, y):
        return y**2 - t

    t_span = [0, 1.0]
    y0 = 1
    h = 0.25

    solver = ODESolver()

    t_euler, y_euler = solver.euler_method(f, t_span, y0, h)
    t_trap, y_trap = solver.explicit_trapezoid(f, t_span, y0, h)
    t_rk4, y_rk4 = solver.rk4_method(f, t_span, y0, h)

    sol_ref = solve_ivp(f, t_span, [y0], method='RK45',
                       rtol=1e-10, atol=1e-12, dense_output=True)
    t_ref = np.linspace(t_span[0], t_span[1], 100)
    y_ref = sol_ref.sol(t_ref)[0]

    y_ref_euler = sol_ref.sol(t_euler)[0]
    y_ref_trap = sol_ref.sol(t_trap)[0]
    y_ref_rk4 = sol_ref.sol(t_rk4)[0]

    err_euler = np.linalg.norm(y_euler - y_ref_euler)
    err_trap = np.linalg.norm(y_trap - y_ref_trap)
    err_rk4 = np.linalg.norm(y_rk4 - y_ref_rk4)

    print(f"Ошибка метода Эйлера: {err_euler:.2e}")
    print(f"Ошибка явного метода трапеций: {err_trap:.2e}")
    print(f"Ошибка метода РК4: {err_rk4:.2e}")

    plt.figure(figsize=(12, 8))
    plt.plot(t_ref, y_ref, 'k-', linewidth=5, label='Эталонное решение (solve_ivp)')
    plt.plot(t_euler, y_euler, 'bo-', markersize=10, label='Метод Эйлера')
    plt.plot(t_trap, y_trap, 'c*-',  markersize=15, label='Явный метод трапеций')
    plt.plot(t_rk4, y_rk4, 'rs-', markersize=5, label='Метод РК4')

    plt.xlabel('t')
    plt.ylabel('y(t)')
    plt.title('Задача 1.4: y\' = y² - t, y(0) = 1')
    plt.legend()
    plt.grid(True, alpha=0.3)
    plt.show()

    return t_euler, y_euler, t_trap, y_trap, t_rk4, y_rk4, t_ref, y_ref


if __name__ == "__main__":
    result1 = problem_1_4()

def problem_2_4():
    print("\n=== ЗАДАЧА 2.4: y' = -2000y + t, y(0) = 0 ===")

    def f(t, y):
        return -2000 * y + t

    t_span = [0, 1]
    y0 = 0
    h = 1

    solver = ODESolver()

    t_euler, y_euler = solver.euler_method(f, t_span, y0, h)
    t_beuler, y_beuler = solver.backward_euler(f, t_span, y0, h)
    t_itrap, y_itrap = solver.implicit_trapezoid(f, t_span, y0, h)

    sol_ref = solve_ivp(f, t_span, [y0], method='Radau',
                       rtol=1e-10, atol=1e-12, dense_output=True)
    t_ref = np.linspace(t_span[0], t_span[1], 100)
    y_ref = sol_ref.sol(t_ref)[0]

    y_ref_euler = sol_ref.sol(t_euler)[0]
    y_ref_beuler = sol_ref.sol(t_beuler)[0]
    y_ref_itrap = sol_ref.sol(t_itrap)[0]

    err_euler = np.linalg.norm(y_euler - y_ref_euler)
    err_beuler = np.linalg.norm(y_beuler - y_ref_beuler)
    err_itrap = np.linalg.norm(y_itrap - y_ref_itrap)

    print(f"Ошибка явного метода Эйлера: {err_euler:.2e}")
    print(f"Ошибка неявного метода Эйлера: {err_beuler:.2e}")
    print(f"Ошибка неявного метода трапеций: {err_itrap:.2e}")

    plt.figure(figsize=(12, 8))
    plt.plot(t_ref, y_ref, 'k-', linewidth=5, label='Эталонное решение (solve_ivp)')
    plt.plot(t_euler, y_euler, 'bo-', markersize=4, label='Явный метод Эйлера')
    plt.plot(t_beuler, y_beuler, 'g^-', markersize=4, label='Неявный метод Эйлера')
    plt.plot(t_itrap, y_itrap, 'm*-',linewidth=2.2, markersize=4, label='Неявный метод трапеций')

    plt.xlabel('t')
    plt.ylabel('y(t)')
    plt.title('Задача 2.4: y\' = -2000y + t, y(0) = 0 (жесткая система)')
    plt.legend()
    plt.grid(True, alpha=0.3)
    plt.show()

    return t_euler, y_euler, t_beuler, y_beuler, t_itrap, y_itrap, t_ref, y_ref


if __name__ == "__main__":
    result2 = problem_2_4()