From c7c68cf1fc978d75220bce1e744c41882a24a637 Mon Sep 17 00:00:00 2001
From: Parthib Roy <proy@lbl.gov>
Date: Fri, 27 Sep 2024 17:07:55 -0700
Subject: [PATCH] Added export functionality for user inputs in dashboard

---
 .../dashboard/Toolbar/exportTemplate.py       | 91 +++++++++++++++++++
 .../impactx/dashboard/Toolbar/toolbarMain.py  | 22 +++++
 2 files changed, 113 insertions(+)
 create mode 100644 src/python/impactx/dashboard/Toolbar/exportTemplate.py

diff --git a/src/python/impactx/dashboard/Toolbar/exportTemplate.py b/src/python/impactx/dashboard/Toolbar/exportTemplate.py
new file mode 100644
index 000000000..005663b8c
--- /dev/null
+++ b/src/python/impactx/dashboard/Toolbar/exportTemplate.py
@@ -0,0 +1,91 @@
+from ..Input.distributionParameters.distributionMain import parameter_input_checker
+from ..Input.latticeConfiguration.latticeMain import parameter_input_checker_for_lattice
+from ..trame_setup import setup_server
+
+server, state, ctrl = setup_server()
+
+# -----------------------------------------------------------------------------
+# Helper Functions
+# -----------------------------------------------------------------------------
+
+
+def build_distribution_list():
+    """
+    Generates an instance of distribution inputs
+    as a string for exporting purposes.
+    """
+    distribution_name = state.selectedDistribution
+    parameters = parameter_input_checker()
+
+    distribution_parameters = ",\n    ".join(
+        f"{key}={value}" for key, value in parameters.items()
+    )
+
+    return (
+        f"distr = distribution.{distribution_name}(\n    {distribution_parameters},\n)"
+    )
+
+
+def build_lattice_list():
+    """
+    Generates a string representation of lattice element
+    inputs for export purposes.
+    """
+
+    lattice_elements = ",\n    ".join(
+        f'elements.{element["name"]}('
+        + ", ".join(
+            f"{key}={value}"
+            for key, value in parameter_input_checker_for_lattice(element).items()
+        )
+        + ")"
+        for element in state.selectedLatticeList
+    )
+
+    return f"lattice_configuration = [\n    {lattice_elements}\n]"
+
+
+# -----------------------------------------------------------------------------
+# Trame setup
+# -----------------------------------------------------------------------------
+
+
+def input_file():
+    """
+    This function creates the template to export
+    dashboard user inputs into a python script.
+    """
+    script = f"""
+from impactx import ImpactX, distribution, elements
+
+sim = ImpactX()
+
+sim.particle_shape = {state.particle_shape}
+sim.space_charge = False
+sim.csr = False
+sim.slice_step_diagnostics = True
+
+sim.init_grids()
+
+# Initialize particle beam
+kin_energy_MeV = {state.kin_energy_MeV}
+bunch_charge_C = {state.bunch_charge_C}
+npart = {state.npart}
+
+# Reference particle
+ref = sim.particle_container().ref_particle()
+ref.set_charge_qe(-1.0).set_mass_MeV(0.510998950).set_kin_energy_MeV(kin_energy_MeV)
+
+{build_distribution_list()}
+sim.add_particles(bunch_charge_C, distr, npart)
+
+{build_lattice_list()}
+sim.lattice.extend(lattice_configuration)
+
+# Simulate
+sim.evolve()
+
+sim.finalize()
+"""
+
+    return script
diff --git a/src/python/impactx/dashboard/Toolbar/toolbarMain.py b/src/python/impactx/dashboard/Toolbar/toolbarMain.py
index a842e6e8b..4b71b0aca 100644
--- a/src/python/impactx/dashboard/Toolbar/toolbarMain.py
+++ b/src/python/impactx/dashboard/Toolbar/toolbarMain.py
@@ -9,11 +9,22 @@
 from trame.widgets import vuetify
 
 from ..trame_setup import setup_server
+from .exportTemplate import input_file
 
 server, state, ctrl = setup_server()
 
 state.show_dashboard_alert = True
 
+# -----------------------------------------------------------------------------
+# Triggers
+# -----------------------------------------------------------------------------
+
+
+@ctrl.trigger("export")
+def on_export_click():
+    return input_file()
+
+
 # -----------------------------------------------------------------------------
 # Common toolbar elements
 # -----------------------------------------------------------------------------
@@ -25,6 +36,15 @@ class ToolbarElements:
     Vuetify UI elements for toolbar.
     """
 
+    @staticmethod
+    def export_input_data():
+        vuetify.VIcon(
+            "mdi-download",
+            style="color: #00313C;",
+            click="utils.download('impactx_simulation.py', trigger('export'), 'text/plain')",
+            disabled=("disableRunSimulationButton", True),
+        )
+
     @staticmethod
     def plot_options():
         vuetify.VSelect(
@@ -75,6 +95,8 @@ def input_toolbar():
         """
 
         (ToolbarElements.dashboard_info(),)
+        vuetify.VSpacer()
+        ToolbarElements.export_input_data()
 
     @staticmethod
     def run_toolbar():