import numpy as np
import math
import time


def thomas_factor(a, b, c):
    """
    Precompute Thomas factors for tridiagonal system.
    a: subdiag length n-1
    b: diag    length n
    c: super   length n-1
    """
    n = len(b)
    denom = np.empty(n, dtype=complex)
    cp = np.empty(n - 1, dtype=complex)

    denom[0] = b[0]
    cp[0] = c[0] / denom[0]
    for i in range(1, n - 1):
        denom[i] = b[i] - a[i - 1] * cp[i - 1]
        cp[i] = c[i] / denom[i]
    denom[n - 1] = b[n - 1] - a[n - 2] * cp[n - 2]
    return cp, denom


def thomas_solve(a, b, c, cp, denom, d):
    """Solve tridiagonal system using precomputed factors."""
    n = len(b)
    dp = np.empty(n, dtype=complex)

    dp[0] = d[0] / denom[0]
    for i in range(1, n):
        dp[i] = (d[i] - a[i - 1] * dp[i - 1]) / denom[i]

    x = np.empty(n, dtype=complex)
    x[-1] = dp[-1]
    for i in range(n - 2, -1, -1):
        x[i] = dp[i] - cp[i] * x[i + 1]
    return x



def precompute_cyclic_solver(n, a_val, b_val, c_val, alpha_corner, beta_corner):
    """
    Matrix A is cyclic tridiagonal:
      diag b_val, sub a_val, super c_val,
      corner A[0,n-1]=alpha_corner, A[n-1,0]=beta_corner
    Return precomputed data to solve A x = r many times.
    """
    b0 = b_val
    gamma = -b0  

    b = np.full(n, b_val, dtype=complex)
    a = np.full(n - 1, a_val, dtype=complex)
    c = np.full(n - 1, c_val, dtype=complex)

    b_mod = b.copy()
    b_mod[0] = b0 - gamma
    b_mod[-1] = b0 - alpha_corner * beta_corner / gamma

    cp, denom = thomas_factor(a, b_mod, c)

    u = np.zeros(n, dtype=complex)
    u[0] = gamma
    u[-1] = alpha_corner
    z = thomas_solve(a, b_mod, c, cp, denom, u)

    return {
        "n": n,
        "a": a, "c": c,
        "b_mod": b_mod,
        "cp": cp, "denom": denom,
        "gamma": gamma,
        "alpha_corner": alpha_corner,
        "beta_corner": beta_corner,
        "z": z
    }


def solve_cyclic(pre, r):
    """Solve A x = r using precomputed cyclic solver data."""
    n = pre["n"]
    a = pre["a"]; c = pre["c"]
    b_mod = pre["b_mod"]
    cp = pre["cp"]; denom = pre["denom"]
    gamma = pre["gamma"]
    beta = pre["beta_corner"]
    z = pre["z"]

    x = thomas_solve(a, b_mod, c, cp, denom, r)

    fact_num = x[0] + beta * x[-1] / gamma
    fact_den = 1.0 + z[0] + beta * z[-1] / gamma
    fact = fact_num / fact_den

    return x - fact * z



def exact_W_fourier(W0, t):
    N = len(W0)
    What0 = np.fft.fft(W0) / N

    if N % 2 == 0:
        k_int = np.concatenate([np.arange(0, N//2), np.arange(-N//2, 0)])
    else:
        k_int = np.concatenate([np.arange(0, (N-1)//2 + 1), np.arange(-(N-1)//2, 0)])
    k_int = k_int.astype(float)

    phase = np.exp(-1j * (k_int**2) * t / 2.0)
    What = What0 * phase
    return np.fft.ifft(What * N)



def crank_nicolson_W(W0, t_end, Nt):
    N = len(W0)
    dpsi = 2.0 * math.pi / N
    dt = t_end / Nt

    alpha = 1j * dt / (4.0 * dpsi * dpsi)

    a_val = -alpha
    c_val = -alpha
    b_val = 1.0 + 2.0 * alpha

    pre = precompute_cyclic_solver(
        n=N,
        a_val=a_val,
        b_val=b_val,
        c_val=c_val,
        alpha_corner=-alpha,
        beta_corner=-alpha
    )

    W = W0.astype(complex).copy()

    for _ in range(Nt):
        Wp = np.roll(W, -1)
        Wm = np.roll(W,  1)
        rhs = (1.0 - 2.0 * alpha) * W + alpha * (Wp + Wm)
        W = solve_cyclic(pre, rhs)

    return W



def run_one(N, Nt, a=1.0, k=40.0):
    t0 = 1.0 / (k * a)
    psi = np.linspace(0.0, 2.0 * math.pi, N, endpoint=False)

    W0 = np.exp(1j * 3 * psi) + 0.3 * np.exp(1j * 11 * psi)

    t_start = time.perf_counter()
    W_num = crank_nicolson_W(W0, t_end=t0, Nt=Nt)
    elapsed = time.perf_counter() - t_start

    W_ex = exact_W_fourier(W0, t=t0)
    rel_err = np.linalg.norm(W_num - W_ex) / np.linalg.norm(W_ex)

    return t0, elapsed, rel_err


def main():
    print("START GRID RUN")

    a = 1.0
    k = 40.0

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

    out_csv = "results_cn_circle.csv"
    with open(out_csv, "w", encoding="utf-8") as f:
        f.write("N,Nt,t0,time_s,rel_L2_error\n")
        for N in N_list:
            for Nt in Nt_list:
                t0, tsec, err = run_one(N, Nt, a=a, k=k)
                f.write(f"{N},{Nt},{t0:.12g},{tsec:.6f},{err:.12e}\n")
                print(f"N={N:4d} Nt={Nt:5d}  time={tsec:.4f}s  err={err:.6e}")

    print(f"END GRID RUN. Saved: {out_csv}")


if __name__ == "__main__":
    main()