---

import numpy as np
import sympy as sp
import matplotlib.pyplot as plt
from matplotlib.widgets import TextBox

# ============================================================
# Условие (вариант 10):
#   u_t = u_xx + f(x,t),  f(x,t)= pi*t*sin(pi*x/2) + x*cos(pi*x/2)*(1+t*pi^2/4)
#   u(x,0)=0,  u(0,t)=u(1,t)=0
# ============================================================

n0, m0, t0 = 20, 25, 1.0


def f(x, t):
    return np.pi * t * np.sin(np.pi * x / 2) + x * np.cos(np.pi * x / 2) * (1 + t * np.pi ** 2 / 4)


def u_exact(x, t):
    """Точное решение, найденное подбором (проверено символьно: u_t-u_xx=f тождественно)."""
    return t * x * np.cos(np.pi * x / 2)


# ============================================================
# МЕТОД 1. Сеточная краевая задача: неявная схема (sigma=1) + прямая прогонка
# (формулы (7.11)-(7.12) методички Матвеева, метод прогонки — формулы (6.7)-(6.10))
# ============================================================

def solve_grid(n, m, t_final):
    h = 1.0 / n
    tau = t_final / m
    x = np.linspace(0, 1, n + 1)
    u = np.zeros(n + 1)  # u(x,0) = 0

    d = -(2 + h ** 2 / tau)
    e = 1.0

    for j in range(m):
        t_next = (j + 1) * tau
        q = np.zeros(n + 1)
        phi = np.zeros(n + 1)
        q[0], phi[0] = 0.0, 0.0  # chi1=mu1=0 (левое условие u(0,t)=0)

        for i in range(1, n):
            g_i = -(h ** 2 * f(x[i], t_next) + (h ** 2 / tau) * u[i])
            denom = q[i - 1] + d
            q[i] = -e / denom
            phi[i] = (g_i - phi[i - 1]) / denom

        u_new = np.zeros(n + 1)
        u_new[n] = 0.0  # chi2=mu2=0 (правое условие u(1,t)=0)
        for i in range(n - 1, -1, -1):
            u_new[i] = q[i] * u_new[i + 1] + phi[i]
        u = u_new

    return x, u


# ============================================================
# МЕТОД 2. Метод Фурье (разложение по sin(k*pi*x), до 3-го члена)
# ============================================================

def solve_fourier_symbolic(num_terms=3):
    x_sym, t_sym, s_sym, k_sym = sp.symbols('x t s k', real=True)
    f_sym = sp.pi * t_sym * sp.sin(sp.pi * x_sym / 2) + \
        x_sym * sp.cos(sp.pi * x_sym / 2) * (1 + t_sym * sp.pi ** 2 / 4)

    vk_funcs = []
    for k in range(1, num_terms + 1):
        ek = sp.sin(k * sp.pi * x_sym)
        gk = sp.simplify(2 * sp.integrate(f_sym * ek, (x_sym, 0, 1)))

        vk = sp.Function('v')
        lam = -(k * sp.pi) ** 2
        eq = sp.Eq(sp.diff(vk(t_sym), t_sym), lam * vk(t_sym) + gk)
        sol = sp.dsolve(eq, vk(t_sym), ics={vk(0): 0})
        vk_expr = sp.simplify(sol.rhs)
        vk_funcs.append((k, sp.lambdify(t_sym, vk_expr, 'numpy'), vk_expr))

    return vk_funcs


VK_FUNCS = solve_fourier_symbolic(num_terms=3)
for k, _, expr in VK_FUNCS:
    print(f"v_{k}(t) = {expr}")


def u_fourier(x, t, vk_funcs=VK_FUNCS):
    x = np.asarray(x, dtype=float)
    result = np.zeros_like(x)
    for k, vk_func, _ in vk_funcs:
        result = result + vk_func(t) * np.sin(k * np.pi * x)
    return result


# ============================================================
# Сравнение
# ============================================================

def print_deviations(n, m, t_final):
    xs, u_grid = solve_grid(n, m, t_final)
    ue = u_exact(xs, t_final)
    uf = u_fourier(xs, t_final)
    err_grid = np.max(np.abs(u_grid - ue))
    err_fourier = np.max(np.abs(uf - ue))
    print(f"--- n={n}, m={m}, t0={t_final} ---")
    print(f"  Сеточный метод (неявная схема + прогонка): max|u-u_точн| = {err_grid:.4e}")
    print(f"  Метод Фурье (3 члена):                     max|u-u_точн| = {err_fourier:.4e}\n")


# ============================================================
# Один график, x от 0 до 1, момент времени t=t0
# ============================================================

N_PLOT_POINTS = 500


def refresh_exact_curve():
    x0, x1 = ax.get_xlim()
    x0c, x1c = max(x0, 0.0), min(x1, 1.0)
    if x1c <= x0c:
        return
    x_local = np.linspace(x0c, x1c, N_PLOT_POINTS)
    line_exact.set_data(x_local, u_exact(x_local, t0))


fig, ax = plt.subplots(figsize=(9, 6))
plt.subplots_adjust(bottom=0.28)

line_exact, = ax.plot([], [], '-', color='black', linewidth=2, label='Точное решение')

xs0, ug0 = solve_grid(n0, m0, t0)
line_grid, = ax.plot(xs0, ug0, linestyle='None', marker='o', markersize=7,
                      markerfacecolor='none', markeredgewidth=1.8, color='tab:red',
                      label=f'Сеточный метод, n={n0}, m={m0}')

x_fine = np.linspace(0, 1, N_PLOT_POINTS)
line_fourier, = ax.plot(x_fine, u_fourier(x_fine, t0), linestyle=(0, (6, 3)),
                         linewidth=2, color='tab:blue', label='Метод Фурье (3 члена)')

ax.set_xlim(0, 1)
refresh_exact_curve()

ax.set_xlabel('x')
ax.set_ylabel('u(x, t0)')
ax.set_title(f"u_t=u_xx+f(x,t),  u(x,0)=0,  u(0,t)=u(1,t)=0\nвариант 10, момент времени t0={t0}")
ax.legend(loc='best')
ax.grid(alpha=0.3)

print_deviations(n0, m0, t0)

# ---------- Поля ввода n, m ----------
axbox_n = plt.axes([0.25, 0.10, 0.08, 0.06])
text_n = TextBox(axbox_n, "n: ", initial=str(n0))

axbox_m = plt.axes([0.55, 0.10, 0.08, 0.06])
text_m = TextBox(axbox_m, "m: ", initial=str(m0))

state = {'n': n0, 'm': m0}


def update(_text):
    try:
        n = int(text_n.text)
        m = int(text_m.text)
        if n < 3 or m < 1:
            raise ValueError
    except ValueError:
        print("Некорректные n или m — введите целые числа (n>=3, m>=1).")
        return

    state['n'], state['m'] = n, m
    xs, ug = solve_grid(n, m, t0)
    line_grid.set_xdata(xs)
    line_grid.set_ydata(ug)
    line_grid.set_label(f'Сеточный метод, n={n}, m={m}')

    ax.legend(loc='best')
    fig.canvas.draw_idle()

    print_deviations(n, m, t0)


text_n.on_submit(update)
text_m.on_submit(update)

MIN_RANGE = 1e-6


def on_scroll(event):
    if event.inaxes != ax:
        return
    scale = 1.2
    factor = (1 / scale) if event.button == 'up' else scale

    xlim = ax.get_xlim()
    ylim = ax.get_ylim()
    xdata, ydata = event.xdata, event.ydata
    if xdata is None or ydata is None:
        return

    new_w = (xlim[1] - xlim[0]) * factor
    new_h = (ylim[1] - ylim[0]) * factor

    if new_w < MIN_RANGE or new_h < MIN_RANGE:
        return

    relx = (xlim[1] - xdata) / (xlim[1] - xlim[0])
    rely = (ylim[1] - ydata) / (ylim[1] - ylim[0])

    ax.set_xlim([xdata - new_w * (1 - relx), xdata + new_w * relx])
    ax.set_ylim([ydata - new_h * (1 - rely), ydata + new_h * rely])

    refresh_exact_curve()
    fig.canvas.draw_idle()


fig.canvas.mpl_connect('scroll_event', on_scroll)

pan_state = {'active': False, 'x0': None, 'y0': None, 'xlim0': None, 'ylim0': None}


def on_press(event):
    if event.inaxes != ax:
        return
    if event.button not in (1, 3):
        return
    pan_state['active'] = True
    pan_state['x0'] = event.xdata
    pan_state['y0'] = event.ydata
    pan_state['xlim0'] = ax.get_xlim()
    pan_state['ylim0'] = ax.get_ylim()


def on_release(_event):
    if pan_state['active']:
        refresh_exact_curve()
        fig.canvas.draw_idle()
    pan_state['active'] = False


def on_motion(event):
    if not pan_state['active'] or event.inaxes != ax:
        return
    if event.xdata is None or event.ydata is None:
        return
    dx = event.xdata - pan_state['x0']
    dy = event.ydata - pan_state['y0']
    xlim0, ylim0 = pan_state['xlim0'], pan_state['ylim0']
    ax.set_xlim(xlim0[0] - dx, xlim0[1] - dx)
    ax.set_ylim(ylim0[0] - dy, ylim0[1] - dy)
    fig.canvas.draw_idle()


fig.canvas.mpl_connect('button_press_event', on_press)
fig.canvas.mpl_connect('button_release_event', on_release)
fig.canvas.mpl_connect('motion_notify_event', on_motion)

plt.show()