---

import numpy as np
import matplotlib.pyplot as plt


# ==========================================================
# МЕТОД ГАЛЕРКИНА
# ==========================================================

# Базовые функции
def phi1(x):
    return x - x**2 / 2


def phi2(x):
    return x - x**3 / 3


# Их производные
def dphi1(x):
    return 1 - x


def dphi2(x):
    return 1 - x**2


def ddphi1(x):
    return -np.ones_like(x)


def ddphi2(x):
    return -2 * x


# Функция u0(x), которая уже удовлетворяет граничным условиям:
# u0(0)=0, u0'(1)=1
def u0(x):
    return x


def du0(x):
    return np.ones_like(x)


def ddu0(x):
    return np.zeros_like(x)


# Оператор L[y] = y'' - 2x y' + 2y
def L(y, dy, ddy, x):
    return ddy - 2 * x * dy + 2 * y


# ==========================================================
# МЕТОД ПРОГОНКИ
# ==========================================================

def progonka(a, b, c, d):
    n = len(d)

    alpha = np.zeros(n)
    beta = np.zeros(n)

    alpha[0] = -c[0] / b[0]
    beta[0] = d[0] / b[0]

    for i in range(1, n):
        denominator = b[i] + a[i] * alpha[i - 1]
        alpha[i] = -c[i] / denominator if i < n - 1 else 0
        beta[i] = (d[i] - a[i] * beta[i - 1]) / denominator

    y = np.zeros(n)
    y[-1] = beta[-1]

    for i in range(n - 2, -1, -1):
        y[i] = alpha[i] * y[i + 1] + beta[i]

    return y


# ==========================================================
# ОСНОВНАЯ ПРОГРАММА
# ==========================================================

# --------------------------
# 1. РЕШЕНИЕ МЕТОДОМ ГАЛЕРКИНА
# --------------------------
x_int = np.linspace(0, 1, 5000)

# Правая часть уравнения
f = x_int

# L[u0] - f
r0 = L(u0(x_int), du0(x_int), ddu0(x_int), x_int) - f

# L[phi1], L[phi2]
L_phi1 = L(phi1(x_int), dphi1(x_int), ddphi1(x_int), x_int)
L_phi2 = L(phi2(x_int), dphi2(x_int), ddphi2(x_int), x_int)

# Составляем систему для c1 и c2
a11 = np.trapz(L_phi1 * phi1(x_int), x_int)
a12 = np.trapz(L_phi2 * phi1(x_int), x_int)
a21 = np.trapz(L_phi1 * phi2(x_int), x_int)
a22 = np.trapz(L_phi2 * phi2(x_int), x_int)

b1 = -np.trapz(r0 * phi1(x_int), x_int)
b2 = -np.trapz(r0 * phi2(x_int), x_int)

A = np.array([[a11, a12],
              [a21, a22]])

B = np.array([b1, b2])

c1, c2 = np.linalg.solve(A, B)

print("Метод Галеркина")
print("c1 =", c1)
print("c2 =", c2)
print()

# Формула решения Галеркина
def y_galerkin(x):
    return u0(x) + c1 * phi1(x) + c2 * phi2(x)


# --------------------------
# 2. РАЗНОСТНОЕ РЕШЕНИЕ
# --------------------------
N = 200
h = 1 / N
x = np.linspace(0, 1, N + 1)

# Неизвестные: y1, y2, ..., yN
n = N

a = np.zeros(n)
b = np.zeros(n)
c = np.zeros(n)
d = np.zeros(n)

# Внутренние узлы
for i in range(1, N):
    row = i - 1
    xi = x[i]

    A_i = 1 / h**2 + xi / h
    B_i = -2 / h**2 + 2
    C_i = 1 / h**2 - xi / h
    D_i = xi

    # y0 = 0
    if i == 1:
        a[row] = 0
        d[row] = D_i
    else:
        a[row] = A_i
        d[row] = D_i

    b[row] = B_i
    c[row] = C_i

# Последнее уравнение из правого граничного условия
row = n - 1
a[row] = 2 / h**2
b[row] = -2 / h**2 + 2
c[row] = 0
d[row] = 3 - 2 / h

y_unknown = progonka(a, b, c, d)

y_fd = np.zeros(N + 1)
y_fd[1:] = y_unknown

print("Разностное решение найдено")
print()

# --------------------------
# 3. ПРОВЕРКА ГРАНИЧНЫХ УСЛОВИЙ
# --------------------------
print("Проверка условий")

# Для Галеркина
print("Галеркин:")
print("y(0) =", y_galerkin(0))

dy_gal_1 = 1 + c1 * (1 - 1) + c2 * (1 - 1**2)
print("y'(1) =", dy_gal_1)
print()

# Для разностного решения
print("Разностное решение:")
print("y(0) =", y_fd[0])

# Центральная формула на правой границе через фиктивный узел формально была заложена в схему
# Но для проверки выведем одностороннюю приближенную производную:
dy_fd = (y_fd[-1] - y_fd[-2]) / h
print("приближенно y'(1) =", dy_fd)
print()


# --------------------------
# 4. ГРАФИК
# --------------------------
x_plot = np.linspace(0, 1, 500)

plt.figure(figsize=(9, 5))
plt.plot(x_plot, y_galerkin(x_plot), label="Галеркин")
plt.plot(x, y_fd, "--", label="Разностное решение")
plt.xlabel("x")
plt.ylabel("u(x)")
plt.title("Задание 1: u'' - 2x u' + 2u = x")
plt.grid(True)
plt.legend()
plt.tight_layout()
plt.show()