"""
Scalar field operators.

Implements pure scalar field operations including:
- Laplacian and gradient operators
- Potential energy and forces
- Kinetic energy density

Note: Core stencil operators (laplacian, gradient_squared) are provided by
`jaxlatt.operators.stencil`. This module re-exports them for backward
compatibility and adds scalar-field-specific operations.
"""

import jax.numpy as jnp
from jax import jit, grad, Array

# Re-export from stencil module for backward compatibility
from jaxlatt.operators.stencil import laplacian as scalar_laplacian
from jaxlatt.operators.stencil import gradient_squared


def scalar_potential_energy(phi: Array, m: float, lambda_: float) -> Array:
    """
    Compute potential energy V(φ) for scalar field.

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

    Args:
        phi: Complex scalar field
        m: Mass parameter
        lambda_: Self-coupling

    Returns:
        Potential energy at each point (same shape as phi)
    """
    phi_sq = jnp.abs(phi) ** 2
    return 0.5 * m**2 * phi_sq + 0.25 * lambda_ * phi_sq**2


@jit
def scalar_potential_force(phi: Array, m: float, lambda_: float) -> Array:
    """
    Compute potential force -dV/dφ* for complex scalar field.

    V = (m²/2)|φ|² + (λ/4)|φ|⁴
    Force = -dV/dφ* = -m²φ - λ|φ|²φ

    Uses JAX autodiff to ensure correctness.

    Args:
        phi: Complex scalar field
        m: Mass parameter
        lambda_: Self-coupling

    Returns:
        Potential force (same shape as phi)
    """

    def total_potential(phi_array):
        return jnp.sum(jnp.real(scalar_potential_energy(phi_array, m, lambda_)))

    # Gradient with respect to conjugate for complex fields
    grad_fn = grad(total_potential, holomorphic=False)
    dV_dphi_conj = grad_fn(phi)

    return -dV_dphi_conj


@jit
def scalar_kinetic_energy_density(pi: Array) -> Array:
    """
    Kinetic energy density (1/2)|π|².

    Args:
        pi: Conjugate momentum (can be complex)

    Returns:
        T: Kinetic energy density (real-valued)
    """
    return 0.5 * jnp.abs(pi) ** 2
