"""
Constraint enforcement for gauge theories.

This module provides comprehensive tools for:
- Computing Gauss law constraint violations
- Projecting fields onto constraint manifold
- Gauge transformations
- Baumgarte-Shapiro constraint stabilization
- Diagnostic utilities

All Jacobians computed automatically via JAX autodiff.

Note on Gauss Constraint Implementations
----------------------------------------
The canonical implementations are:

- Pure gauge (no sources): `jaxlatt.operators.gauge.gauss_constraint`
  Signature: (links, E, dx) -> G = div(E)

- With scalar sources: `jaxlatt.operators.coupled.gauss_constraint_with_source`
  Signature: (links, E, phi, pi, g, dx) -> G = div(E) + g*Im(phi* pi)

This module re-exports these for convenience and provides additional
constraint enforcement tools (projection, gauge transformations, etc.).
"""

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

from jaxlatt.core import CoupledLattice
from jaxlatt.operators.coupled import scalar_charge_density

# Import canonical Gauss constraint implementations
from jaxlatt.operators.gauge import divergence_3d
from jaxlatt.operators.coupled import (
    gauss_constraint_with_source as _gauss_constraint_with_source,
)


# =============================================================================
# Gauss Law Constraint - Canonical Re-exports
# =============================================================================


@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.

    This is the canonical implementation for coupled scalar-gauge systems.
    Gauss law: G = div(E) + g*rho = 0, where rho = Im(phi* pi).

    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)

    See Also:
        jaxlatt.operators.coupled.gauss_constraint_with_source
    """
    return _gauss_constraint_with_source(links, E, phi, pi, g, dx)


# =============================================================================
# Constraint Projection Methods
# =============================================================================


def project_gauss_constraint(
    E: jnp.ndarray,
    phi: jnp.ndarray,
    pi: jnp.ndarray,
    g: float,
    dx: float,
    tolerance: float = 1e-10,
    max_iterations: int = 100,
) -> jnp.ndarray:
    """
    Project electric field to satisfy Gauss constraint div(E) + g*rho = 0.

    Uses iterative Jacobi method to solve Poisson equation:
        laplacian(phi_gauge) = div(E) + g*rho
        E_corrected = E - grad(phi_gauge)

    This ensures Gauss law is satisfied to numerical precision without
    changing the transverse part of E.

    Parameters:
    -----------
    E : float array (3, N, N, N)
        Electric field to project
    phi : complex array (N, N, N)
        Scalar field
    pi : complex array (N, N, N)
        Conjugate momentum
    g : float
        Gauge coupling
    dx : float
        Lattice spacing
    tolerance : float
        Convergence tolerance for iterative solver
    max_iterations : int
        Maximum number of iterations

    Returns:
    --------
    E_projected : float array (3, N, N, N)
        Electric field satisfying Gauss constraint

    Note:
        Uses the canonical Gauss constraint form from operators/coupled.py:
        G = div(E) + g*Im(phi* pi) = 0
    """
    # Compute charge density using canonical function
    rho = scalar_charge_density(phi, pi, g)

    # Compute divergence of E using canonical function
    div_E = divergence_3d(E, dx)

    # Source term: div(E) + g*rho (this is what we want to set to zero)
    source = div_E + rho

    # Solve nabla^2 phi_gauge = source using Jacobi iteration
    phi_gauge = jnp.zeros_like(phi, dtype=jnp.float32)

    # Jacobi iteration: solve nabla^2 phi = source
    # In discrete form: (phi_{i+1} + phi_{i-1} + ... - 6*phi_i) / dx^2 = source
    # Rearrange: phi_i = (phi_{i+1} + phi_{i-1} + ... - source*dx^2) / 6
    for _ in range(max_iterations):
        # Neighbor sum
        neighbor_sum = jnp.zeros_like(phi_gauge)
        for i in range(3):
            neighbor_sum = (
                neighbor_sum
                + jnp.roll(phi_gauge, -1, axis=i)
                + jnp.roll(phi_gauge, 1, axis=i)
            )

        # Jacobi update
        phi_gauge_new = (neighbor_sum - source * dx**2) / 6.0

        # Check convergence
        diff = phi_gauge_new - phi_gauge
        error = jnp.sqrt(jnp.mean(diff**2))

        phi_gauge = phi_gauge_new

        if error < tolerance:
            break

    # Compute gradient of phi_gauge
    E_correction = jnp.zeros_like(E)
    for i in range(3):
        phi_forward = jnp.roll(phi_gauge, -1, axis=i)
        grad_i = (phi_forward - phi_gauge) / dx
        E_correction = E_correction.at[i].set(grad_i)

    # Project E
    E_projected = E - E_correction

    return E_projected


def project_gauss_coupled_lattice(
    lattice: CoupledLattice,
    tolerance: float = 1e-10,
    max_iterations: int = 100,
) -> CoupledLattice:
    """
    Project electric field in coupled lattice to satisfy Gauss constraint.

    Creates a new lattice with E field corrected to enforce div(E) + g*rho = 0.

    Parameters:
    -----------
    lattice : CoupledLattice
        Lattice to project
    tolerance : float
        Convergence tolerance
    max_iterations : int
        Maximum iterations for solver

    Returns:
    --------
    CoupledLattice with projected E field
    """
    E_projected = project_gauss_constraint(
        lattice.E,
        lattice.phi,
        lattice.pi,
        lattice.g,
        lattice.dx,
        tolerance,
        max_iterations,
    )

    return CoupledLattice(
        phi=lattice.phi,
        pi=lattice.pi,
        links=lattice.links,
        E=E_projected,
        m=lattice.m,
        lambda_=lattice.lambda_,
        g=lattice.g,
        dx=lattice.dx,
        size=lattice.size,
        length=lattice.length,
    )


# =============================================================================
# Constraint Stabilization (Baumgarte-Shapiro)
# =============================================================================


@jit
def stabilized_constraint_evolution(
    G: Array,
    dG_dt: Array,
    alpha: float = 0.1,
    beta: float = 0.01,
) -> Array:
    """
    Baumgarte-Shapiro constraint damping.

    Instead of enforcing dG/dt = 0, enforce:
    d^2 G/dt^2 + alpha * dG/dt + beta * G = 0

    This exponentially damps constraint violations.

    Args:
        G: Current constraint violation
        dG_dt: Time derivative of constraint
        alpha: Linear damping coefficient
        beta: Constraint restoring coefficient

    Returns:
        Modified d^2 G/dt^2 with damping
    """
    return -alpha * dG_dt - beta * G


# =============================================================================
# Gauge Transformations
# =============================================================================


@jit
def gauge_transform_scalar(
    phi: jnp.ndarray, alpha: jnp.ndarray, g: float
) -> jnp.ndarray:
    """
    Apply gauge transformation to scalar field.

    Under U(1) gauge transformation with phase alpha(x):
        phi(x) -> e^{ig alpha(x)} phi(x)

    Parameters:
    -----------
    phi : complex array (N, N, N)
        Scalar field
    alpha : float array (N, N, N)
        Gauge transformation parameter
    g : float
        Gauge coupling

    Returns:
    --------
    phi' : complex array (N, N, N)
        Transformed scalar field
    """
    return jnp.exp(1j * g * alpha) * phi


@jit
def gauge_transform_links(
    links: jnp.ndarray, alpha: jnp.ndarray, g: float
) -> jnp.ndarray:
    """
    Apply gauge transformation to gauge links.

    Under U(1) gauge transformation:
        U_i(n) -> e^{ig alpha(n)} U_i(n) e^{-ig alpha(n+i)}

    This keeps the gauge-covariant derivative properly transformed.

    Parameters:
    -----------
    links : complex array (3, N, N, N)
        Gauge links
    alpha : float array (N, N, N)
        Gauge transformation parameter
    g : float
        Gauge coupling

    Returns:
    --------
    U' : complex array (3, N, N, N)
        Transformed gauge links
    """
    links_new = jnp.zeros_like(links)

    for i in range(3):
        alpha_n = alpha
        alpha_n_plus_i = jnp.roll(alpha, -1, axis=i)
        phase = jnp.exp(1j * g * (alpha_n - alpha_n_plus_i))
        links_new = links_new.at[i].set(phase * links[i])

    return links_new


def gauge_transform_coupled_lattice(
    lattice: CoupledLattice, alpha: jnp.ndarray
) -> CoupledLattice:
    """
    Apply gauge transformation to entire coupled lattice.

    Transforms:
        phi -> e^{ig alpha} phi
        pi -> e^{ig alpha} pi
        U_i -> e^{ig alpha(n)} U_i e^{-ig alpha(n+i)}
        E unchanged (gauge invariant)

    Parameters:
    -----------
    lattice : CoupledLattice
        Lattice to transform
    alpha : float array (N, N, N)
        Gauge transformation parameter

    Returns:
    --------
    CoupledLattice after gauge transformation
    """
    phi_new = gauge_transform_scalar(lattice.phi, alpha, lattice.g)
    pi_new = gauge_transform_scalar(lattice.pi, alpha, lattice.g)
    links_new = gauge_transform_links(lattice.links, alpha, lattice.g)

    return CoupledLattice(
        phi=phi_new,
        pi=pi_new,
        links=links_new,
        E=lattice.E,  # E is gauge invariant
        m=lattice.m,
        lambda_=lattice.lambda_,
        g=lattice.g,
        dx=lattice.dx,
        size=lattice.size,
        length=lattice.length,
    )


# =============================================================================
# Diagnostics
# =============================================================================


@jit
def check_gauge_invariance(
    observable_before: float, observable_after: float, tolerance: float = 1e-10
) -> jnp.ndarray:
    """
    Check if an observable is gauge invariant.

    Parameters:
    -----------
    observable_before : float
        Observable value before gauge transformation
    observable_after : float
        Observable value after gauge transformation
    tolerance : float
        Tolerance for equality

    Returns:
    --------
    is_invariant : bool (as JAX array)
        True if observable is gauge invariant
    """
    return jnp.abs(observable_before - observable_after) < tolerance


def compute_constraint_violation_norm(
    phi: Array,
    pi: Array,
    E: Array,
    links: Array,
    g: float,
    dx: float,
) -> Dict[str, float]:
    """
    Compute various norms of constraint violation.

    Args:
        phi: Scalar field
        pi: Conjugate momentum
        E: Electric field
        links: Gauge links
        g: Gauge coupling
        dx: Lattice spacing

    Returns:
        Dict with L1, L2, Linf norms of G
    """
    G = gauss_constraint_with_source(links, E, phi, pi, g, dx)

    return {
        "L1_norm": float(jnp.mean(jnp.abs(G))),
        "L2_norm": float(jnp.sqrt(jnp.mean(G**2))),
        "Linf_norm": float(jnp.max(jnp.abs(G))),
        "total_violation": float(jnp.sum(jnp.abs(G)) * dx**3),
    }


def check_constraint_preservation(
    phi_old: Array,
    pi_old: Array,
    E_old: Array,
    links_old: Array,
    phi_new: Array,
    pi_new: Array,
    E_new: Array,
    links_new: Array,
    g: float,
    dx: float,
) -> Dict[str, float]:
    """
    Check if time evolution preserves Gauss constraint.

    Returns violation before and after, plus drift rate.
    """
    G_old = gauss_constraint_with_source(links_old, E_old, phi_old, pi_old, g, dx)
    G_new = gauss_constraint_with_source(links_new, E_new, phi_new, pi_new, g, dx)

    violation_old = float(jnp.sqrt(jnp.mean(G_old**2)))
    violation_new = float(jnp.sqrt(jnp.mean(G_new**2)))
    drift = violation_new - violation_old

    return {
        "violation_before": violation_old,
        "violation_after": violation_new,
        "drift": drift,
        "drift_rate": drift / (violation_old + 1e-10),
    }


def charge_current_continuity(
    lattice: CoupledLattice,
) -> Tuple[jnp.ndarray, jnp.ndarray]:
    """
    Compute charge and current densities for continuity equation check.

    The continuity equation for conserved charge:
        d rho/dt + div(j) = 0

    where:
        rho = g Im(phi* pi)  (charge density)
        j_i = (g/2dx) Im[phi*(n) U_i(n) phi(n+i)]  (current density)

    Parameters:
    -----------
    lattice : CoupledLattice
        Lattice state

    Returns:
    --------
    rho : float array (N, N, N)
        Charge density
    div_j : float array (N, N, N)
        Divergence of current
    """
    from jaxlatt.operators.coupled import scalar_current_density

    # Charge density
    rho = scalar_charge_density(lattice.phi, lattice.pi, lattice.g)

    # Current divergence
    div_j = jnp.zeros_like(lattice.phi, dtype=jnp.float32)

    for i in range(3):
        j_i = scalar_current_density(
            lattice.phi, lattice.links, i, lattice.g, lattice.dx
        )
        j_i_backward = jnp.roll(j_i, 1, axis=i)
        div_j = div_j + (j_i - j_i_backward) / lattice.dx

    return rho, div_j
