import jax.numpy as jnp
from jax import random
import jaxlatt as jl


def higgs_2D(potential, filename, size=(256, 256), length=(100, 100), steps=4000):
    lattice = jl.create_initial_lattice_2d(
        size=size,
        length=length,
        field_type="random",
        amplitude=0.3,  # Small random fluctuations near unstable origin
        velocity=0.0,
        key=random.PRNGKey(42),
    )

    times, snapshots = jl.evolve(
        lattice=lattice,
        potential=potential,
        dt=0.01,
        steps=steps,
        snapshot_interval=10,
        verbose=True,
    )

    print("\n   Creating animation...")
    anim2d = jl.field_2d(
        times=times,
        snapshots=snapshots,
        potential=potential,
        interval=50,
        vmin=-3.0,
        vmax=3.0,
        save_path=f"../animations/{filename}" if filename else None,
        # show_energy=True,
        # title="2D Higgs: Phase Transition & Domain Formation",
    )

    print(f"   ✓ Saved {filename}\n")
    return anim2d
