"""
Численное решение параболического уравнения Малюжинца
методом Кранка–Николсона для контура в виде окружности.

Модельная задача (раздел C2):
    dW/dt= (i/2) * d^2W/dpsi^2, psi in [0, 2*pi),  t in [0, T]

Периодические граничные условия:
    W(0, t)= W(2*pi, t), dW/dpsi(0,t) = dW/dpsi(2*pi, t)

Начальное условие (раздел C2):
    W0(psi) =exp(3i*psi) + 0.3*exp(11i*psi)

Точное решение (раздел C3):
    W_ex(psi, t) =exp(3i*psi - 9i*t/2) + 0.3*exp(11i*psi - 121i*t/2)

Параметры (раздел C7):
    a = 1,  k = 40,T = t0 = 1/(k*a) =0.025
"""

import numpy as np
import time
import os
from scipy.linalg import solve_banded

A_RADIUS = 1.0
K_WAVE   = 40.0
KA       = K_WAVE * A_RADIUS   
T_FINAL  = 1.0 / KA            



def initial_condition(psi):
    """W0(psi) = exp(3i*psi) + 0.3*exp(11i*psi)."""
    return np.exp(3j * psi) + 0.3 * np.exp(11j * psi)


def exact_solution(psi, t):
    """
    W_ex(psi, t) = exp(3i*psi - 9i*t/2)
                  + 0.3*exp(11i*psi - 121i*t/2).

    Выводится из того, что W(psi, 0) содержит только две гармоники
    m=3 и m=11, каждая эволюционирует как exp(-i*m^2*t/2).
    """
    return (np.exp(3j * psi) * np.exp(-9j * t / 2)
            + 0.3 * np.exp(11j * psi) * np.exp(-121j * t / 2))



def cn_step(W, dpsi, dt):
    """
    Один шаг схемы Кранка–Николсона для уравнения
        dW/dt = (i/2) * d^2W/dpsi^2
    с периодическими граничными условиями.

    Параметр схемы:
        alpha = i*dt / (4*dpsi^2)

    Трёхдиагональная система на каждом шаге:
        A * W^{n+1} = B * W^n
    где
        A = diag(1+2*alpha) + off-diag(-alpha) + corners(-alpha)
        B = diag(1-2*alpha) + off-diag(+alpha) + corners(+alpha)

    Угловые элементы обрабатываются методом Шермана–Моррисона.

    Параметры
    ----------
    W    : ndarray shape (N,), комплексный, решение на шаге n
    dpsi : float, шаг по psi
    dt   : float, шаг по t

    Возвращает
    ----------
    W_new : ndarray shape (N,), решение на шаге n+1
    """
    N = len(W)
    alpha = 1j * dt / (4.0 * dpsi**2)

    Wp  = np.roll(W, -1)  
    Wm  = np.roll(W,  1)  
    rhs = alpha * Wm + (1 - 2 * alpha) * W + alpha * Wp

    b   = 1 + 2 * alpha  
    off = -alpha          

    gamma = -b

    ab      = np.zeros((3, N), dtype=complex)
    ab[0, 1:]  = off   
    ab[1, :]   = b     
    ab[2, :-1] = off   

    ab_mod = ab.copy()
    ab_mod[1, 0]   -= gamma
    ab_mod[1, N-1] -= off**2 / gamma

    u = np.zeros(N, dtype=complex)
    u[0]   = gamma
    u[N-1] = off

    y = solve_banded((1, 1), ab_mod, rhs)
    z = solve_banded((1, 1), ab_mod, u)

    vTy = y[0] + (off / gamma) * y[N-1]
    vTz = z[0] + (off / gamma) * z[N-1]
    return y - z * vTy / (1 + vTz)



def solve(N, Nt, T=T_FINAL):
    """
    Решение задачи на сетке N x Nt.

    Параметры
    ----------
    N  : int, число угловых узлов
    Nt : int, число временных шагов
    T  : float, конечное время

    Возвращает
    ----------
    W_num   : ndarray (N,), численное решение при t = T
    psi_arr : ndarray (N,), сетка по psi
    runtime : float, время счёта в секундах
    """
    psi_arr = np.linspace(0, 2 * np.pi, N, endpoint=False)
    dpsi    = psi_arr[1] - psi_arr[0]
    dt      = T / Nt

    W = initial_condition(psi_arr).astype(complex)

    t_start = time.time()
    for _ in range(Nt):
        W = cn_step(W, dpsi, dt)
    runtime = time.time() - t_start

    return W, psi_arr, runtime



def relative_L2_error(W_num, W_ref):
    """
    eps_L2 = ||W_num - W_ref||_2 / ||W_ref||_2
    """
    diff = W_num - W_ref
    return np.linalg.norm(diff) / np.linalg.norm(W_ref)



if __name__ == "__main__":
    os.makedirs("data",    exist_ok=True)
    os.makedirs("figures", exist_ok=True)

    N_list  = [128, 256, 512, 1024]
    Nt_list = [250, 500, 1000, 2000, 4000]

    print("=" * 60)
    print(f"Параметры: a={A_RADIUS}, k={K_WAVE}, ka={KA}, T={T_FINAL}")
    print("=" * 60)
    print("\nТаблица C8: фиксированный Nt = 4000")
    print(f"{'N':>6}  {'eps_L2':>14}  {'runtime, s':>12}")
    print("-" * 40)

    Nt_fixed = 4000
    eps_c8   = []
    rt_c8    = []
    for N in N_list:
        W_num, psi, rt = solve(N, Nt_fixed)
        W_ref = exact_solution(psi, T_FINAL)
        eps   = relative_L2_error(W_num, W_ref)
        eps_c8.append(eps)
        rt_c8.append(rt)
        print(f"{N:>6}  {eps:>14.4e}  {rt:>12.6f}")

    print("\nТаблица C9: порядок сходимости")
    print(f"{'Переход':>12}  {'ratio':>8}  {'p_N':>6}")
    print("-" * 32)
    for i in range(1, len(N_list)):
        ratio = eps_c8[i-1] / eps_c8[i]
        p     = np.log(ratio) / np.log(2)
        print(f"{N_list[i-1]:>5} -> {N_list[i]:>5}  {ratio:>8.4f}  {p:>6.3f}")

    print("\nТаблица C11: eps_L2(N, Nt)")
    header = f"{'Nt':>5}  " + "  ".join(f"N={N:>5}" for N in N_list)
    print(header)
    print("-" * 70)

    table = {}
    for Nt in Nt_list:
        row = f"{Nt:>5}  "
        for N in N_list:
            W_num, psi, rt = solve(N, Nt)
            W_ref = exact_solution(psi, T_FINAL)
            eps   = relative_L2_error(W_num, W_ref)
            table[(N, Nt)] = (eps, rt)
            row += f"  {eps:.4e}"
        print(row)

    print("\nТаблица D10: время выполнения при Nt = 4000")
    print(f"{'N':>6}  {'runtime, s':>12}  {'ratio':>8}")
    print("-" * 35)
    print(f"{N_list[0]:>6}  {rt_c8[0]:>12.6f}  {'–':>8}")
    for i in range(1, len(N_list)):
        r = rt_c8[i] / rt_c8[i-1]
        print(f"{N_list[i]:>6}  {rt_c8[i]:>12.6f}  {r:>8.2f}")

    np.savez("data/results_CD.npz",
             N_list=N_list, Nt_list=Nt_list,
             eps_c8=np.array(eps_c8),
             rt_c8=np.array(rt_c8),
             ka=KA, T=T_FINAL, k=K_WAVE, a=A_RADIUS)
    print("\nДанные сохранены: data/results_CD.npz")