import marimo

__generated_with = "0.18.0"
app = marimo.App(width="full")

with app.setup(hide_code=True):
    # Initialization code that runs before all other cells
    import jax
    import marimo as mo
    import numpy as np
    import matplotlib.pyplot as plt
    import scienceplots
    plt.style.use(['science'])

    import jaxlatt as jl
    import jaxlatt.animation as animate



@app.cell
def _():
    mo.md("""
    # JAXLatt - Animate
    """)
    return


@app.cell(hide_code=True)
def _():
    print(f"JAX backend: {jax.default_backend()}\n")
    return


@app.cell
def _():
    # Parameters
    mu2 = 1.0
    lam = 0.5
    v_theory = np.sqrt(mu2 / lam)
    potential = jl.ScalarPotential.double_well(mu2=mu2, lam=lam)
    return (potential,)


@app.cell
def _():
    size_options=[(32,32),(64,64),(128,128),(256,256)]
    lattice_ui=mo.ui.dictionary({'size': mo.ui.dropdown(size_options,(64,64),label="Lattice Size"),
                                 'length': mo.ui.number(start=0,stop=1000,value=10,step=1,label="Length"),
                                 'field_type': mo.ui.dropdown(['gaussian','random','zero','sine'], 'random', label="Field Type"),
                                 'amplitude': mo.ui.number(start=0,stop=10,value=0.3,step=0.1,label="Amplitude"),
                                 'velocity': mo.ui.number(start=0,stop=3,value=0.,step=0.01,label="Velocity"),
                                })
    return (lattice_ui,)


@app.cell
def _(lattice_ui):
    # mo.md
    lattice_ui.hstack()
    return


@app.cell
def _():
    # lattice_kwargs={**lattice_ui.value}
    return


@app.cell
def _(lattice_ui):
    lattice = jl.create_initial_lattice_2d(
        key=jax.random.PRNGKey(42),
        **lattice_ui.value
    )

    # size=(64, 64),
    # length=(10.0, 10.0),
    # field_type="random",
    # amplitude=0.3,  # Small random fluctuations near unstable origin
    # velocity=0.0,
    return (lattice,)


@app.cell
def _(lattice):
    lattice
    return


@app.cell
def _():
    evolve=mo.ui.run_button(label="Evolve Lattice")
    evolve
    return (evolve,)


@app.cell
def _(evolve, lattice, potential):
    mo.stop(not evolve.value, mo.md("Click the button to run the animation"))
    with mo.status.spinner("Evolving lattice...") as _spinner:
        times, snapshots = jl.evolve(
            lattice=lattice,
            potential=potential,
            dt=0.01,
            steps=3000,
            snapshot_interval=10,
            verbose=True,
        )
        _spinner.update(subtitle='Simulation complete!')
    return snapshots, times


@app.cell
def _(snapshots, times):
    import jaxlatt.plotting as plot
    import functools

    @functools.cache
    def snapshots_2d(index,**kwargs):
        _=plot.snapshot2d(snapshots[index],**kwargs)
        ax=plt.gca()
        ax.text(0.05,0.95,f" $t = {times[index]:.2f}$",transform=ax.transAxes,color='k',ha='left',va='top',fontsize=16)
        return ax
    return (snapshots_2d,)


@app.cell
def _(snapshots):
    selected_index=mo.ui.slider(0,len(snapshots)-1,1,0,label="Snapshot Index")
    selected_index
    return (selected_index,)


@app.cell
def _(selected_index, snapshots_2d):
    snapshots_2d(selected_index.value,vmin=-2,vmax=2)
    return


@app.cell
def _():
    run=mo.ui.run_button(label="Animate & Export to GIF")
    run
    return (run,)


@app.cell
def _(potential, run, snapshots, times):
    mo.stop(not run.value,mo.md("Click the button to run the animation"))
    with mo.status.spinner("Exporting animation to .gif ...") as _spinner:
        _anim_2d=animate.field_2d(
            times=times,
            snapshots=snapshots,
            potential=potential,
            interval=10,
            save_path="dark_phase_transition_field_higgs2d.gif",
            vmin=-2,
            vmax=2,
            dark=True
        )
    return


@app.cell
def _():
    # anim_2d.to_jshtml()
    return


@app.cell
def _():
    # mo.as_html(anim_2d.to_html5_video())
    return


@app.cell
def _():
    return


if __name__ == "__main__":
    app.run()
