"""Check (iii) of 09-anchor-formation: the uniform part U(l) = place average of da(x,l) (item 3).

Exact solution of the linear equations of the means (M -> infinity),
    d^2 q/dl^2 = -K_a q - e_{x_s},  q(0) = dq/dl(0) = 0,  K_a = kappa_h*Lap + kappa_a*P_h.
Cubic torus L = 8, 12, 16, 20 (kappa_a = 1; RK4, dt = 0.05) against (9.14):
    U(l) = U_st [1 - cos(l/T_U)] + O(1/L) U_st,  T_U = sqrt(L^3 c_phi^(L)),  c_phi^(L) = 1/kappa_a + g_L(0),
    U_st = -[c_phi^(L) - g_L(x_s-h)];
  plus two variations at L = 12 (anchor farther, kappa_a = 4), the lowest normal mode of K_a from
  a sparse eigen-solver, and the slow image: place-dependent part of psi_0(x)psi_0(x_s)/nu_0 against
  -(1-beta_L) g_L(x-h), beta_L = g_L(x_s-h)/c_phi^(L).
Ring L = 128, 256, 512 (kappa_a = 1, d_s = L/4; exact eigen-decomposition) against (9.15):
    U(l) = -(L/2kappa_h) [f(alpha) - f(alpha+tau)/2 - f(alpha-tau)/2],  tau = c_0 l/L, alpha = d_s/L,
  f the odd 2-periodic extension of x(1-x).
Ring, discrete derivation of (9.15): the even in-band labels from the secular equation (9.19),
  tan(kL/2) = kappa_a/(2 kappa_h sin k), against the eigenvalues of K_a; the exact mode sum of Step 9
  (weights N_j^2 cos(delta_j) cot(k_j/2)/(4 kappa_h L sin^2(k_j/2)), phases sin(k_j d_s + delta_j)) against
  the eigen-decomposition (difference = localized mode, bound 2/(kappa_a L)); the remainder
  R = U - U_cont for alpha = 1/8, 1/4, 1/2, kappa_a = 1, 4 over tau <= 8, scaled by kappa_a/(1+tau)
  (bound C[(1+tau)/kappa_a + tau^(2/3) L^(-1/3)/kappa_h + log L/(kappa_h L)]), and for a stiff anchor
  (kappa_a = 1e4) over tau <= 8 and 32 against tau^(2/3) L^(-1/3).
Ring with a weak anchor (kappa_a L << kappa_h; L = 32, 64): the slow-mode form of (9.14) with
  c_phi^(L) = 1/kappa_a + (L^2-1)/(12 kappa_h L), T_U = sqrt(L c_phi^(L)).
"""
from pathlib import Path
import numpy as np
import scipy.sparse as sp
from scipy.sparse.linalg import eigsh
from scipy.optimize import brentq

OUT = Path(__file__).parent / "output"
OUT.mkdir(exist_ok=True)
KH = 1.0
C0 = np.sqrt(KH)


def g_torus(L, ys):
    """g_L(y) = L^-3 sum_{k != 0} cos(k.y) / (kappa_h Lap)(k), for an array ys of shape (3, M)."""
    ks = 2 * np.pi * np.arange(L) / L
    k1, k2, k3 = [k.ravel() for k in np.meshgrid(ks, ks, ks, indexing="ij")]
    om2 = 2 * KH * (3 - np.cos(k1) - np.cos(k2) - np.cos(k3))
    nz = om2 > 1e-12
    out = []
    for c in range(0, ys.shape[1], 512):
        y = ys[:, c:c + 512]
        ph = np.cos(k1[nz, None] * y[0] + k2[nz, None] * y[1] + k3[nz, None] * y[2])
        out.append((ph / om2[nz, None]).sum(0) / L ** 3)
    return np.concatenate(out)


def lap3(q):
    return 6 * q - sum(np.roll(q, s, axis=a) for a in range(3) for s in (1, -1))


def run_cubic(L, h, ka, lmax, dt=0.05):
    q = np.zeros((L, L, L))
    p = np.zeros_like(q)
    src = np.zeros_like(q)
    src[0, 0, 0] = 1.0
    hh = tuple(v % L for v in h)

    def acc(qq):
        a = -KH * lap3(qq) - src
        a[hh] -= ka * qq[hh]
        return a

    n = int(round(lmax / dt))
    U = np.zeros(n + 1)
    for m in range(1, n + 1):
        k1q, k1p = p, acc(q)
        k2q, k2p = p + 0.5 * dt * k1p, acc(q + 0.5 * dt * k1q)
        k3q, k3p = p + 0.5 * dt * k2p, acc(q + 0.5 * dt * k2q)
        k4q, k4p = p + dt * k3p, acc(q + dt * k3q)
        q = q + dt / 6 * (k1q + 2 * k2q + 2 * k3q + k4q)
        p = p + dt / 6 * (k1p + 2 * k2p + 2 * k3p + k4p)
        U[m] = q.mean()
    return np.arange(n + 1) * dt, U


def sparse_Ka(L, h, ka):
    one = sp.identity(L, format="csr")
    shift = sp.csr_matrix((np.ones(L), (np.arange(L), (np.arange(L) + 1) % L)), shape=(L, L))
    lap1 = 2 * one - shift - shift.T
    lap = (sp.kron(sp.kron(lap1, one), one) + sp.kron(sp.kron(one, lap1), one)
           + sp.kron(sp.kron(one, one), lap1))
    K = (KH * lap).tolil()
    ih = np.ravel_multi_index(tuple(v % L for v in h), (L, L, L))
    K[ih, ih] += ka
    return K.tocsc(), ih


def cubic_case(L, h, ka, lines, csv, tag):
    coords = np.array(np.unravel_index(np.arange(L ** 3), (L, L, L)))
    g0 = g_torus(L, np.zeros((3, 1)))[0]
    gsh = g_torus(L, -np.array(h, dtype=float)[:, None])[0]
    cphi = 1 / ka + g0
    Ust = -(cphi - gsh)
    TU = np.sqrt(L ** 3 * cphi)
    lam, U = run_cubic(L, h, ka, 2.2 * np.pi * TU)
    pred = Ust * (1 - np.cos(lam / TU))
    dev = np.max(np.abs(U - pred)) / abs(Ust)
    first = lam < 2 * np.pi * TU
    lmin = lam[first][np.argmin(U[first])]
    # lowest normal mode (independent, sparse shift-invert)
    K, ih = sparse_Ka(L, h, ka)
    nu, psi = eigsh(K, k=1, sigma=0, which="LM")
    nu0, psi0 = nu[0], psi[:, 0]
    A0 = psi0 * psi0[0] / nu0
    A0perp = A0 - A0.mean()
    beta = gsh / cphi
    gxh = g_torus(L, coords - np.array(h)[:, None] % L)
    img = -(1 - beta) * gxh
    dimg = np.max(np.abs(A0perp - img)) / np.max(np.abs(img))
    lines.append(f"  L={L:2d} kappa_a={ka} h-x_s={h}: c_phi^(L)={cphi:.5f} U_st={Ust:.5f} T_U={TU:.2f}; "
                 f"1/sqrt(nu_0)={1/np.sqrt(nu0):.2f} (ratio {1/np.sqrt(nu0)/TU:.5f}); "
                 f"max|U-pred|/|U_st|={dev:.2e} (x L = {dev*L:.4f}); "
                 f"first minimum at l/(pi T_U)={lmin/(np.pi*TU):.4f}, U_min/U_st={U[first].min()/Ust:.4f}; "
                 f"U(2 pi T_U)/U_st={np.interp(2*np.pi*TU, lam, U)/Ust:+.4f}; "
                 f"slow image: max|A0perp+(1-beta)g(x-h)|/max|.|={dimg:.2e} (x L^2 = {dimg*L*L:.2f}), beta_L={beta:.4f}")
    for j in range(0, len(lam), 20):
        csv.append(f"3,{tag},{L},{ka},{lam[j]:.3f},{U[j]:.6e},{pred[j]:.6e}")


def fext(x):
    y = np.mod(x + 1, 2) - 1  # in [-1, 1)
    return np.where(y >= 0, y * (1 - y), y * (1 + y))


def ring_case(L, lines, csv, ka=1.0):
    ds = L // 4
    i = np.arange(L)
    K = np.zeros((L, L))
    K[i, i] = 2 * KH
    K[i, (i + 1) % L] = -KH
    K[i, (i - 1) % L] = -KH
    K[ds, ds] += ka
    w2, V = np.linalg.eigh(K)
    tau = np.linspace(0, 4, 2001)
    lam = tau * L / C0
    U = -((V.sum(0) / L) * V[0] / w2) @ (1 - np.cos(np.sqrt(w2)[:, None] * lam[None, :]))
    alpha = ds / L
    Uc = -(L / (2 * KH)) * (fext(alpha) - 0.5 * fext(alpha + tau) - 0.5 * fext(alpha - tau))
    Ust = -(1 / ka + ds * (L - ds) / (2 * KH * L))
    dev = np.max(np.abs(U - Uc)) / abs(Ust)
    dev1 = np.max(np.abs(U - Uc)[tau <= 2]) / abs(Ust)
    first = tau < 2
    tmin = tau[first][np.argmin(U[first])]
    lines.append(f"  L={L:3d} d_s={ds}: U_st={Ust:.3f}; max|U-U_cont|/|U_st| over tau<=2: {dev1:.4f} (x L = {dev1*L:.2f}), "
                 f"over tau<=4: {dev:.4f}; first minimum at c0 l/L = {tmin:.4f} (pred 1), U_min/U_st={U[first].min()/Ust:.4f} (pred 2); "
                 f"U(c0 l = 2L)/U_st = {np.interp(2.0, tau, U)/Ust:+.4f} (pred 0)")
    for j in range(0, len(tau), 4):
        csv.append(f"1,ring,{L},{ka},{lam[j]:.3f},{U[j]:.6e},{Uc[j]:.6e}")


def ring_eig(L, ds, ka):
    """Anchor at place ds, source at place 0."""
    i = np.arange(L)
    K = np.zeros((L, L))
    K[i, i] = 2 * KH
    K[i, (i + 1) % L] = -KH
    K[i, (i - 1) % L] = -KH
    K[ds % L, ds % L] += ka
    return np.linalg.eigh(K)


def U_eig(w2, V, L, lam):
    return -((V.sum(0) / L) * V[0] / w2) @ (1 - np.cos(np.sqrt(w2)[:, None] * lam[None, :]))


def even_labels(L, ka):
    """Roots of -2 kappa_h sin k sin(kL/2) + kappa_a cos(kL/2) = 0 in (2j pi/L, (2j+1) pi/L), L even."""
    f = lambda k: -2 * KH * np.sin(k) * np.sin(k * L / 2) + ka * np.cos(k * L / 2)
    return np.array([brentq(f, max(2 * j * np.pi / L, 1e-12), (2 * j + 1) * np.pi / L, xtol=1e-15)
                     for j in range(L // 2)])


def U_modesum(L, ka, ds, lam):
    k = even_labels(L, ka)
    dl = np.arctan(2 * KH * np.sin(k) / ka)
    N2 = 1 / (L / 2 + np.sin(2 * dl) * np.cos(k) / (2 * np.sin(k)))
    w = N2 * np.cos(dl) / np.tan(k / 2) / (4 * KH * L * np.sin(k / 2) ** 2)
    om = 2 * C0 * np.sin(k / 2)
    return -(w * np.sin(k * ds + dl)) @ (1 - np.cos(om[:, None] * lam[None, :])), k, dl


def U_cont(L, alpha, tau):
    return -(L / (2 * KH)) * (fext(alpha) - 0.5 * fext(alpha + tau) - 0.5 * fext(alpha - tau))


def ring_discrete(L, ka, lines):
    tau = np.linspace(0, 8, 4001)
    lam = tau * L / C0
    out = []
    for frac in (8, 4, 2):
        ds = L // frac
        w2, V = ring_eig(L, ds, ka)
        U = U_eig(w2, V, L, lam)
        Um, k, dl = U_modesum(L, ka, ds, lam)
        if frac == 8:
            ev = np.sort(w2[(np.abs(V[ds]) > 1e-9) & (w2 < 4 * KH)])
            dq = np.max(np.abs(ev - 4 * KH * np.sin(k / 2) ** 2)) if len(ev) == len(k) else np.nan
            dd = np.max(np.abs(k * L - (2 * np.arange(len(k)) + 1) * np.pi + 2 * dl))
            nub = w2[-1]
            out.append(f"    L={L} kappa_a={ka}: {len(k)} even in-band labels from (9.19) vs {len(ev)} eigenvalues, "
                       f"max|nu_eig - 4 kappa_h sin^2(k_j/2)| = {dq:.1e}; max|k_j L - (2j+1)pi + 2 delta_j| = {dd:.1e}; "
                       f"top eigenvalue {nub:.6f} (line value 2kh+sqrt(4kh^2+ka^2) = {2*KH+np.sqrt(4*KH**2+ka**2):.6f})")
        loc = np.max(np.abs(U - Um))
        R = U - U_cont(L, ds / L, tau)
        sc = np.max(np.abs(R) * ka / (1 + tau))
        out.append(f"      alpha=1/{frac}: max|U_eig - U_modesum| = {loc:.2e} (localized mode; bound 2/(kappa_a L) = {2/(ka*L):.2e}); "
                   f"max_(tau<=8)|R| = {np.max(np.abs(R)):.3f}, max_(tau<=2)|R| = {np.max(np.abs(R[tau <= 2])):.3f}, "
                   f"max_tau |R| kappa_a/(1+tau) = {sc:.3f}, |R(tau=8)| = {abs(R[-1]):.3f}, |U_st| = {1/ka + ds*(L-ds)/(2*KH*L):.2f}")
    lines += out


def ring_stiff(L, ka, lines):
    tau = np.linspace(0, 32, 16001)
    lam = tau * L / C0
    ds = L // 4
    w2, V = ring_eig(L, ds, ka)
    R = U_eig(w2, V, L, lam) - U_cont(L, 0.25, tau)
    r8, r32 = np.max(np.abs(R[tau <= 8])), np.max(np.abs(R))
    lines.append(f"    stiff anchor kappa_a={ka:.0e}, L={L}, alpha=1/4: max_(tau<=8)|R| = {r8:.3f} "
                 f"(8^(2/3) L^(-1/3) = {8**(2/3)*L**(-1/3):.3f}, ratio {r8/(8**(2/3)*L**(-1/3)):.3f}); max_(tau<=32)|R| = {r32:.3f} "
                 f"(32^(2/3) L^(-1/3) = {32**(2/3)*L**(-1/3):.3f}, ratio {r32/(32**(2/3)*L**(-1/3)):.3f})")


def ring_weak(L, ka, lines, csv):
    ds = L // 4
    i = np.arange(L)
    K = np.zeros((L, L))
    K[i, i] = 2 * KH
    K[i, (i + 1) % L] = -KH
    K[i, (i - 1) % L] = -KH
    K[ds, ds] += ka
    w2, V = np.linalg.eigh(K)
    cphi = 1 / ka + (L * L - 1) / (12 * KH * L)
    TU = np.sqrt(L * cphi)
    Ust = -(1 / ka + ds * (L - ds) / (2 * KH * L))
    lam = np.linspace(0, 2.2 * np.pi * TU, 4001)
    U = -((V.sum(0) / L) * V[0] / w2) @ (1 - np.cos(np.sqrt(w2)[:, None] * lam[None, :]))
    pred = Ust * (1 - np.cos(lam / TU))
    dev = np.max(np.abs(U - pred)) / abs(Ust)
    lines.append(f"  L={L} kappa_a={ka} (kappa_a L/kappa_h = {ka*L/KH:.3f}): T_U={TU:.1f} (wrap L/c0 = {L/C0:.0f}), "
                 f"1/sqrt(nu_0)/T_U = {1/np.sqrt(w2[0])/TU:.5f}, max|U-pred|/|U_st| = {dev:.4f}")
    for j in range(0, len(lam), 8):
        csv.append(f"1,ring_weak,{L},{ka},{lam[j]:.3f},{U[j]:.6e},{pred[j]:.6e}")


def main():
    lines = ["Check (iii): the uniform part (item 3)", f"kappa_h = {KH}, c_0 = {C0}", "",
             "CUBIC TORUS, kappa_a = 1, h - x_s = (2,1,0): prediction (9.14)"]
    csv = ["D,case,L,kappa_a,lambda,U_exact,U_pred"]
    for L in (8, 12, 16, 20):
        cubic_case(L, (2, 1, 0), 1.0, lines, csv, "sizes")
    lines.append("CUBIC TORUS, L = 12: anchor farther and stiffer anchor")
    cubic_case(12, (5, 4, 3), 1.0, lines, csv, "far_anchor")
    cubic_case(12, (2, 1, 0), 4.0, lines, csv, "kappa_a_4")
    lines += ["", "RING, kappa_a = 1, d_s = L/4: prediction (9.15)"]
    for L in (128, 256, 512):
        ring_case(L, lines, csv)
    lines += ["", "RING, discrete secular equation (9.19), exact mode sum and remainder of (9.15)"]
    for L in (128, 256, 512):
        for ka in (1.0, 4.0):
            ring_discrete(L, ka, lines)
    for L in (128, 256, 512):
        ring_stiff(L, 1e4, lines)
    lines += ["", "RING, weak anchor kappa_a L << kappa_h: slow-mode form (9.14) with c_phi^(L) = 1/kappa_a + (L^2-1)/(12 kappa_h L)"]
    for L, ka in ((32, 0.004), (64, 0.002), (64, 0.0005)):
        ring_weak(L, ka, lines, csv)
    (OUT / "uniform_check.txt").write_text("\n".join(lines) + "\n", encoding="utf-8")
    (OUT / "uniform_U.csv").write_text("\n".join(csv) + "\n", encoding="utf-8")
    print("\n".join(lines))


if __name__ == "__main__":
    main()