"""
JAX-based Lattice Field Theory Simulation Package
==================================================

A modern implementation of lattice field theory simulations using JAX,
designed for early-universe cosmology and gauge theory dynamics.

**Module Structure**:
- `jaxlatt.core` - Lattice classes, cosmology, field initialization, potentials
- `jaxlatt.operators` - Field operators (scalar, gauge, coupled, spectral, autodiff)
- `jaxlatt.evolution` - Time integrators (leapfrog, RK4, expanding space)
- `jaxlatt.animation` - Animation scripts & tools for visualizing field evolution
- `jaxlatt.plotting` - Plotting utilities for phase portraits and comparisons
- `jaxlatt.observables` - Physics observables, spectra, energy tracking

**Usage Example**:
```python
from jaxlatt.core import create_random_coupled_lattice, FRWUniverse
from jaxlatt.operators import coupled_energy
from jaxlatt.evolution import coupled_leapfrog_step
from jaxlatt.observables import power_spectrum_3d

# Create lattice
lattice = create_random_coupled_lattice(key, size=(32,32,32), length=10.0)

# Evolve
for _ in range(100):
    lattice = coupled_leapfrog_step(lattice, dt=0.01)
```
"""

# =============================================================================
# Main Public API
# =============================================================================

# Core: Lattice classes and cosmology
from jaxlatt.core import (
    # Lattice data structures
    Lattice,
    CoupledLattice,
    GaugeLattice,
    FRWUniverse,
    # Field initialization
    create_vacuum_coupled_lattice,
    create_random_coupled_lattice,
    create_higgs_vev_lattice,
    create_vacuum_gauge_lattice,
    create_random_gauge_lattice,
    # Cosmology
    create_radiation_universe,
    friedmann_step_predictor_corrector,
    # Potentials
    potentials,
)

# Evolution: Time integrators
from jaxlatt.evolution import (
    # Simple scalar field evolution
    evolve,
    evolve_step,
    # Flat spacetime
    coupled_leapfrog_step,
    coupled_evolve,
    gauge_leapfrog_step,
    gauge_evolve,
    # Expanding universe (FLRW)
    coupled_leapfrog_step_expanding,
    coupled_evolve_expanding,
    make_radiation_rho,
)

# Utilities for creating simple lattices
from jaxlatt.utils import (
    create_initial_lattice_1d,
    create_initial_lattice_2d,
    create_initial_lattice_3d,
)

# Animation utilities
from jaxlatt.animation import (
    field_1d,
    field_2d,
    field_3d,
    slices_3d,
    phase_space_1d,
    create_comparison_animation,
    animate_multi_snapshot_grid,
)

# Plotting utilities
from jaxlatt.plotting import phase_portrait, dimensional_comparison

# Operators: Most commonly used
from jaxlatt.operators import (
    # Energies
    coupled_energy,
    gauge_energy,
    # Forces and operators
    scalar_force,
    covariant_laplacian,
    # Gauge theory
    plaquette,
    magnetic_field,
    gauss_constraint,
)

# Observables: Physics analysis
from jaxlatt.observables import (
    # Spectra
    power_spectrum_3d,
    power_spectrum_with_binning,
    # Energy tracking
    EnergyTracker,
    total_energy,
    energy_components,
    # Constraints
    project_gauss_constraint,
    gauge_transform_coupled_lattice,
)

# I/O: Checkpointing
from jaxlatt.io import (
    save_checkpoint,
    load_checkpoint,
)

# =============================================================================
# Package Metadata
# =============================================================================

try:
    from importlib.metadata import version, PackageNotFoundError

    try:
        __version__ = version("jaxlatt")
    except PackageNotFoundError:
        # Package is not installed (e.g., running from source in development)
        __version__ = "0.0.1-dev"
except ImportError:
    # Fallback for Python < 3.8 (shouldn't happen with requires-python >= 3.12)
    __version__ = "0.0.1-dev"

__author__ = "Rodrigo Calderon"
__all__ = [
    # === Metadata ===
    "__version__",
    "__author__",
    # === Core ===
    "Lattice",
    "CoupledLattice",
    "GaugeLattice",
    "FRWUniverse",
    # Field initialization
    "create_vacuum_coupled_lattice",
    "create_random_coupled_lattice",
    "create_higgs_vev_lattice",
    "create_vacuum_gauge_lattice",
    "create_random_gauge_lattice",
    # Cosmology
    "create_radiation_universe",
    "friedmann_step_predictor_corrector",
    # Potentials
    "potentials",
    # === Evolution ===
    "evolve",
    "evolve_step",
    "coupled_leapfrog_step",
    "coupled_evolve",
    "gauge_leapfrog_step",
    "gauge_evolve",
    "coupled_leapfrog_step_expanding",
    "coupled_evolve_expanding",
    "make_radiation_rho",
    # === Utilities ===
    "create_initial_lattice_1d",
    "create_initial_lattice_2d",
    "create_initial_lattice_3d",
    # === Animation ===
    "field_1d",
    "field_2d",
    "field_3d",
    "slices_3d",
    "phase_space_1d",
    "create_comparison_animation",
    "animate_multi_snapshot_grid",
    # === Plotting ===
    "phase_portrait",
    "dimensional_comparison",
    # === Operators ===
    "coupled_energy",
    "gauge_energy",
    "scalar_force",
    "covariant_laplacian",
    "plaquette",
    "magnetic_field",
    "gauss_constraint",
    # === Observables ===
    "power_spectrum_3d",
    "power_spectrum_with_binning",
    "EnergyTracker",
    "total_energy",
    "energy_components",
    "project_gauss_constraint",
    "gauge_transform_coupled_lattice",
    # === I/O ===
    "save_checkpoint",
    "load_checkpoint",
]


def main() -> None:
    """Entry point for CLI."""
    from jaxlatt.cli import app

    app()
