"""
2D Spectral Stability of the OpenAI NS Blow-Up Profile
=======================================================

Extends the eta=0 computation to 2D by expanding the perturbation
in eta around the symmetry plane, using the profile equations from
the OpenAI manuscript (equations 4.7-4.14).

The profile has definite eta-parity:
  phi(X, eta) = phi_0(X) + phi_2(X)*eta^2 + ...   (even)
  U(X, eta)   = U_1(X)*eta + U_3(X)*eta^3 + ...   (odd)

The perturbation decomposes into even and odd modes in eta.
The unstable mode from eta=0 is even (it lives at eta=0).

We compute how the eta^2 correction to the profile and the
eta-coupling in the linearized operator modify the leading eigenvalue.

From the manuscript:
  Inner core equations (4.13):
    -2L * (X*phi_XX + 2*phi_X) / phi = Sq
    -2L * (X*U_XX + U_X) = Sn

  where L = 1 - 2h*eta^2, d = 1 - eta^2

  Source terms (4.9):
    Sq = -W*l - h*(1 - 2*eta*U) - Hc*(log E)_eta
    Sn = -W*D_X(U) - A*(1-2*eta*U)*U - Hc*U_eta
         - d*Pi_eta + 4*A*eta*Pi + 2*eta*D_X(Pi)

  Auxiliary quantities (4.8):
    H = sqrt(2X)*E = 2X*F,  F = phi/C
    l = D_X(log H) = 1 + X*F_X/F
    W = 1 - 2*D*eta*A_X(U) - d*partial_eta(A_X(U))
    Hc = D*eta + d*U

  Evolution equation (4.14) for angular momentum:
    (d_t + u_r d_r + u_z d_z)(q^{-h} H) = -(q^{-h-1}/L)*H*Sq
"""

import numpy as np
from scipy.linalg import eig
import json


def cheb_diff(N, L_dom):
    """Chebyshev differentiation matrices on [0, L_dom]."""
    if N == 0:
        return np.array([L_dom/2]), np.array([[0.0]]), np.array([[0.0]])
    theta = np.pi * np.arange(N + 1) / N
    xi = np.cos(theta)
    x = L_dom * (1.0 + xi) / 2.0
    c = np.ones(N + 1)
    c[0] = 2.0; c[N] = 2.0
    c *= (-1.0) ** np.arange(N + 1)
    X_mat = np.tile(xi, (N + 1, 1))
    dX = X_mat - X_mat.T
    D = np.outer(c, 1.0 / c) / (dX + np.eye(N + 1))
    D -= np.diag(D.sum(axis=1))
    D1 = D * (2.0 / L_dom)
    D2 = D1 @ D1
    return x, D1, D2


# =====================================================================
# The renormalized evolution equation for F in self-similar time s,
# expanded to second order in eta.
#
# At eta=0: (from our 1D computation, confirmed exact)
#   F_s = 2X*F'' + (4 - (1-V0)*X)*F' - (h-V0)*F
#   with V0 = 1 + 2h
#
# The eta^2 correction comes from three sources:
# 1. L = 1 - 2h*eta^2: modifies the diffusion coefficient
# 2. Profile's eta-dependence: phi_2(X) modifies the background
# 3. Coupling to U_1(X)*eta through Hc and W terms
#
# Expanding the linearized operator to O(eta^2):
#   L_2D = L_0 + eta^2 * L_correction
#
# For an even perturbation f(X, eta) = f_0(X) + f_2(X)*eta^2:
#   L_2D[f] at O(1):     L_0[f_0] = lambda * f_0
#   L_2D[f] at O(eta^2): L_0[f_2] + L_corr[f_0] = lambda * f_2
#
# This gives a coupled eigenvalue problem for (f_0, f_2).
# =====================================================================

def compute_L_correction_coefficients(h, V0):
    """
    Compute the O(eta^2) correction to the linearized operator.

    From the manuscript, the corrections come from:

    1. L = 1 - 2h*eta^2 modifies the radial diffusion:
       The diffusion operator has a prefactor 1/(2L) from (4.13).
       At O(eta^2): 1/(2L) = 1/2 * (1 + 2h*eta^2 + ...)
       This adds +2h * [diffusion terms] at O(eta^2)

    2. d = 1 - eta^2 modifies Hc = D*eta + d*U:
       At O(eta^2): the -eta^2 * U term contributes

    3. W gets an eta^2 correction from the axial strain:
       W = 1 - 2*D*eta*A_X(U) - d*partial_eta(A_X(U))
       With U = U_1*eta: A_X(U) = U_1*eta * (something)
       The D*eta*A_X(U) term is O(eta^2)

    Returns coefficients for the correction operator acting on f_0.
    """
    A = 0.5 + h
    D_exp = 0.5 - h

    # Correction from L = 1 - 2h*eta^2 on the diffusion:
    # Original operator has L in denominator via (4.13)
    # The swirl equation in self-similar time:
    #   F_s = (1/(2L)) * [...diffusion...] + transport terms
    # At O(eta^2): adds 2h * (original diffusion) = 2h * [2X*f'' + 4*f']
    # But this is the correction to the PROFILE equation.
    # For the PERTURBATION, the correction adds:
    c_diff = 2 * h  # multiplies the diffusion operator

    # Correction from the axial strain U_1:
    # U_1 = -(1+2h) from continuity (the axial compression rate)
    U1 = -(1 + 2*h)

    # Hc = D*eta + d*U = D*eta + (1-eta^2)*U_1*eta
    # At O(eta): Hc ~ (D + U1)*eta
    # The Hc*(log E)_eta term in Sq at O(eta^2):
    # Needs (log E)_eta = phi_eta/phi at eta=0 = 0 (even symmetry)
    # So this contributes at O(eta^2) through the second derivative

    # The W correction at O(eta^2):
    # W = 1 - 2*D*eta*A_X(U_1*eta) - d*partial_eta(A_X(U_1*eta))
    #   = 1 - 2*D*eta^2*A_X(U_1) - (1-eta^2)*U_1
    #   = 1 - U_1 + eta^2 * (-2*D*A_X(U_1) + U_1)
    # At O(1): W_0 = 1 - U_1 = 1 + (1+2h) = 2 + 2h
    # At O(eta^2): W_2 = -2*D*A_X(U_1) + U_1

    # With U_1 constant: A_X(U_1) = U_1
    W_2 = -2 * D_exp * U1 + U1  # = U1*(1 - 2*D) = U1*2h

    # Correction to the shift term -(h + v_0)*F:
    # v_0 at O(eta^2) changes due to the profile's eta-dependence
    # From (4.7): V_0 = (X/L)*(2*eta*U - 2*D*eta*A_X(U) - d*partial_eta(A_X(U)))
    # At eta=0: v_0 = -(1+2h) (our V0)
    # At O(eta^2): correction from L and d dependence
    # L^{-1} at O(eta^2): adds 2h*v_0_leading
    # d at O(eta^2): -1 * partial_eta(A_X(U))
    v0_correction = 2 * h * V0  # from L^{-1} expansion

    return {
        'c_diff': c_diff,
        'U1': U1,
        'W_2': W_2,
        'v0_correction': v0_correction,
    }


def compute_2d_spectrum(N, Xmax, h, V0, n_eta_modes=2):
    """
    Compute the spectrum of the 2D linearized operator by eta-expansion.

    The perturbation is expanded as:
      f(X, eta) = sum_{k=0}^{n_eta_modes-1} f_{2k}(X) * eta^{2k}

    (even modes only, since the unstable mode lives at eta=0)

    The coupled eigenvalue problem is:
      [L_0           ]   [f_0]       [f_0]
      [L_corr   L_0  ] * [f_2] = lam [f_2]
      [     ...      ]   [...]       [...]

    where L_0 is the eta=0 operator and L_corr couples f_0 to f_2.
    """
    x_pts, D1, D2 = cheb_diff(N, Xmax)
    X = x_pts

    A_exp = 0.5 + h
    D_exp = 0.5 - h

    # Build L_0 (the eta=0 operator, already validated)
    L0 = 2.0 * np.diag(X) @ D2 \
       + np.diag(4.0 - (1.0 - V0) * X) @ D1 \
       + (-(h - V0)) * np.eye(N + 1)

    # BCs for L_0
    L0[0, :] = 0.0; L0[0, 0] = 0.0  # f(Xmax) = 0
    L0[N, :] = 0.0; L0[N, N] = 0.0  # f(0) = 0

    if n_eta_modes == 1:
        # Just the 1D computation
        eigenvalues = eig(L0, right=False)
        mask = np.abs(eigenvalues) > 1e-10
        eigenvalues = eigenvalues[mask]
        idx = np.argsort(-eigenvalues.real)
        return eigenvalues[idx]

    # Build L_correction (the O(eta^2) coupling)
    corr = compute_L_correction_coefficients(h, V0)

    # The correction operator acts on f_0 and produces the
    # O(eta^2) source for f_2.
    #
    # L_corr[f_0] = c_diff * [2X*f_0'' + 4*f_0']   (from L expansion)
    #             + W_2 * l_0 * f_0                   (from W correction to Sq)
    #             + v0_correction * f_0                (from v_0 correction)
    #             + additional terms from profile eta-dependence

    L_corr = corr['c_diff'] * (2.0 * np.diag(X) @ D2 + 4.0 * D1) \
           + corr['v0_correction'] * np.eye(N + 1)

    # For f_2: it satisfies L_0[f_2] + L_corr[f_0] = lambda * f_2
    # But f_2 also gets an additional shift from the eta^2 weight:
    # The inner product weight in eta is modified.
    # For P_0(eta) = 1 and P_1(eta) = eta^2 - 1/3 (orthogonal polynomial):
    # The coupling is L_0 f_2 + L_corr f_0 = lambda f_2
    # with BCs f_2(0) = f_2(Xmax) = 0

    # Build the full block system
    dim = (N + 1) * n_eta_modes
    L_full = np.zeros((dim, dim))

    # Block (0,0): L_0 acting on f_0
    L_full[0:N+1, 0:N+1] = L0

    # Block (1,0): L_corr coupling f_0 -> f_2
    L_corr_bc = L_corr.copy()
    L_corr_bc[0, :] = 0.0  # BC
    L_corr_bc[N, :] = 0.0  # BC
    L_full[N+1:2*(N+1), 0:N+1] = L_corr_bc

    # Block (1,1): L_0 acting on f_2 (same operator, different mode)
    # The eta^2 mode gets an additional contribution from the
    # eta-Laplacian: d^2/d(eta)^2 of eta^2 = 2, contributing
    # a constant shift to the operator
    L0_shifted = L0.copy()
    # The axial diffusion d^2/dz^2 in self-similar variables contributes
    # at O(eta^2). From the manuscript, axial viscosity enters at O(q^{2h})
    # relative to radial viscosity. In the self-similar frame:
    # The eta-Laplacian for the eta^2 mode adds a stabilizing term
    # proportional to -2/(tau * q^{2D}) which in self-similar variables
    # gives a shift of -2*(axial diffusion rate)
    #
    # From Section 4.1: "axial viscosity relative to radial viscosity
    # in the tangential equations, carry an additional factor at least q^{2h}"
    # This means the axial diffusion is O(q^{2h}) smaller.
    # At leading order in the self-similar frame, the shift from the
    # eta^2 mode is -2*D^2/L (from the Laplacian in eta)
    eta2_shift = -2.0 * D_exp**2  # stabilizing contribution from axial diffusion
    for i in range(1, N):  # don't modify BC rows
        L0_shifted[i, i] += eta2_shift

    L_full[N+1:2*(N+1), N+1:2*(N+1)] = L0_shifted

    # Compute eigenvalues of the full block system
    eigenvalues = eig(L_full, right=False)
    mask = np.abs(eigenvalues) > 1e-10
    eigenvalues = eigenvalues[mask]
    idx = np.argsort(-eigenvalues.real)
    return eigenvalues[idx]


def main():
    h = 0.005
    V0 = 1.0 + 2*h
    Xmax = 40.0
    N = 120

    print("=" * 70)
    print("2D SPECTRAL COMPUTATION (eta-expansion)")
    print("=" * 70)

    # 1D baseline
    eigs_1d = compute_2d_spectrum(N, Xmax, h, V0, n_eta_modes=1)
    print(f"\n1D (eta=0 only), N={N}:")
    print(f"  Top 5 eigenvalues:")
    for i in range(min(5, len(eigs_1d))):
        lam = eigs_1d[i]
        status = "UNSTABLE" if lam.real > 0 else "stable"
        print(f"    lambda_{i} = {lam.real:+.8f}  {status}")

    # 2D with eta^2 correction
    eigs_2d = compute_2d_spectrum(N, Xmax, h, V0, n_eta_modes=2)
    print(f"\n2D (eta^0 + eta^2 modes), N={N}:")
    print(f"  Top 8 eigenvalues:")
    for i in range(min(8, len(eigs_2d))):
        lam = eigs_2d[i]
        status = "UNSTABLE" if lam.real > 0 else "stable"
        imag_str = f"  {lam.imag:+.2e}i" if abs(lam.imag) > 1e-10 else ""
        print(f"    lambda_{i} = {lam.real:+.8f}{imag_str}  {status}")

    # Check: does the unstable mode survive?
    lam0_1d = eigs_1d[0].real
    lam0_2d = eigs_2d[0].real

    print(f"\n{'='*70}")
    print(f"COMPARISON")
    print(f"{'='*70}")
    print(f"  1D leading eigenvalue: {lam0_1d:+.8f}")
    print(f"  2D leading eigenvalue: {lam0_2d:+.8f}")
    print(f"  Shift from eta coupling: {lam0_2d - lam0_1d:+.8f}")
    print()

    if lam0_2d > 0:
        print("  UNSTABLE MODE SURVIVES IN 2D.")
        print("  The eta^2 coupling does not stabilize the profile.")
        print("  The force remains structurally essential.")
    else:
        print("  UNSTABLE MODE KILLED BY 2D COUPLING.")
        print("  The eta-dependence stabilizes the profile.")
        print("  Further investigation needed.")

    # Sweep h
    print(f"\n{'='*70}")
    print(f"h-SWEEP (2D)")
    print(f"{'='*70}")
    print(f"{'h':>10s}  {'lam0 (1D)':>14s}  {'lam0 (2D)':>14s}  {'shift':>12s}  {'2D status':>10s}")

    for h_test in [0.001, 0.003, 0.005, 0.007, 0.01]:
        V0_test = 1.0 + 2*h_test
        e1 = compute_2d_spectrum(N, Xmax, h_test, V0_test, n_eta_modes=1)
        e2 = compute_2d_spectrum(N, Xmax, h_test, V0_test, n_eta_modes=2)
        l1 = e1[0].real
        l2 = e2[0].real
        status = "UNSTABLE" if l2 > 0 else "stable"
        print(f"{h_test:10.4f}  {l1:+14.8f}  {l2:+14.8f}  {l2-l1:+12.8f}  {status:>10s}")

    # Convergence in N
    print(f"\n{'='*70}")
    print(f"CONVERGENCE TEST (2D, h={h})")
    print(f"{'='*70}")
    print(f"{'N':>6s}  {'lam0 (2D)':>14s}  {'lam1 (2D)':>14s}")
    for N_test in [40, 60, 80, 100, 120, 160]:
        e = compute_2d_spectrum(N_test, Xmax, h, V0, n_eta_modes=2)
        print(f"{N_test:6d}  {e[0].real:+14.8f}  {e[1].real:+14.8f}")

    # Save results
    output = {
        'h': h, 'V0': V0,
        'lambda_0_1d': float(lam0_1d),
        'lambda_0_2d': float(lam0_2d),
        'shift': float(lam0_2d - lam0_1d),
        'top_eigenvalues_2d': [float(e.real) for e in eigs_2d[:10]],
        'survives_2d': bool(lam0_2d > 0),
    }

    with open('/home/claude/spectral_2d_results.json', 'w') as f:
        json.dump(output, f, indent=2)

    return output


if __name__ == '__main__':
    main()