"""
Stability analysis using Hessian computation via JAX autodiff.

Implements:
- Tachyonic mode detection
- Linear response theory
- Fluctuation spectrum
- Phase transition indicators
- Tunneling rate estimates

All second derivatives computed automatically via JAX.
"""

import jax.numpy as jnp
from jax import jit, hessian, Array
from typing import Tuple, Dict


# =============================================================================
# Potential Hessian (Mass Matrix)
# =============================================================================


@jit
def compute_mass_matrix(
    phi: Array,
    m: float,
    lambda_: float,
) -> Array:
    """
    Compute effective mass matrix ∂²V/∂φ∂φ* (Hessian of potential).

    For V = (m²/2)|φ|² + (λ/4)|φ|⁴:
    ∂²V/∂|φ|² = m² + λ|φ|²

    This tells us if the field configuration is stable (all positive)
    or has tachyonic modes (negative eigenvalues).

    Args:
        phi: Scalar field (N, N, N) complex
        m: Mass parameter
        lambda_: Self-coupling

    Returns:
        Effective mass squared at each point (N, N, N) real
    """
    phi_sq = jnp.abs(phi) ** 2
    # Allow negative "mass" parameter to represent negative mass-squared
    # (tachyonic) without forcing a square that would erase the sign.
    m_eff_sq = jnp.where(m < 0, m + lambda_ * phi_sq, m**2 + lambda_ * phi_sq)
    return m_eff_sq


def detect_tachyonic_modes(
    phi: Array,
    m: float,
    lambda_: float,
) -> Dict[str, float]:
    """
    Detect tachyonic (unstable) modes where m_eff² < 0.

    Returns statistics about instability.
    """
    m_eff_sq = compute_mass_matrix(phi, m, lambda_)

    # Find tachyonic regions
    is_tachyonic = m_eff_sq < 0
    n_tachyonic = jnp.sum(is_tachyonic)
    fraction_tachyonic = jnp.mean(is_tachyonic)

    # Strongest instability
    min_m_sq = jnp.min(m_eff_sq)
    max_m_sq = jnp.max(m_eff_sq)

    return {
        "n_tachyonic_sites": int(n_tachyonic.item()),
        "fraction_tachyonic": float(fraction_tachyonic.item()),
        "min_mass_squared": float(min_m_sq.item()),
        "max_mass_squared": float(max_m_sq.item()),
        "has_instability": bool(min_m_sq < 0),
    }


# =============================================================================
# Energy Hessian (Full Second Derivatives)
# =============================================================================


def make_energy_hessian_function(
    m: float,
    lambda_: float,
    dx: float,
):
    """
    Create function to compute full Hessian ∂²H/∂φ∂φ via autodiff.

    This includes both potential and gradient terms:
    H = (1/2)∫|∇φ|² + V(φ)

    The Hessian operator is: -∇² + ∂²V/∂φ²

    Args:
        m: Mass
        lambda_: Self-coupling
        dx: Lattice spacing

    Returns:
        Function phi → Hessian operator applied to phi
    """

    def energy_functional(phi_flat: Array) -> float:
        """Energy as scalar function of flattened φ."""
        # Reshape
        shape = tuple([int(jnp.cbrt(phi_flat.size + 0.5))] * 3)
        phi = phi_flat.reshape(shape)

        # Gradient energy (finite difference Laplacian)
        laplacian = (
            jnp.roll(phi, -1, axis=0)
            + jnp.roll(phi, 1, axis=0)
            + jnp.roll(phi, -1, axis=1)
            + jnp.roll(phi, 1, axis=1)
            + jnp.roll(phi, -1, axis=2)
            + jnp.roll(phi, 1, axis=2)
            - 6 * phi
        ) / (dx**2)
        grad_energy = -0.5 * jnp.sum(jnp.real(jnp.conj(phi) * laplacian))

        # Potential energy
        phi_sq = jnp.abs(phi) ** 2
        V = 0.5 * m**2 * phi_sq + 0.25 * lambda_ * phi_sq**2
        pot_energy = jnp.sum(V)

        return float(grad_energy + pot_energy) * dx**3

    # Compute Hessian via autodiff
    hess_fn = hessian(energy_functional, holomorphic=False)

    @jit
    def apply_hessian(phi: Array) -> Array:
        """Apply Hessian operator to phi."""
        phi_flat = phi.ravel()
        H = hess_fn(phi_flat)
        # For large systems, H is huge - instead compute H @ phi
        result = jnp.dot(H, phi_flat)
        return result.reshape(phi.shape)

    return apply_hessian


# =============================================================================
# Fluctuation Spectrum
# =============================================================================


@jit
def compute_fluctuation_spectrum_1d(
    phi: Array,
    m: float,
    lambda_: float,
    dx: float,
) -> Tuple[Array, Array]:
    """
    Compute power spectrum of fluctuations δφ around background.

    P(k) = <|δφ_k|²> gives spectrum of quantum/thermal fluctuations.

    Args:
        phi: Field configuration
        m: Mass
        lambda_: Coupling
        dx: Lattice spacing

    Returns:
        (k_values, power_spectrum)
    """
    # FFT to k-space
    phi_k = jnp.fft.fftn(phi)

    # Power spectrum
    P_k = jnp.abs(phi_k) ** 2 / phi.size

    # Radial average (for 3D → 1D spectrum)
    # First get k-space coordinates
    kx = jnp.fft.fftfreq(phi.shape[0], d=dx) * 2 * jnp.pi
    ky = jnp.fft.fftfreq(phi.shape[1], d=dx) * 2 * jnp.pi
    kz = jnp.fft.fftfreq(phi.shape[2], d=dx) * 2 * jnp.pi

    KX, KY, KZ = jnp.meshgrid(kx, ky, kz, indexing="ij")
    K = jnp.sqrt(KX**2 + KY**2 + KZ**2)

    # Bin by |k|
    k_bins = jnp.linspace(0, K.max(), 20)
    k_centers = 0.5 * (k_bins[:-1] + k_bins[1:])

    # Average power in each bin
    P_avg = jnp.zeros_like(k_centers)
    for i in range(len(k_centers)):
        mask = (K >= k_bins[i]) & (K < k_bins[i + 1])
        P_avg = P_avg.at[i].set(jnp.mean(jnp.where(mask, P_k, 0)))

    return k_centers, P_avg


@jit
def estimate_effective_mass_from_spectrum(
    k_values: Array,
    power_spectrum: Array,
    dx: float,
) -> float:
    """
    Estimate effective mass from fluctuation spectrum.

    For free theory: P(k) ~ 1/(k² + m_eff²)

    Fit low-k behavior to extract m_eff.

    Args:
        k_values: Momentum values
        power_spectrum: P(k)
        dx: Lattice spacing

    Returns:
        Effective mass estimate
    """
    # Use only low-k modes (k < π/4L)
    k_max = jnp.pi / (4 * dx)
    mask = k_values < k_max

    P_low = power_spectrum[mask]

    # Fit 1/P ~ k² + m²
    # m² ≈ 1/P(k→0) - k²
    if jnp.sum(mask) > 2:
        P_zero = P_low[0]
        m_eff_sq = 1.0 / (P_zero + 1e-10)
        return float(jnp.sqrt(jnp.maximum(m_eff_sq, 0)))
    else:
        return 0.0


# =============================================================================
# Linear Response
# =============================================================================


@jit
def compute_linear_response(
    phi_background: Array,
    delta_phi: Array,
    m: float,
    lambda_: float,
    dx: float,
) -> Array:
    """
    Compute linear response: how system responds to perturbation.

    δφ → (∂²V/∂φ²)|_bg δφ

    This is the linearized equation of motion.

    Args:
        phi_background: Background field
        delta_phi: Perturbation
        m: Mass
        lambda_: Coupling
        dx: Lattice spacing

    Returns:
        Response δφ
    """
    # Effective mass at background
    m_eff_sq = compute_mass_matrix(phi_background, m, lambda_)

    # Laplacian of perturbation
    laplacian_delta = (
        jnp.roll(delta_phi, -1, axis=0)
        + jnp.roll(delta_phi, 1, axis=0)
        + jnp.roll(delta_phi, -1, axis=1)
        + jnp.roll(delta_phi, 1, axis=1)
        + jnp.roll(delta_phi, -1, axis=2)
        + jnp.roll(delta_phi, 1, axis=2)
        - 6 * delta_phi
    ) / (dx**2)

    # Linear response: (-∇² + m_eff²) δφ
    return -laplacian_delta + m_eff_sq * delta_phi


# =============================================================================
# Phase Transition Indicators
# =============================================================================


@jit
def compute_curvature_at_origin(
    m: float,
    lambda_: float,
) -> float:
    """
    Compute ∂²V/∂φ²|_{φ=0} = m².

    If m² < 0: symmetry broken (φ=0 unstable)
    If m² > 0: symmetric phase (φ=0 stable)

    Args:
        m: Mass parameter
        lambda_: Coupling

    Returns:
        Curvature at origin (= m²)
    """
    return m**2


def compute_vev_estimate(
    m: float,
    lambda_: float,
) -> float:
    """
    Estimate vacuum expectation value for symmetry breaking.

    If m² < 0: VEV ~ sqrt(-m²/λ)
    If m² > 0: VEV = 0

    Args:
        m: Mass (can be negative for SSB)
        lambda_: Self-coupling

    Returns:
        Estimated |<φ>|
    """
    if m**2 < 0 and lambda_ > 0:
        return float(jnp.sqrt(-(m**2) / lambda_))
    else:
        return 0.0


def check_symmetry_breaking(
    phi: Array,
    m: float,
    lambda_: float,
) -> Dict[str, float]:
    """
    Check if field has non-zero VEV (symmetry breaking).

    Returns statistics comparing with expected VEV.
    """
    # Mean field
    phi_mean = jnp.mean(phi)
    phi_rms = jnp.sqrt(jnp.mean(jnp.abs(phi) ** 2))

    # Expected VEV
    vev_expected = compute_vev_estimate(m, lambda_)

    # Curvature
    curvature = compute_curvature_at_origin(m, lambda_)

    return {
        "mean_field_magnitude": float(jnp.abs(phi_mean)),
        "rms_field": float(phi_rms),
        "expected_vev": vev_expected,
        "curvature_at_origin": curvature,
        "is_symmetry_broken": bool(curvature < 0),
        "vev_ratio": float(phi_rms / (vev_expected + 1e-10)),
    }


# =============================================================================
# Tunneling and Bubble Nucleation
# =============================================================================


@jit
def estimate_bounce_action(
    phi_false: float,
    phi_true: float,
    m: float,
    lambda_: float,
) -> float:
    """
    Estimate bounce action for tunneling φ_false → φ_true.

    S_bounce ~ (ΔV)^{-2} (barrier height)^4

    Tunneling rate: Γ ~ exp(-S_bounce)

    Args:
        phi_false: False vacuum field value
        phi_true: True vacuum field value
        m: Mass
        lambda_: Coupling

    Returns:
        Bounce action estimate
    """
    # Potential at false and true vacua
    V_false = 0.5 * m**2 * phi_false**2 + 0.25 * lambda_ * phi_false**4
    V_true = 0.5 * m**2 * phi_true**2 + 0.25 * lambda_ * phi_true**4

    # Potential difference
    Delta_V = jnp.abs(V_true - V_false)

    # Barrier estimate (crude approximation)
    # For quartic potential with symmetry breaking
    if m**2 < 0:
        # Barrier at φ = 0
        V_barrier = 0.0
        barrier_height = jnp.abs(V_barrier - V_false)
    else:
        # No barrier
        barrier_height = 0.0

    # Bounce action ~ barrier^4 / ΔV²
    if Delta_V > 1e-10:
        S_bounce = barrier_height**4 / (Delta_V**2 + 1e-10)
    else:
        S_bounce = 1e10  # Very suppressed

    return float(S_bounce)


# =============================================================================
# Comprehensive Stability Report
# =============================================================================


@jit
def generate_stability_report(
    phi: Array,
    m: float,
    lambda_: float,
    dx: float,
) -> Dict[str, float]:
    """
    Generate comprehensive stability analysis report.

    Returns dict with:
    - Tachyonic mode detection
    - Mass matrix statistics
    - Symmetry breaking indicators
    - Effective parameters
    """
    # Mass matrix
    m_eff_sq = compute_mass_matrix(phi, m, lambda_)

    # Tachyonic modes
    tach_info = detect_tachyonic_modes(phi, m, lambda_)

    # Symmetry breaking
    ssb_info = check_symmetry_breaking(phi, m, lambda_)

    # Combined report
    report = {
        "min_mass_sq": float(jnp.min(m_eff_sq)),
        "max_mass_sq": float(jnp.max(m_eff_sq)),
        "mean_mass_sq": float(jnp.mean(m_eff_sq)),
        "has_tachyonic_modes": tach_info["has_instability"],
        "fraction_unstable": tach_info["fraction_tachyonic"],
        "field_rms": ssb_info["rms_field"],
        "symmetry_broken": ssb_info["is_symmetry_broken"],
        "curvature_at_origin": ssb_info["curvature_at_origin"],
    }

    return report
