"""
Gauge field operators for U(1) lattice gauge theory.

Implements:
- Plaquettes and staples
- Magnetic field extraction
- Electric field operations
- Gauge forces and constraints
"""

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


@partial(jit, static_argnums=(1, 2))
def plaquette(links: Array, i: int, j: int) -> Array:
    """
    Compute spatial plaquette U_ij for all lattice sites.

    U_ij(n) = U_i(n) * U_j(n+i) * U_i^*(n+j) * U_j^*(n)

    For U(1), this is a complex number with |U_ij| ≤ 1.
    In continuum limit: U_ij ≈ exp(-i g dx² F_ij)

    Args:
        links: Link variables (3, N, N, N) complex
        i: First spatial direction (0, 1, or 2) - static for JIT
        j: Second spatial direction (0, 1, or 2) - static for JIT

    Returns:
        Plaquette values at all sites (N, N, N) complex
    """
    Ui = links[i]
    Uj = links[j]

    # U_j(n+i): shift j-links forward in i direction
    Uj_shift_i = jnp.roll(Uj, shift=-1, axis=i)

    # U_i(n+j): shift i-links forward in j direction
    Ui_shift_j = jnp.roll(Ui, shift=-1, axis=j)

    # Plaquette: forward-forward-back-back path
    U_plaq = Ui * Uj_shift_i * jnp.conj(Ui_shift_j) * jnp.conj(Uj)

    return U_plaq


@jit
def all_plaquettes(links: Array) -> Tuple[Array, Array, Array]:
    """
    Compute all three spatial plaquettes (xy, yz, zx).

    Args:
        links: Link variables (3, N, N, N)

    Returns:
        Tuple of (U_01, U_12, U_20) plaquettes, each (N, N, N) complex
    """
    U_xy = plaquette(links, 0, 1)
    U_yz = plaquette(links, 1, 2)
    U_zx = plaquette(links, 2, 0)

    return U_xy, U_yz, U_zx


@jit
def magnetic_field(links: Array, dx: float, g: float) -> Array:
    """
    Compute magnetic field B_k from plaquettes.

    B_k = (1 / (2 g dx²)) * ε_kij * Im(U_ij)

    where ε_kij is Levi-Civita symbol.

    Args:
        links: Link variables (3, N, N, N)
        dx: Lattice spacing
        g: Gauge coupling

    Returns:
        Magnetic field components (3, N, N, N) real
    """
    U_xy, U_yz, U_zx = all_plaquettes(links)

    prefactor = 1.0 / (2.0 * g * dx * dx)

    # B_x from yz plaquette, B_y from zx, B_z from xy
    Bx = prefactor * jnp.imag(U_yz)
    By = prefactor * jnp.imag(U_zx)
    Bz = prefactor * jnp.imag(U_xy)

    B = jnp.stack([Bx, By, Bz], axis=0)

    return B.astype(jnp.float32)


@partial(jit, static_argnums=(1, 2))
def _staple_contribution(links: Array, i: int, j: int) -> Array:
    """Compute staple contribution from perpendicular direction j for link i."""
    Ui = links[i]
    Uj = links[j]

    # Forward staple: goes up in j, back in i, down in j
    Uj_shift_i = jnp.roll(Uj, shift=-1, axis=i)
    Ui_shift_j = jnp.roll(Ui, shift=-1, axis=j)
    forward = Uj_shift_i * jnp.conj(Ui_shift_j) * jnp.conj(Uj)

    # Backward staple: goes down in j, back in i, up in j
    Uj_shift_i_back_j = jnp.roll(jnp.roll(Uj, shift=-1, axis=i), shift=1, axis=j)
    Ui_back_j = jnp.roll(Ui, shift=1, axis=j)
    Uj_back_j = jnp.roll(Uj, shift=1, axis=j)
    backward = jnp.conj(Uj_shift_i_back_j) * jnp.conj(Ui_back_j) * Uj_back_j

    return forward + backward


@partial(jit, static_argnums=(1,))
def staple(links: Array, i: int) -> Array:
    """
    Compute staple S_i(n) for link i at all sites.

    The staple is the sum of 4 plaquettes touching link U_i(n),
    needed for computing the force ∂H/∂A_i.

    Args:
        links: Link variables (3, N, N, N)
        i: Direction of link - static for JIT

    Returns:
        Staple values at all sites (N, N, N) complex
    """
    if i == 0:
        return _staple_contribution(links, 0, 1) + _staple_contribution(links, 0, 2)
    elif i == 1:
        return _staple_contribution(links, 1, 0) + _staple_contribution(links, 1, 2)
    else:  # i == 2
        return _staple_contribution(links, 2, 0) + _staple_contribution(links, 2, 1)


@jit
def all_staples(links: Array) -> Array:
    """
    Compute staples for all three link directions.

    Args:
        links: Link variables (3, N, N, N)

    Returns:
        Staples for all directions (3, N, N, N) complex
    """
    S0 = staple(links, 0)
    S1 = staple(links, 1)
    S2 = staple(links, 2)

    return jnp.stack([S0, S1, S2], axis=0)


@jit
def gauge_force(links: Array, dx: float, g: float) -> Array:
    """
    Compute force F_i = -∂H_mag/∂A_i for magnetic Hamiltonian.

    F_i = (g dx) * Im(U_i * S_i^*)

    Args:
        links: Link variables (3, N, N, N)
        dx: Lattice spacing
        g: Gauge coupling

    Returns:
        Force components (3, N, N, N) real
    """
    staples = all_staples(links)

    # Force from magnetic energy: Im(U_i S_i^*)
    force = g * dx * jnp.imag(links * jnp.conj(staples))

    return force.astype(jnp.float32)


@jit
def magnetic_energy_density(links: Array, dx: float, g: float) -> Array:
    """
    Compute total magnetic energy E_B = (1/2) ∫ B² dV.

    On lattice: E_B = (1/2) Σ_n Σ_k B_k(n)² * dx³

    Args:
        links: Link variables
        dx: Lattice spacing
        g: Gauge coupling

    Returns:
        Total magnetic energy (as JAX scalar)
    """
    B = magnetic_field(links, dx, g)
    dV = dx**3
    E_mag = 0.5 * jnp.sum(B**2) * dV
    return E_mag


@jit
def electric_energy_density(E: Array, dx: float) -> Array:
    """
    Compute total electric energy E_E = (1/2) ∫ E² dV.

    On lattice: E_E = (1/2) Σ_n Σ_i E_i(n)² * dx³

    Args:
        E: Electric field components (3, N, N, N)
        dx: Lattice spacing

    Returns:
        Total electric energy (as JAX scalar)
    """
    dV = dx**3
    E_elec = 0.5 * jnp.sum(E**2) * dV
    return E_elec


@jit
def divergence_3d(field: Array, dx: float) -> Array:
    """
    Compute divergence of a 3-component vector field.

    ∇·F = ∂F_x/∂x + ∂F_y/∂y + ∂F_z/∂z

    Uses centered difference: (F_i(n) - F_i(n-i)) / dx

    Args:
        field: Vector field components (3, N, N, N)
        dx: Lattice spacing

    Returns:
        Divergence at each site (N, N, N)
    """
    div = (
        (field[0] - jnp.roll(field[0], 1, axis=0))
        + (field[1] - jnp.roll(field[1], 1, axis=1))
        + (field[2] - jnp.roll(field[2], 1, axis=2))
    ) / dx
    return div


@jit
def gauss_constraint(links: Array, E: Array, dx: float) -> Array:
    """
    Compute Gauss constraint violation G(n) = ∇·E for pure gauge (no sources).

    Args:
        links: Link variables (not used, kept for interface consistency)
        E: Electric field (3, N, N, N)
        dx: Lattice spacing

    Returns:
        Gauss violation at each site (N, N, N)
    """
    return divergence_3d(E, dx)


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

    Returns sqrt(mean(|∇·E|²))

    Args:
        links: Link variables (not used, kept for interface consistency)
        E: Electric field (3, N, N, N)
        dx: Lattice spacing

    Returns:
        RMS Gauss violation (scalar)
    """
    G = gauss_constraint(links, E, dx)
    return jnp.sqrt(jnp.mean(G**2))


def gauge_energy(links: Array, E: Array, dx: float, g: float) -> dict:
    """
    Compute all energy components for pure gauge field.

    Args:
        links: Gauge links (3, N, N, N)
        E: Electric field (3, N, N, N)
        dx: Lattice spacing
        g: Gauge coupling

    Returns:
        Dictionary with 'electric', 'magnetic', and 'total' energies
    """
    E_elec = float(electric_energy_density(E, dx))
    E_mag = float(magnetic_energy_density(links, dx, g))

    return {
        "electric": E_elec,
        "magnetic": E_mag,
        "total": E_elec + E_mag,
    }


@jit
def evolve_links_half_step(links: Array, E: Array, dt: float) -> Array:
    """
    Evolve gauge links by half time step: U → U exp(i E dt/2).

    This is the position update in the leapfrog scheme.
    Links are normalized to maintain unitarity.

    Args:
        links: Gauge links (3, N, N, N) complex
        E: Electric field (3, N, N, N) real
        dt: Time step

    Returns:
        Updated links (3, N, N, N), normalized to unit modulus
    """
    phase_kick = jnp.exp(1j * E * dt / 2)
    links_new = links * phase_kick
    return links_new / jnp.abs(links_new)


@jit
def evolve_links_full_step(links: Array, E: Array, dt: float) -> Array:
    """
    Evolve gauge links by full time step: U → U exp(i E dt).

    Args:
        links: Gauge links (3, N, N, N) complex
        E: Electric field (3, N, N, N) real
        dt: Time step

    Returns:
        Updated links (3, N, N, N), normalized to unit modulus
    """
    phase_kick = jnp.exp(1j * E * dt)
    links_new = links * phase_kick
    return links_new / jnp.abs(links_new)
