"""
Coupled scalar-gauge operators for Abelian Higgs model.

Implements gauge-covariant operations for charged scalar fields:
- Covariant derivatives and Laplacian
- Covariant gradient energy
- Scalar forces with gauge coupling
- Charge density and current
- Gauss law with scalar sources
"""

import jax.numpy as jnp
from jax import jit, Array
from functools import partial

from .scalar import (
    scalar_potential_energy,
    scalar_potential_force,
    scalar_kinetic_energy_density,
)
from .gauge import divergence_3d, magnetic_energy_density


# Re-export for backward compatibility
scalar_potential = scalar_potential_energy


@partial(jit, static_argnums=(2,))
def covariant_derivative(phi: Array, links: Array, direction: int, dx: float) -> Array:
    """
    Compute covariant derivative D_i φ = (U_i(n) φ(n+i) - φ(n)) / dx.

    For a charged field, the covariant derivative includes
    the gauge link to maintain gauge invariance.

    Args:
        phi: Complex scalar field (N, N, N)
        links: Gauge links (3, N, N, N)
        direction: Direction (0=x, 1=y, 2=z) - static for JIT
        dx: Lattice spacing

    Returns:
        D_i φ: Covariant derivative (N, N, N)
    """
    phi_shifted = jnp.roll(phi, -1, axis=direction)
    U_i = links[direction]

    return (U_i * phi_shifted - phi) / dx


@jit
def covariant_gradient_squared(phi: Array, links: Array, dx: float) -> Array:
    """
    Compute |D_i φ|² = sum_i |D_i φ|² at each site.

    This is the kinetic energy density for the scalar field
    in the presence of gauge fields.

    Args:
        phi: Complex scalar field (N, N, N)
        links: Gauge links (3, N, N, N)
        dx: Lattice spacing

    Returns:
        |∇_cov φ|²: Covariant gradient squared (N, N, N)
    """
    D0 = covariant_derivative(phi, links, 0, dx)
    D1 = covariant_derivative(phi, links, 1, dx)
    D2 = covariant_derivative(phi, links, 2, dx)

    return jnp.abs(D0) ** 2 + jnp.abs(D1) ** 2 + jnp.abs(D2) ** 2


@jit
def covariant_laplacian(phi: Array, links: Array, dx: float) -> Array:
    """
    Compute covariant Laplacian Δ_cov φ for gauge-coupled scalar field.

    Δ_cov φ = Σ_i [U_i(n) φ(n+i) + U_i†(n-i) φ(n-i) - 2φ(n)] / dx²

    This is the gauge-covariant kinetic term in the scalar equation of motion.

    Args:
        phi: Complex scalar field (N, N, N)
        links: Gauge links U_i (3, N, N, N)
        dx: Lattice spacing

    Returns:
        Covariant Laplacian (N, N, N)
    """
    laplacian = jnp.zeros_like(phi)

    # Direction 0 (x)
    laplacian = laplacian + (
        links[0] * jnp.roll(phi, -1, axis=0)
        + jnp.roll(jnp.conj(links[0]), 1, axis=0) * jnp.roll(phi, 1, axis=0)
        - 2 * phi
    )

    # Direction 1 (y)
    laplacian = laplacian + (
        links[1] * jnp.roll(phi, -1, axis=1)
        + jnp.roll(jnp.conj(links[1]), 1, axis=1) * jnp.roll(phi, 1, axis=1)
        - 2 * phi
    )

    # Direction 2 (z)
    laplacian = laplacian + (
        links[2] * jnp.roll(phi, -1, axis=2)
        + jnp.roll(jnp.conj(links[2]), 1, axis=2) * jnp.roll(phi, 1, axis=2)
        - 2 * phi
    )

    return laplacian / (dx**2)


@jit
def scalar_force(
    phi: Array,
    links: Array,
    m: float,
    lambda_: float,
    dx: float,
) -> Array:
    """
    Compute force on scalar field: F_φ = -δH/δφ*.

    F_φ = Δ_cov φ - m² φ - λ |φ|² φ

    Includes:
    - Covariant Laplacian (gauge-covariant kinetic term)
    - Mass term
    - Self-interaction term

    Args:
        phi: Complex scalar field (N, N, N)
        links: Gauge links (3, N, N, N)
        m: Mass parameter
        lambda_: Self-coupling
        dx: Lattice spacing

    Returns:
        F_φ: Force on phi (N, N, N)
    """
    laplacian = covariant_laplacian(phi, links, dx)
    potential_force = scalar_potential_force(phi, m, lambda_)
    return laplacian + potential_force


@jit
def scalar_charge_density(phi: Array, pi: Array, g: float) -> Array:
    """
    Compute scalar field charge density ρ = g * Im(φ* π).

    This appears as a source in the Gauss law: ∇·E = -g ρ

    Args:
        phi: Complex scalar field (N, N, N)
        pi: Conjugate momentum (N, N, N)
        g: Gauge coupling

    Returns:
        ρ: Charge density (N, N, N)
    """
    return g * jnp.imag(jnp.conj(phi) * pi)


@partial(jit, static_argnums=(2,))
def scalar_current_density(
    phi: Array, links: Array, direction: int, g: float, dx: float
) -> Array:
    """
    Compute scalar field current density in direction i.

    j_i = (g/2) Im[φ*(n) U_i(n) φ(n+i)]

    This is the conserved current for U(1) gauge symmetry.

    Args:
        phi: Complex scalar field (N, N, N)
        links: Gauge links (3, N, N, N)
        direction: Direction (0=x, 1=y, 2=z) - static for JIT
        g: Gauge coupling
        dx: Lattice spacing

    Returns:
        j_i: Current density in direction i (N, N, N)
    """
    phi_shifted = jnp.roll(phi, -1, axis=direction)
    U_i = links[direction]

    return (g / (2 * dx)) * jnp.imag(jnp.conj(phi) * U_i * phi_shifted)


@jit
def gauss_constraint_with_source(
    links: Array,
    E: Array,
    phi: Array,
    pi: Array,
    g: float,
    dx: float,
) -> Array:
    """
    Compute Gauss constraint with scalar field source.

    Gauss law: ∇·E = -g ρ
    where ρ = Im(φ* π) is the charge density.

    Args:
        links: Gauge links (not used, kept for interface consistency)
        E: Electric field (3, N, N, N)
        phi: Complex scalar field (N, N, N)
        pi: Conjugate momentum (N, N, N)
        g: Gauge coupling
        dx: Lattice spacing

    Returns:
        G: Gauss constraint violation (N, N, N)
    """
    div_E = divergence_3d(E, dx)
    rho = scalar_charge_density(phi, pi, g)

    # Gauss constraint: ∇·E + g ρ = 0
    return div_E + rho


@jit
def gauss_violation_norm_with_source(
    links: Array,
    E: Array,
    phi: Array,
    pi: Array,
    g: float,
    dx: float,
) -> Array:
    """
    Compute RMS norm of Gauss constraint violation.

    Returns sqrt(mean(|∇·E + g ρ|²))

    Args:
        links: Gauge links (not used, kept for interface consistency)
        E: Electric field (3, N, N, N)
        phi: Complex scalar field (N, N, N)
        pi: Conjugate momentum (N, N, N)
        g: Gauge coupling
        dx: Lattice spacing

    Returns:
        ||G||_2: RMS Gauss violation (scalar)
    """
    G = gauss_constraint_with_source(links, E, phi, pi, g, dx)
    return jnp.sqrt(jnp.mean(G**2))


def coupled_energy(
    phi: Array,
    pi: Array,
    links: Array,
    E: Array,
    dx: float,
    m: float,
    lambda_: float,
    g: float,
) -> dict:
    """
    Compute total energy of coupled scalar-gauge system.

    H = ∫ d³x [
        (1/2)|π|² +                     (scalar kinetic)
        (1/2)|D_i φ|² +                 (scalar gradient)
        V(φ) +                          (scalar potential)
        (1/2)E_i² +                     (electric energy)
        (1/2)B_i²                       (magnetic energy)
    ]

    Args:
        phi: Scalar field (N, N, N)
        pi: Conjugate momentum (N, N, N)
        links: Gauge links (3, N, N, N)
        E: Electric field (3, N, N, N)
        dx: Lattice spacing
        m: Scalar mass
        lambda_: Self-coupling
        g: Gauge coupling

    Returns:
        dict with energy components and total
    """
    dV = dx**3

    # Scalar kinetic energy
    T_scalar = scalar_kinetic_energy_density(pi)
    E_scalar_kinetic = float(jnp.sum(T_scalar) * dV)

    # Scalar gradient energy (covariant)
    grad_sq = covariant_gradient_squared(phi, links, dx)
    E_scalar_gradient = float(0.5 * jnp.sum(grad_sq) * dV)

    # Scalar potential energy
    V = scalar_potential_energy(phi, m, lambda_)
    E_scalar_potential = float(jnp.sum(V) * dV)

    # Electric energy
    E_electric = float(0.5 * jnp.sum(E**2) * dV)

    # Magnetic energy (now returns total, not density)
    E_magnetic = float(magnetic_energy_density(links, dx, g))

    E_total = (
        E_scalar_kinetic
        + E_scalar_gradient
        + E_scalar_potential
        + E_electric
        + E_magnetic
    )

    return {
        "scalar_kinetic": E_scalar_kinetic,
        "scalar_gradient": E_scalar_gradient,
        "scalar_potential": E_scalar_potential,
        "electric": E_electric,
        "magnetic": E_magnetic,
        "total": E_total,
    }
