"""3-D Poisson solver for linear and nonlinear pressure perturbations.

The governing equation (anelastic/Boussinesq, Trapp 2013) is:

    ∇²p' = F_lin + F_spin + F_splat + F_buoy

where each forcing term is solved independently so the contributions
can be compared and displayed separately.

Solver:  2-D FFT in the (periodic) horizontal  +  vectorized Thomas
         algorithm for the resulting 1-D BVP in z with Neumann BCs.
"""

from __future__ import annotations

import numpy as np

from .sounding import G


# --------------------------------------------------------------------------
# Grid derivative helpers
# --------------------------------------------------------------------------

def _grad_x(f: np.ndarray, dx: float) -> np.ndarray:
    """∂f/∂x with periodic x (centered differences)."""
    return (np.roll(f, -1, axis=0) - np.roll(f, 1, axis=0)) / (2.0 * dx)


def _grad_y(f: np.ndarray, dx: float) -> np.ndarray:
    """∂f/∂y with periodic y (centered differences)."""
    return (np.roll(f, -1, axis=1) - np.roll(f, 1, axis=1)) / (2.0 * dx)


def _grad_z(f: np.ndarray, dz: float) -> np.ndarray:
    """∂f/∂z with one-sided differences at boundaries."""
    out = np.empty_like(f)
    out[:, :, 1:-1] = (f[:, :, 2:] - f[:, :, :-2]) / (2.0 * dz)
    out[:, :, 0] = (f[:, :, 1] - f[:, :, 0]) / dz
    out[:, :, -1] = (f[:, :, -1] - f[:, :, -2]) / dz
    return out


# --------------------------------------------------------------------------
# Forcing term constructors
# --------------------------------------------------------------------------

def forcing_linear(
    rho0: np.ndarray,
    dudz_env: np.ndarray,
    dvdz_env: np.ndarray,
    w3d: np.ndarray,
    dx: float,
) -> np.ndarray:
    """Linear (shear-interaction) forcing.

    F_lin = -2ρ₀ [(∂U/∂z)(∂w'/∂x) + (∂V/∂z)(∂w'/∂y)]
    """
    dwdx = _grad_x(w3d, dx)
    dwdy = _grad_y(w3d, dx)
    rho0_3d = rho0[np.newaxis, np.newaxis, :]  # broadcast to (1,1,Nz)
    dUdz = dudz_env[np.newaxis, np.newaxis, :]
    dVdz = dvdz_env[np.newaxis, np.newaxis, :]
    return -2.0 * rho0_3d * (dUdz * dwdx + dVdz * dwdy)


def _strain_and_rotation(
    u3d: np.ndarray,
    v3d: np.ndarray,
    w3d: np.ndarray,
    env_u: np.ndarray,
    env_v: np.ndarray,
    dx: float,
    dz: float,
) -> tuple[np.ndarray, np.ndarray]:
    """Compute ΣΣ Sᵢⱼ² and ΣΣ Rᵢⱼ² from the perturbation velocity field."""
    # Perturbation winds
    up = u3d - env_u[np.newaxis, np.newaxis, :]
    vp = v3d - env_v[np.newaxis, np.newaxis, :]
    wp = w3d  # no environmental w

    # Velocity gradients of perturbation
    dudx = _grad_x(up, dx)
    dudy = _grad_y(up, dx)
    dudz = _grad_z(up, dz)
    dvdx = _grad_x(vp, dx)
    dvdy = _grad_y(vp, dx)
    dvdz = _grad_z(vp, dz)
    dwdx = _grad_x(wp, dx)
    dwdy = _grad_y(wp, dx)
    dwdz = _grad_z(wp, dz)

    # Strain rate tensor Sᵢⱼ = ½(∂uᵢ/∂xⱼ + ∂uⱼ/∂xᵢ)
    S11 = dudx
    S22 = dvdy
    S33 = dwdz
    S12 = 0.5 * (dudy + dvdx)
    S13 = 0.5 * (dudz + dwdx)
    S23 = 0.5 * (dvdz + dwdy)
    S2 = S11**2 + S22**2 + S33**2 + 2.0 * (S12**2 + S13**2 + S23**2)

    # Rotation rate tensor ΣΣ Rᵢⱼ² = ½|ω|² where ω = ∇×u'
    # ζ_x = ∂w/∂y − ∂v/∂z,  ζ_y = ∂u/∂z − ∂w/∂x,  ζ_z = ∂v/∂x − ∂u/∂y
    zx = dwdy - dvdz
    zy = dudz - dwdx
    zz = dvdx - dudy
    R2 = 0.5 * (zx**2 + zy**2 + zz**2)  # = ΣΣ Rᵢⱼ²

    return S2, R2


def forcing_splat(
    rho0: np.ndarray,
    u3d: np.ndarray,
    v3d: np.ndarray,
    w3d: np.ndarray,
    env_u: np.ndarray,
    env_v: np.ndarray,
    dx: float,
    dz: float,
) -> np.ndarray:
    """Nonlinear splat (deformation) forcing  F_splat = -ρ₀ ΣΣ Sᵢⱼ²."""
    S2, _ = _strain_and_rotation(u3d, v3d, w3d, env_u, env_v, dx, dz)
    return -rho0[np.newaxis, np.newaxis, :] * S2


def forcing_spin(
    rho0: np.ndarray,
    u3d: np.ndarray,
    v3d: np.ndarray,
    w3d: np.ndarray,
    env_u: np.ndarray,
    env_v: np.ndarray,
    dx: float,
    dz: float,
) -> np.ndarray:
    """Nonlinear spin (rotation) forcing  F_spin = +ρ₀ ΣΣ Rᵢⱼ² = ρ₀/2 |ω'|²."""
    _, R2 = _strain_and_rotation(u3d, v3d, w3d, env_u, env_v, dx, dz)
    return rho0[np.newaxis, np.newaxis, :] * R2


def forcing_buoyancy(
    rho0: np.ndarray,
    theta0: np.ndarray,
    theta_prime3d: np.ndarray,
    dz: float,
) -> np.ndarray:
    """Buoyancy pressure forcing  F_buoy = -ρ₀ (g/θ₀) ∂θ'/∂z."""
    dthp_dz = _grad_z(theta_prime3d, dz)
    return -rho0[np.newaxis, np.newaxis, :] * (G / theta0[np.newaxis, np.newaxis, :]) * dthp_dz


# --------------------------------------------------------------------------
# Vectorized Thomas algorithm (TDMA)
# --------------------------------------------------------------------------

def _tdma_batch(
    a: np.ndarray,
    b: np.ndarray,
    c: np.ndarray,
    d: np.ndarray,
) -> np.ndarray:
    """Thomas algorithm for M independent tridiagonal systems of size N.

    a, b, c, d : (M, N) arrays (complex or real).
    a[:,0] and c[:,-1] are unused (boundary rows).
    Returns x : (M, N).
    """
    _, N = d.shape
    cp = np.zeros_like(d)
    dp = np.zeros_like(d)

    # Forward sweep
    w = b[:, 0].copy()
    cp[:, 0] = c[:, 0] / w
    dp[:, 0] = d[:, 0] / w

    for k in range(1, N):
        w = b[:, k] - a[:, k] * cp[:, k - 1]
        cp[:, k] = c[:, k] / w
        dp[:, k] = (d[:, k] - a[:, k] * dp[:, k - 1]) / w

    # Back substitution
    x = np.zeros_like(d)
    x[:, N - 1] = dp[:, N - 1]
    for k in range(N - 2, -1, -1):
        x[:, k] = dp[:, k] - cp[:, k] * x[:, k + 1]

    return x


# --------------------------------------------------------------------------
# 3-D Poisson solver
# --------------------------------------------------------------------------

def solve_poisson_3d(F: np.ndarray, dx: float, dz: float) -> np.ndarray:
    """Solve ∇²p = F on a periodic (x,y) × Neumann-z domain.

    F  : (Nx, Ny, Nz) real forcing array.
    dx : horizontal grid spacing (same in x and y) in meters.
    dz : vertical grid spacing in meters.

    Returns p : (Nx, Ny, Nz) real array with zero mean.
    """
    Nx, Ny, Nz = F.shape

    # --- 2-D real FFT in the horizontal ---
    F_hat = np.fft.rfft2(F, axes=(0, 1))  # (Nx, Ny//2+1, Nz)
    Nkx = Nx
    Nky = Ny // 2 + 1
    Nmodes = Nkx * Nky

    kx = np.fft.fftfreq(Nx, d=dx) * (2.0 * np.pi)  # (Nx,)
    ky = np.fft.rfftfreq(Ny, d=dx) * (2.0 * np.pi)  # (Nky,)
    KX, KY = np.meshgrid(kx, ky, indexing="ij")     # (Nx, Nky)
    lam2 = (KX**2 + KY**2).reshape(Nmodes)           # (Nmodes,)

    # Reshape F_hat → (Nmodes, Nz) for batch solve
    rhs = F_hat.reshape(Nmodes, Nz) * dz**2          # (Nmodes, Nz), complex

    # --- Build tridiagonal coefficients (Nmodes, Nz) ---
    main_diag = -(2.0 + lam2[:, np.newaxis] * dz**2) * np.ones((Nmodes, Nz), dtype=complex)
    upper_diag = np.ones((Nmodes, Nz), dtype=complex)
    lower_diag = np.ones((Nmodes, Nz), dtype=complex)

    # Neumann BC at z=0: ghost-point reflection → super-diagonal = 2 at row 0
    upper_diag[:, 0] = 2.0
    upper_diag[:, -1] = 0.0   # unused

    # Neumann BC at z=top: sub-diagonal = 2 at last row
    lower_diag[:, -1] = 2.0
    lower_diag[:, 0] = 0.0    # unused

    # --- Handle the (kx=0, ky=0) mode separately ---
    # ∇²p = F with Neumann BCs is solvable only if ∫F dz = 0.
    # Enforce solvability by subtracting mean, then fix gauge p[0]=0.
    m00 = 0
    rhs[m00] -= rhs[m00].mean()
    main_diag[m00, 0] = 1.0
    upper_diag[m00, 0] = 0.0
    lower_diag[m00, 0] = 0.0
    rhs[m00, 0] = 0.0

    # --- Batch TDMA ---
    P_hat_2d = _tdma_batch(lower_diag, main_diag, upper_diag, rhs)  # (Nmodes, Nz)

    # --- Inverse 2-D FFT ---
    P_hat = P_hat_2d.reshape(Nkx, Nky, Nz)
    p = np.fft.irfft2(P_hat, s=(Nx, Ny), axes=(0, 1))   # (Nx, Ny, Nz)
    p -= p.mean()
    return p


# --------------------------------------------------------------------------
# Acceleration diagnostics
# --------------------------------------------------------------------------

def pressure_accelerations(
    p_lin: np.ndarray,
    p_spin: np.ndarray,
    p_splat: np.ndarray,
    p_buoy: np.ndarray,
    rho0: np.ndarray,
    dz: float,
) -> dict:
    """Compute vertical acceleration from each pressure component.

    a = -(1/ρ₀) ∂p'/∂z   at each (x, y, z).
    """
    def _accel(p):
        dpdz = _grad_z(p, dz)
        return -dpdz / rho0[np.newaxis, np.newaxis, :]

    return {
        "a_lin": _accel(p_lin),
        "a_spin": _accel(p_spin),
        "a_splat": _accel(p_splat),
        "a_buoy": _accel(p_buoy),
        "a_total": _accel(p_lin + p_spin + p_splat + p_buoy),
    }