code: # Standard Library Imports
  import ast
  import json
  import os
  import re
  
  # Third-Party Imports
  import textwrap
  from typing import Any, Dict, List, Optional, Type
  
  import langchain
  import streamlit as st
  from langchain.chains import LLMChain
  from langchain.prompts import PromptTemplate
  from langchain.tools import BaseTool
  from openmm import (
      AndersenThermostat,
      BrownianIntegrator,
      LangevinIntegrator,
      LangevinMiddleIntegrator,
      MonteCarloBarostat,
      OpenMMException,
      Platform,
      VerletIntegrator,
      app,
      unit,
  )
  from openmm.app import (
      PME,
      AllBonds,
      CutoffNonPeriodic,
      CutoffPeriodic,
      DCDReporter,
      Ewald,
      ForceField,
      HAngles,
      HBonds,
      Modeller,
      NoCutoff,
      PDBFile,
      PDBReporter,
      PDBxFile,
      Simulation,
      StateDataReporter,
  )
  from openmm.unit import bar, femtoseconds, kelvin, nanometers, picosecond, picoseconds
  from pydantic import BaseModel, Field
  
  from mdagent.tools.base_tools.preprocess_tools import CleaningTools
  
  # Local Library/Application Imports
  from mdagent.utils import FileType, PathRegistry
  
  # TODO delete files created from the simulation if not needed.
  
  FORCEFIELD_LIST = [
      "amber14/DNA.OL15.xml",
      "amber14/DNA.bsc1.xml",
      "amber14/RNA.OL3.xml",
      "amber14/lipid17.xml",
      "amber14/protein.ff14SB.xml",
      "amber14/protein.ff15ipq.xml",
      "amber14/spce.xml",
      "amber14/tip3p.xml",
      "amber14/tip3pfb.xml",
      "amber14/tip4pew.xml",
      "amber14/tip4pfb.xml",
      "charmm36/spce.xml",
      "charmm36/tip3p-pme-b.xml",
      "charmm36/tip3p-pme-f.xml",
      "charmm36/tip4p2005.xml",
      "charmm36/tip4pew.xml",
      "charmm36/tip5p.xml",
      "charmm36/tip5pew.xml",
      "charmm36/water.xml",
      "absinth.xml",
      "amber03.xml",
      "amber03_obc.xml",
      "amber10.xml",
      "amber10_obc.xml",
      "amber14-all",
      "amber14-all.xml",
      "amber96.xml",
      "amber96_obc.xml",
      "amber99Test.xml",
      "amber99_obc.xml",
      "amber99sb.xml",
      "amber99sbildn.xml",
      "amber99sbnmr.xml",
      "amberfb15.xml",
      "amoeba2009.xml",
      "amoeba2009_gk.xml",
      "amoeba2013.xml",
      "amoeba2013_gk.xml",
      "charmm36.xml",
      "charmm_polar_2013.xml",
      "hydrogens.xml",
      "iamoeba.xml",
      "pdbNames.xml",
      "residues.xml",
      "spce.xml",
      "swm4ndp.xml",
      "tip3p.xml",
      "tip3pfb.xml",
      "tip4pew.xml",
      "tip4pfb.xml",
      "tip5p.xml",
  ]
  
  
  class SimulationFunctions:
      def __init__(
          self,
          path_registry,
          temperature: float = 0.05,
          model_name: str = "gpt-4",
          request_timeout: int = 1000,
          max_tokens: int = 2000,
      ):
          self.path_registry = path_registry
          self.temperature = temperature
          self.model_name = model_name
          self.request_timeout = request_timeout
          self.max_tokens = max_tokens
  
          self.llm = langchain.chat_models.ChatOpenAI(
              temperature=self.temperature,
              model_name=self.model_name,
              request_timeout=self.request_timeout,
              max_tokens=self.request_timeout,
          )
  
      #######==================System Congifuration==================########
      # System Configuration initialization.
  
      def _create_system(
          pdb,
          forcefield,
          nonbondedMethod="NoCutoff",
          nonbondedCutoff=None,
          ewaldErrorTolerance=None,
          constraints="None",
          rigidWater=False,
          constraintTolerance=None,
          **kwargs,
      ):
          # Create a dictionary to hold system parameters
          system_params = {
              "nonbondedMethod": nonbondedMethod,
              "constraints": constraints,
              "rigidWater": rigidWater,
          }
  
          # Set nonbondedCutoff if applicable
          if (
              nonbondedMethod in ["PME", "CutoffNonPeriodic", "CutoffPeriodic"]
              and nonbondedCutoff is not None
          ):
              system_params["nonbondedCutoff"] = nonbondedCutoff
  
          # Set ewaldErrorTolerance if PME is used
          if nonbondedMethod == "PME" and ewaldErrorTolerance is not None:
              system_params["ewaldErrorTolerance"] = ewaldErrorTolerance
  
          # Set constraintTolerance if constraints are used
          if constraints in ["HBonds", " AllBonds"] and constraintTolerance is not None:
              system_params["constraintTolerance"] = constraintTolerance
          elif system_params["rigidWater"] and constraintTolerance is not None:
              system_params["constraintTolerance"] = constraintTolerance
  
          # Update system_params with any additional parameters provided
          system_params.update(kwargs)
          system = forcefield.createSystem(pdb.topology, **system_params)
          return system
  
      ########==================Integrator==================########
      # Integrator
      def _define_integrator(
          integrator_type="LangevinMiddle",
          temperature=300 * kelvin,
          friction=1.0 / picoseconds,
          timestep=0.004 * picoseconds,
          **kwargs,
      ):
          # Create a dictionary to hold integrator parameters
          integrator_params = {
              "temperature": temperature,
              "friction": friction,
              "timestep": timestep,
          }
  
          # Update integrator_params with any additional parameters provided
          integrator_params.update(kwargs)
  
          # Create the integrator
          if integrator_type == "LangevinMiddle":
              integrator = LangevinMiddleIntegrator(**integrator_params)
          elif integrator_type == "Verlet":
              integrator = VerletIntegrator(**integrator_params)
          elif integrator_type == "Brownian":
              integrator = BrownianIntegrator(**integrator_params)
          else:
              raise Exception("Integrator type not recognized")
  
          return integrator
  
      def _prompt_summary(self, query: str):
          prompt_template = """Your input is the original query. Your
                              task is to parse through the user query.
                              and provide a summary of the file path input,
                              the type of preprocessing needed (this is the
                              same as cleaning the file), the forcefield
                              used for the simulation,
                              the ensemble of the simulation, the integrator needed,
                              the number of steps, the timestep, the temperature,
                              and other instructions.
                              and follow the format "name: description.
  
                              File Path: what is the file path of the file
                              you are using? it must include a .cif or .pdb extension.
                              Preprocessing: what preprocessing is needed?
                              you can choose from the following: standard cleaning,
                              remove water, add hydrogens, add hydrogens and remove
                              water. The default is add hydrogens and remove water.
                              Forcefield: what forcefields are you using?
                              you can choose from the following: AMBER, CHARMM,
                              OPLS, GROMACS. Default -->  "amber14-all.xml, tip3p.xml".
                              Ensemble: what ensemble are you using?
                              you can choose from the following:
                              NPT, NVT, NVE. Default --> "NVT".
                              Integrator: what integrator are you using?
                              you can choose from the following:
                              Langevin, Verlet, Brownian.
                              The default depends on the ensemble
                              (NPT -> Langevin, NVT -> Langevin, NVE -> Verlet).
                              Number of Steps: how many steps
                              are you using? The default is 10000.
                              Timestep: what is the timestep?
                              Default --> "1 fs".
                              Temperature: what is the temperature?
                              Default --> "300 K".
                              Pressure: What is the pressure?
                              If NPT ensemble, the default is 1.0 bar, otherwise None.
                              Friction: what is the friction coefficient?
                              Default --> "1.0"
                              record_params: what parameters do you want to record?
                              you can choose from the following:
                              step, time, potentialEnergy, kineticEnergy,
                              totalEnergy, temperature, volume, density,
                              progress, remainingTime, speed, elapsedTime,
                              separator, systemMass, totalSteps, append.
                              Default --> ["step", "potentialEnergy", "temperature"].
                              Other Instructions: what other instructions do you have?
                              The default is none.
                              Example of the final output:
                              File Path: 1a1p.pdb
                              Preprocessing: standard cleaning
                              Forcefield: amber14-all.xml, tip3p.xml
                              Ensemble: NPT
                              Integrator: Langevin
                              Number of Steps: 10000
                              Timestep: 1 fs
                              Temperature: 300 K
                              Pressure: 1.0 bar
                              Friction: 1.0
                              record_params: ["step", "potentialEnergy", "temperature"]
                              Other Instructions: none
                              If there is not enough information in a category,
                              you may fill in with the default, but explicitly state so.
                              Here is the information:{query}"""
          prompt = PromptTemplate(template=prompt_template, input_variables=["query"])
          llm_chain = LLMChain(prompt=prompt, llm=self.llm)
  
          return llm_chain.run(" ".join(query))
  
      def _save_to_file(self, summary: str, filename: str):
          """Parse the summary string and
          save it to a file in JSON format."""
          # Split the summary into lines
          lines = summary.strip().split("\n")
  
          # Parse each line into a key and a value
          summary_dict = {}
          for line in lines:
              key, value = line.split(":")
              summary_dict[key.strip()] = value.strip()
  
          # Save the dictionary to a file
          with open(filename, "w") as f:
              json.dump(summary_dict, f)
  
          # add filename to registry
          file_description = "Simulation Parameters"
          self.path_registry.map_path(filename, filename, file_description)
  
      def _instruction_summary(self, query: str):
          summary = self._prompt_summary(query)
          self._save_to_file(summary, "simulation_parameters.json")
          return summary
  
      def _setup_simulation_from_json(self, file_name):
          # Open the json file and load the parameters
          with open(file_name, "r") as f:
              params = json.load(f)
          return params
  
      def _setup_and_run_simulation(self, query):
          # Load the force field
          # ask for inputs from the user
          params = self._setup_simulation_from_json(query)
  
          # forcefield key can be forcefield_files or Forcefield
          if "forcefield_files" in params:
              params["forcefield_files"] = (
                  params["forcefield_files"]
                  .replace("(default)", "")
                  .replace(" and ", ",")
                  .strip()
              )
              Forcefield_files = [
                  file.strip() for file in params["forcefield_files"].split(",")
              ]
              Forcefield = Forcefield_files[0]
              Water_model = Forcefield_files[1]
          else:
              params["Forcefield"] = (
                  params["Forcefield"]
                  .replace("(default)", "")
                  .replace(" and ", ",")
                  .strip()
              )
              Forcefield_files = [
                  file.strip() for file in params["Forcefield"].split(",")
              ]
              Forcefield = Forcefield_files[0]
              Water_model = Forcefield_files[1]
          print("Setting up forcefields :", Forcefield, Water_model)
          st.markdown("Setting up forcefields", unsafe_allow_html=True)
          # check if forcefields end in .xml
          if Forcefield.endswith(".xml") and Water_model.endswith(".xml"):
              forcefield = ForceField(Forcefield, Water_model)
          # adding forcefield to registry
  
          # Load the PDB file
          CleaningTools(self.path_registry)
          pdbfile = self.path_registry.get_mapped_path(params["File Path"])
          name = pdbfile.split(".")[0]
          end = pdbfile.split(".")[1]
          if end == "pdb":
              pdb = PDBFile(pdbfile)
          elif end == "cif":
              pdb = PDBxFile(pdbfile)
  
          modeller = Modeller(pdb.topology, pdb.positions)
          system = forcefield.createSystem(
              modeller.topology,
              nonbondedMethod=app.PME,
              nonbondedCutoff=1.0 * nanometers,
              constraints=app.PME,
          )
  
          _integrator = params["Integrator"].split(" ")[0].strip()
          _temp = params["Temperature"].split(" ")[0].strip()
          _friction_coef = params["Friction"].split(" ")[0].strip()
          _timestep = params["Timestep"].split(" ")[0].strip()
  
          if _integrator == "Langevin":
              print(
                  "Setting up Langevin integrator with Parameters:",
                  _temp,
                  "K",
                  _friction_coef,
                  "1/ps",
                  _timestep,
                  "fs",
              )
              st.markdown("Setting up Langevin integrator", unsafe_allow_html=True)
              if params["Ensemble"] == "NPT":
                  _pressure = params["Pressure"].split(" ")[0].strip()
                  system.addForce(MonteCarloBarostat(_pressure * bar, _temp * kelvin))
              integrator = LangevinIntegrator(
                  float(_temp) * kelvin,
                  float(_friction_coef) / picosecond,
                  float(_timestep) * femtoseconds,
              )
          elif _integrator == "Verlet":
              if params["Ensemble"] == "NPT":
                  _pressure = params["Pressure"].split(" ")[0].strip()
                  system.addForce(AndersenThermostat(_temp * kelvin, 1 / picosecond))
                  system.addForce(MonteCarloBarostat(_pressure * bar, _temp * kelvin))
                  print(
                      "Setting up Verlet integrator with Parameters:",
                      _timestep,
                      "fs",
                      _temp,
                      "K",
                      _pressure,
                      "bar",
                  )
              print("Setting up Verlet integrator with Parameters:", _timestep, "fs")
              st.markdown("Setting up Verlet integrator", unsafe_allow_html=True)
              integrator = VerletIntegrator(float(_timestep) * picoseconds)
  
          simulation = Simulation(modeller.topology, system, integrator)
          simulation.context.setPositions(modeller.positions)
          simulation.minimizeEnergy()
          # save initial positions to registry
          file_name = "initial_positions.pdb"
          with open(file_name, "w") as f:
              PDBFile.writeFile(
                  simulation.topology,
                  simulation.context.getState(getPositions=True).getPositions(),
                  f,
              )
          print("Initial Positions saved to initial_positions.pdb")
          simulation.reporters.append(PDBReporter(f"{name}.pdb", 1000))
          # reporter_args = {"reportInterval": 1000}
          reporter_args = {}
          params["record_params"] = ast.literal_eval(params["record_params"])
          for param in params["record_params"]:
              if param in [
                  "step",
                  "time",
                  "potentialEnergy",
                  "kineticEnergy",
                  "totalEnergy",
                  "temperature",
                  "volume",
                  "density",
                  "progress",
                  "remainingTime",
                  "speed",
                  "elapsedTime",
                  "separator",
                  "systemMass",
                  "totalSteps",
                  "append",
              ]:
                  # The params from the json file should be booleans
                  reporter_args[param] = True
          simulation.reporters.append(
              StateDataReporter(f"{name}.csv", 1000, **reporter_args)
          )
  
          simulation.step(int(params["Number of Steps"].split(" ")[0].strip()))
  
          # add filenames to registry
          file_name1 = "simulation_trajectory.pdb"
          file_description1 = "Simulation PDB, containing the simulation trajectory"
          self.path_registry.map_path(file_name1, f"{name}.pdb", file_description1)
          file_name2 = "simulation_data.csv"
          file_description2 = (
              "Simulation Data, containing step, potential energy, and temperature"
          )
          self.path_registry.map_path(file_name2, f"{name}.csv", file_description2)
  
          return simulation
  
      def _extract_parameters_path(self):
          """Check directory for parameters.json file."""
          # Check if there is a parameters.json file in the directory.
          if os.path.exists("simulation_parameters_summary.json"):
              return "simulation_parameters_summary.json"
          # If there's no exact match, check for
          # any JSON file that contains 'parameters' in its name.
          else:
              for file in os.listdir("."):
                  if "parameters" in file and file.endswith(".json"):
                      return file
              # If no matching file is found, raise an exception.
              raise ValueError("No parameters.json file found in directory.")
  
  
  class SetUpAndRunTool(BaseTool):
      name = "SetUpAndRunTool"
      description = """This tool will set up the simulation objects
                      and run the simulation.
                      It will ask for the parameters path.
                      input:  json file
                      """
      path_registry: Optional[PathRegistry]
  
      def __init__(
          self,
          path_registry: Optional[PathRegistry],
      ):
          super().__init__()
          self.path_registry = path_registry
  
      def _run(self, query: str) -> str:
          """Use the tool"""
          # find the parameters in the directory
          try:
              if self.path_registry is None:  # this should not happen
                  return "Registry not initialized"
              sim_fxns = SimulationFunctions(path_registry=self.path_registry)
              parameters = sim_fxns._extract_parameters_path()
  
          except ValueError as e:
              return (
                  str(e)
                  + """\nPlease use the Instruction summary tool with the
                  to create a parameters.json file in the directory."""
              )
          self.log("This are the parameters:")
          self.log(parameters)
          # print the parameters in json file
          with open(parameters) as f:
              params = json.load(f)
          for key, value in params.items():
              print(key, ":", value)
  
          self.log("Are you sure you want to run the simulation? (y/n)")
          response = input("yes or no: ")
          if response.lower() in ["yes", "y"]:
              sim_fxns._setup_and_run_simulation(parameters)
          else:
              return "Simulation interrupted due to human input"
          return "Simulation Completed, simulation trajectory and data files saved."
  
      def log(self, text, color="blue"):
          if color == "blue":
              print("\033[1;34m\t{}\033[00m".format(text))
          if color == "red":
              print("\033[31m\t{}\033[00m".format(text))
  
      async def _arun(self, query: str) -> str:
          """Use the tool asynchronously."""
          raise NotImplementedError("custom_search does not support async")
  
  
  #######==================System Configuration==================########
  # System Configuration
  class SetUpandRunFunctionInput(BaseModel):
      pdb_id: str
      forcefield_files: List[str]
      save: bool = Field(
          True,
          description=(
              (
                  "Set to 'True' (default) to save the log files and trajectories "
                  "of the simulation. "
                  "If set to 'False', "
                  "the simulation is considered as being in a testing "
                  "or preliminary scripting stage, utilizing default parameters and "
                  "results are not saved. "
                  "This second setting is ideal for initial experimentation or "
                  "basic script development before customizing the "
                  "script for final use."
              )
          ),
      )
  
      system_params: Dict[str, Any] = Field(
          {
              "nonbondedMethod": "NoCutoff",
              "nonbondedCutoff": "1 * nanometers",
              "ewaldErrorTolerance": None,
              "constraints": "None",
              "rigidWater": False,
              "constraintTolerance": None,
              "solvate": False,
          },
          description=(
              "Parameters for the openmm system. "
              "For nonbondedMethod, you can choose from the following:\n"
              "NoCutoff, CutoffNonPeriodic, CutoffPeriodic, Ewald, PME. "
              "If anything but NoCutoff is chosen,"
              "you have to include a nonbondedCutoff"
              "and a constrainTolerance.\n"
              "If PME is chosen,"
              "you have to include an ewaldErrorTolerance too."
              "For constraints, you can choose from the following:\n"
              "None, HBonds, AllBonds or OnlyWater."
              "For rigidWater, you can choose from the following:\n"
              "True, False.\n"
              "Finally, if you want to solvate the system, before the simulation,"
              "you can set solvate to True.\n"
              "Example1:\n"
              "{'nonbondedMethod': 'NoCutoff',\n"
              "'constraints': 'None',\n"
              "'rigidWater': False}\n"
              "Example2:\n"
              "{'nonbondedMethod': 'CutoffPeriodic',\n"
              "'nonbondedCutoff': 1.0,\n"
              "'constraints': 'HBonds',\n"
              "'rigidWater': True,\n"
              "'constraintTolerance': 0.00001,\n"
              "'solvate': True} "
          ),
      )
      integrator_params: Dict[str, Any] = Field(
          {
              "integrator_type": "LangevinMiddle",
              "Temperature": "300 * kelvin",
              "Friction": "1.0 / picoseconds",
              "Timestep": "0.002 * picoseconds",
              "Pressure": "1.0 * bar",
          },
          description="""Parameters for the openmm integrator.""",
      )
      simulation_params: Dict[str, Any] = Field(
          {
              "Ensemble": "NVT",
              "Number of Steps": 5000,
              "record_interval_steps": 100,
              "record_params": ["step", "potentialEnergy", "temperature"],
          },
          description="""Parameters for the openmm simulation.
          The ensemble can be NPT, NVT or NVE.
          The number of steps is the number of steps the simulation will run for.
          record_interval_steps is the number of steps between each record:
          hould be the number of steps divided by 100.
          The record_params is a list of parameters that will
          be recorded during the simulation The options are:
          [Step,Time,Speed,Progress,RemainingTime,ElapsedTime,
          PotentialEnergy,KineticEnergy,TotalEnergy,
          Temperature,Volume,Density]""",
      )
  
  
  #########===================================================================############
  
  
  class OpenMMSimulation:
      def __init__(
          self,
          input_params: SetUpandRunFunctionInput,
          path_registry: PathRegistry,
          save: bool,
          sim_id: str,
          pdb_id: str,
      ):
          self.params = input_params
          self.save = save
          self.sim_id = sim_id
          self.pdb_id = pdb_id
          self.int_params = (
              self.params.integrator_params
              if self.params.integrator_params is not None
              else {
                  "integrator_type": "LangevinMiddle",
                  "Temperature": 300 * kelvin,
                  "Friction": 1.0 / picoseconds,
                  "Timestep": 0.002 * picoseconds,
                  "Pressure": 1.0 * bar,
              }
          )
  
          self.sys_params = (
              self.params.system_params
              if self.params.system_params is not None
              else {
                  "nonbondedMethod": NoCutoff,
                  "nonbondedCutoff": 1 * nanometers,
                  "ewaldErrorTolerance": None,
                  "constraints": AllBonds,
                  "rigidWater": True,
                  "constraintTolerance": 0.000001,
                  "solvate": False,
              }
          )
  
          self.sim_params = (
              self.params.simulation_params
              if self.params.simulation_params is not None
              else {
                  "Ensemble": "NVT",
                  "Number of Steps": 5000,
                  "record_interval_steps": 100,
                  "record_params": ["step", "potentialEnergy", "temperature"],
              }
          )
  
          self.path_registry = path_registry
  
      def setup_system(self):
          print("Building system...")
          st.markdown("Building system", unsafe_allow_html=True)
          self.pdb_id = self.params.pdb_id
          self.pdb_path = self.path_registry.get_mapped_path(self.pdb_id)
          self.pdb = PDBFile(self.pdb_path)
          self.forcefield = ForceField(*self.params.forcefield_files)
          self.system = self._create_system(self.pdb, self.forcefield, **self.sys_params)
  
          if self.sys_params.get("nonbondedMethod", None) in [
              CutoffPeriodic,
              PME,
          ]:
              if self.sim_params["Ensemble"] == "NPT":
                  self.system.addForce(
                      MonteCarloBarostat(
                          self.int_params["Pressure"],
                          self.int_params["Temperature"],
                          self.sim_params.get("barostatInterval", 25),
                      )
                  )
  
      def setup_integrator(self):
          print("Setting up integrator...")
          st.markdown("Setting up integrator", unsafe_allow_html=True)
          int_params = self.int_params
          integrator_type = int_params.get("integrator_type", "LangevinMiddle")
  
          if integrator_type == "LangevinMiddle":
              self.integrator = LangevinMiddleIntegrator(
                  int_params["Temperature"],
                  int_params["Friction"],
                  int_params["Timestep"],
              )
          elif integrator_type == "LangevinIntegrator":
              self.integrator = LangevinIntegrator(
                  int_params["Temperature"],
                  int_params["Friction"],
                  int_params["Timestep"],
              )
          else:
              raise ValueError("Invalid integrator type")
  
          self.integrator.setConstraintTolerance(
              self.sys_params.get("constraintTolerance", 0.000001)
          )
  
      def create_simulation(self):
          print("Creating simulation...")
          st.markdown("Creating simulation", unsafe_allow_html=True)
          self.simulation = Simulation(
              self.modeller.topology,
              self.system,
              self.integrator,
              Platform.getPlatformByName("CPU"),
          )
          self.simulation.context.setPositions(self.modeller.positions)
  
          # TEMPORARY FILE MANAGEMENT OR PATH REGISTRY MAPPING
          if self.save:
              trajectory_name = self.path_registry.write_file_name(
                  type=FileType.RECORD,
                  record_type="TRAJ",
                  protein_file_id=self.pdb_id,
                  Sim_id=self.sim_id,
                  term="dcd",
              )
              topology_name = self.path_registry.write_file_name(
                  type=FileType.RECORD,
                  record_type="TOP",
                  protein_file_id=self.pdb_id,
                  Sim_id=self.sim_id,
                  term="pdb",
              )
  
              log_name = self.path_registry.write_file_name(
                  type=FileType.RECORD,
                  record_type="LOG",
                  protein_file_id=self.pdb_id,
                  Sim_id=self.sim_id,
                  term="txt",
              )
  
              traj_desc = (
                  f"Simulation trajectory for protein {self.pdb_id}"
                  f" and simulation {self.sim_id}"
              )
              top_desc = (
                  f"Simulation topology for protein"
                  f"{self.pdb_id} and simulation {self.sim_id}"
              )
              log_desc = (
                  f"Simulation state log for protein {self.pdb_id} "
                  f"and simulation {self.sim_id}"
              )
  
              self.simulation.reporters.append(
                  DCDReporter(
                      f"{trajectory_name}",
                      self.sim_params["record_interval_steps"],
                  )
              )
              self.simulation.reporters.append(
                  PDBReporter(
                      f"{topology_name}",
                      self.sim_params["record_interval_steps"],
                  )
              )
              self.simulation.reporters.append(
                  StateDataReporter(
                      f"{log_name}",
                      self.sim_params["record_interval_steps"],
                      step=True,
                      potentialEnergy=True,
                      temperature=True,
                      separator="\t",
                  )
              )
              self.registry_records = [
                  ("holder", f"files/records/{trajectory_name}", traj_desc),
                  ("holder", f"files/records/{log_name}", log_desc),
                  ("holder", f"files/records/{topology_name}", top_desc),
              ]
  
              # TODO add checkpoint too?
  
          else:
              self.simulation.reporters.append(
                  DCDReporter(
                      "temp_trajectory.dcd",
                      self.sim_params["record_interval_steps"],
                  )
              )
              self.simulation.reporters.append(
                  PDBReporter(
                      "temp_topology.pdb",
                      self.sim_params["record_interval_steps"],
                  )
              )
              self.simulation.reporters.append(
                  StateDataReporter(
                      "temp_log.txt",
                      self.sim_params["record_interval_steps"],
                      step=True,
                      potentialEnergy=True,
                      temperature=True,
                      separator="\t",
                  )
              )
  
      def _create_system(
          self,
          pdb,
          forcefield,
          nonbondedMethod="NoCutoff",
          nonbondedCutoff=None,
          ewaldErrorTolerance=None,
          constraints="None",
          rigidWater=False,
          constraintTolerance=None,
          solvate=False,
          **kwargs,
      ):
          # Create a dictionary to hold system parameters
          system_params = {
              "nonbondedMethod": nonbondedMethod,
              "constraints": constraints,
              "rigidWater": rigidWater,
          }
  
          # Set nonbondedCutoff if applicable Had to double if pre-commit
          if nonbondedMethod in ["PME", "CutoffNonPeriodic", "CutoffPeriodic"]:
              if nonbondedCutoff is not None:
                  system_params["nonbondedCutoff"] = nonbondedCutoff
  
          # Set ewaldErrorTolerance if PME is used
          if nonbondedMethod == "PME" and ewaldErrorTolerance is not None:
              system_params["ewaldErrorTolerance"] = ewaldErrorTolerance
  
          # Set constraintTolerance if constraints are used
          if constraints in ["HBonds", "AllBonds"] and constraintTolerance is not None:
              pass
          elif system_params["rigidWater"] and constraintTolerance is not None:
              pass
  
          # Update system_params with any additional parameters provided
          system_params.update(kwargs)
  
          # if use_constraint_tolerance:
          #    constraintTolerance = system_params.pop('constraintTolerance')
          self.modeller = Modeller(pdb.topology, pdb.positions)
          if solvate:
              try:
                  self.modeller.addSolvent(forcefield)
              except ValueError as e:
                  print("Error adding solvent", type(e).__name__, "–", e)
                  if "No Template for" in str(e):
                      raise ValueError(str(e))
              except AttributeError as e:
                  print("Error adding solvent: ", type(e).__name__, "–", e)
                  print("Trying to add solvent with 1 nm padding")
                  if "NoneType" and "value_in_unit" in str(e):
                      try:
                          self.modeller.addSolvent(forcefield, padding=1 * nanometers)
                      except Exception as e:
                          print("Error adding solvent", type(e).__name__, "–", e)
                          raise (e)
              system = forcefield.createSystem(self.modeller.topology, **system_params)
          else:
              system = forcefield.createSystem(self.modeller.topology, **system_params)
  
          return system
  
      def unit_to_string(self, unit):
          """Needed to convert units to strings for the script
          Otherwise internal __str()__ method makes the script
          not runnable"""
          return f"{unit.value_in_unit(unit.unit)}*{unit.unit.get_name()}"
  
      def _construct_script_content(
          self,
          pdb_path,
          forcefield_files,
          nonbonded_method,
          constraints,
          rigid_water,
          constraint_tolerance,
          nonbonded_cutoff,
          ewald_error_tolerance,
          hydrogen_mass,
          time_step,
          temperature,
          friction,
          ensemble,
          pressure,
          record_interval_steps,
          solvate,
          integrator_type,
      ):
          script_content = f"""
          # This script was generated by MDagent-Setup.
  
          from openmm import *
          from openmm.app import *
          from openmm.unit import *
  
          # Input Files
          pdb = PDBFile('{pdb_path}')
          forcefield = ForceField({forcefield_files})
  
          # System Configuration
          nonbondedMethod = {nonbonded_method}
          constraints = {constraints}
          rigidWater = {rigid_water}
          """
          if rigid_water and constraint_tolerance is not None:
              script_content += f"constraintTolerance = {constraint_tolerance}\n"
  
          # Conditionally add nonbondedCutoff
  
          if nonbonded_method != NoCutoff:
              script_content += f"nonbondedCutoff = {nonbonded_cutoff}\n"
          if nonbonded_method == PME:
              script_content += f"ewaldErrorTolerance = {ewald_error_tolerance}\n"
          if hydrogen_mass:
              script_content += f"hydrogenMass = {hydrogen_mass}\n"
  
          # ... other configurations ...
          script_content += f"""
          # Integration Options
          dt = {time_step}
          temperature = {temperature}
          friction = {friction}
          """
          if ensemble == "NPT":
              script_content += f"""
              pressure = {pressure}
              barostatInterval = {self.sim_params.get("barostatInterval", 25)}
              """
  
          # ... other integration options ...
          script_content += f"""
          # Simulation Options
          steps = {self.sim_params.get("Number of Steps", record_interval_steps)}
          equilibrationSteps = 1000
          platform = Platform.getPlatformByName('CPU')
          dcdReporter = DCDReporter('trajectory.dcd', 1000)
          pdbReporter = PDBReporter('trajectory.pdb', 1000)
          dataReporter = StateDataReporter('log.txt', {record_interval_steps},
              totalSteps=steps,
              step=True, speed=True, progress=True, elapsedTime=True, remainingTime=True,
              potentialEnergy=True, temperature=True, volume=True, density=True,
              separator='\t')
          checkpointReporter = CheckpointReporter('checkpoint.chk', 5000)
  
          # Minimize and Equilibrate
          # ... code for minimization and equilibration ...
  
          # Simulate
  
          print('Building system...')
          modeller = Modeller(pdb.topology, pdb.positions)
          """
          if solvate:
              script_content += (
                  """modeller.addSolvent(forcefield, padding=1*nanometers)"""
              )
  
          if nonbonded_method == NoCutoff:
              if hydrogen_mass:
                  script_content += """
              system = forcefield.createSystem(modeller.topology,
              nonbondedMethod=nonbondedMethod, constraints=constraints,
              rigidWater=rigidWater, hydrogenMass=hydrogenMass)
              """
              else:
                  script_content += """
              system = forcefield.createSystem(modeller.topology,
              nonbondedMethod=nonbondedMethod, constraints=constraints,
              rigidWater=rigidWater)
              """
          if nonbonded_method == CutoffNonPeriodic or nonbonded_method == CutoffPeriodic:
              if hydrogen_mass:
                  script_content += """
                  system = forcefield.createSystem(modeller.topology,
                  nonbondedMethod=nonbondedMethod, nonbondedCutoff=nonbondedCutoff,
                  constraints=constraints, rigidWater=rigidWater,
                  hydrogenMass=hydrogenMass)
              """
              else:
                  script_content += """
                  system = forcefield.createSystem(modeller.topology,
                  nonbondedMethod=nonbondedMethod, nonbondedCutoff=nonbondedCutoff,
                  constraints=constraints, rigidWater=rigidWater)
              """
          if nonbonded_method == PME:
              if hydrogen_mass:
                  script_content += """
              system = forcefield.createSystem(modeller.topology,
              nonbondedMethod=nonbondedMethod,
              nonbondedCutoff=nonbondedCutoff, ewaldErrorTolerance=ewaldErrorTolerance,
              constraints=constraints, rigidWater=rigidWater, hydrogenMass=hydrogenMass)
              """
              else:
                  script_content += """
              system = forcefield.createSystem(modeller.topology,
              nonbondedMethod=nonbondedMethod,
              nonbondedCutoff=nonbondedCutoff, ewaldErrorTolerance=ewaldErrorTolerance,
              constraints=constraints, rigidWater=rigidWater)
              """
          if ensemble == "NPT":
              script_content += """
              system.addForce(MonteCarloBarostat(pressure, temperature, barostatInterval))
              """
  
          if integrator_type == "LangevinMiddle" and constraints != "None":
              script_content += """
          integrator = LangevinMiddleIntegrator(temperature, friction, dt)
          integrator.setConstraintTolerance(constraintTolerance)
          simulation = Simulation(modeller.topology, system, integrator, platform)
          simulation.context.setPositions(modeller.positions)
          """
          if integrator_type == "LangevinMiddle" and constraints == "None":
              script_content += """
              integrator = LangevinMiddleIntegrator(temperature, friction, dt)
              simulation = Simulation(modeller.topology, system, integrator, platform)
              simulation.context.setPositions(modeller.positions)
          """
  
          script_content += """
          # Minimize and Equilibrate
  
          print('Performing energy minimization...')
          simulation.minimizeEnergy()
          print('Equilibrating...')
          simulation.context.setVelocitiesToTemperature(temperature)
          simulation.step(equilibrationSteps)
  
          # Simulate
  
          print('Simulating...')
          simulation.reporters.append(dcdReporter)
          simulation.reporters.append(pdbReporter)
          simulation.reporters.append(dataReporter)
          simulation.reporters.append(checkpointReporter)
          simulation.currentStep = 0
          simulation.step(steps)
  """
          return script_content
  
      def write_standalone_script(self, filename="reproduce_simulation.py"):
          """Extracting parameters from the class instance
          Inspired by the code snippet provided from openmm-setup
          https://github.com/openmm/openmm-setup
          """
  
          pdb_path = self.pdb_path
          forcefield_files = ", ".join(
              f"'{file}'" for file in self.params["forcefield_files"]
          )
          nonbonded_method = self.sys_params.get("nonbondedMethod", NoCutoff)
          nbCo = self.sys_params.get("nonbondedCutoff", 1 * nanometers)
          nonbonded_cutoff = self.unit_to_string(nbCo)
          constraints = self.sys_params.get("constraints", "None")
          rigid_water = self.sys_params.get("rigidWater", False)
          ewald_error_tolerance = self.sys_params.get("ewaldErrorTolerance", 0.0005)
          constraint_tolerance = self.sys_params.get("constraintTolerance", None)
          hydrogen_mass = self.sys_params.get("hydrogenMass", None)
          solvate = self.sys_params.get("solvate", False)
  
          integrator_type = self.int_params.get("integrator_type", "LangevinMiddle")
          friction = self.int_params.get("Friction", 1.0 / picoseconds)
          friction = f"{friction.value_in_unit(friction.unit)}{friction.unit.get_name()}"
          _temp = self.int_params.get("Temperature", 300 * kelvin)
          temperature = self.unit_to_string(_temp)
  
          t_step = self.int_params.get("Timestep", 0.004 * picoseconds)
          time_step = self.unit_to_string(t_step)
          press = self.int_params.get("Pressure", 1.0 * bar)
          pressure = self.unit_to_string(press)
          ensemble = self.sim_params.get("Ensemble", "NVT")
          self.sim_params.get("Number of Steps", 10000)
          record_interval_steps = self.sim_params.get("record_interval_steps", 1000)
  
          script_content = self._construct_script_content(
              pdb_path,
              forcefield_files,
              nonbonded_method,
              constraints,
              rigid_water,
              constraint_tolerance,
              nonbonded_cutoff,
              ewald_error_tolerance,
              hydrogen_mass,
              time_step,
              temperature,
              friction,
              ensemble,
              pressure,
              record_interval_steps,
              solvate,
              integrator_type,
          )
  
          # Remove leading spaces for proper formatting
          def remove_leading_spaces(text):
              lines = text.split("\n")
              stripped_lines = [line.lstrip() for line in lines]
              return "\n".join(stripped_lines)
  
          script_content = remove_leading_spaces(script_content)
          script_content = textwrap.dedent(script_content).strip()
  
          # Write to file
          directory = "files/simulations"
          if not os.path.exists(directory):
              os.makedirs(directory)
  
          with open(f"{directory}/{filename}", "w") as file:
              file.write(script_content)
  
          print(f"Standalone simulation script written to {directory}/{filename}")
          st.markdown("Standalone simulation script written", unsafe_allow_html=True)
  
      def run(self):
          # Minimize and Equilibrate
          print("Performing energy minimization...")
          st.markdown("Performing energy minimization", unsafe_allow_html=True)
  
          self.simulation.minimizeEnergy()
          print("Minimization complete!")
          top_name = f"files/pdb/{self.sim_id}_initial_positions.pdb"
          top_description = f"Initial positions for simulation {self.sim_id}"
          with open(top_name, "w") as f:
              PDBFile.writeFile(
                  self.simulation.topology,
                  self.simulation.context.getState(getPositions=True).getPositions(),
                  f,
              )
          self.path_registry.map_path(f"top_{self.sim_id}", top_name, top_description)
          print("Initial Positions saved to initial_positions.pdb")
          st.markdown("Minimization complete! Equilibrating...", unsafe_allow_html=True)
          print("Equilibrating...")
          _temp = self.int_params["Temperature"]
          self.simulation.context.setVelocitiesToTemperature(_temp)
          _eq_steps = self.sim_params.get("equilibrationSteps", 1000)
          self.simulation.step(_eq_steps)
          # Simulate
          print("Simulating...")
          st.markdown("Simulating...", unsafe_allow_html=True)
          self.simulation.currentStep = 0
          self.simulation.step(self.sim_params["Number of Steps"])
          print("Done!")
          st.markdown("Done!", unsafe_allow_html=True)
          if not self.save:
              if os.path.exists("temp_trajectory.dcd"):
                  os.remove("temp_trajectory.dcd")
              if os.path.exists("temp_log.txt"):
                  os.remove("temp_log.txt")
              if os.path.exists("temp_checkpoint.chk"):
                  os.remove("temp_checkpoint.chk")
  
          return "Simulation done!"
  
  
  class SetUpandRunFunction(BaseTool):
      name: str = "SetUpandRunFunction"
      description: str = (
          "This tool will set up and run a short simulation of a protein. "
          "Then will write a standalone script that can be used "
          "to reproduce the simulation or change accordingly for "
          "a more elaborate simulation. It only runs short simulations because, "
          "if there are errors, you can try again changing the input"
      )
  
      args_schema: Type[BaseModel] = SetUpandRunFunctionInput
  
      path_registry: Optional[PathRegistry]
  
      def __init__(self, path_registry: Optional[PathRegistry]):
          super().__init__()
          self.path_registry = path_registry
  
      def _run(self, **input_args):
          if self.path_registry is None:
              return "Path registry not initialized"
          input = self.check_system_params(input_args)
          error = input.get("error", None)
          if error:
              print(f"error found: {error}")
              return error
  
          try:
              pdb_id = input["pdb_id"]
              # check if pdb_id is in the registry or as 1XYZ_112233 format
              if pdb_id not in self.path_registry.list_path_names():
                  return "No pdb_id found in input, use the file id not the file name"
          except KeyError:
              return "No pdb_id found in input"
          try:
              save = input["save"]  # either this simulation
              # to save or not the output files from this simulation
          except KeyError:
              save = True
              print(
                  "No 'save' key found in input, setting to True. "
                  "Record files will be deleted after script is written."
              )
          try:
              file_name = self.path_registry.write_file_name(
                  type=FileType.SIMULATION,
                  type_of_sim=input["simulation_params"]["Ensemble"],
                  protein_file_id=pdb_id,
              )
  
              sim_id = self.path_registry.get_fileid(file_name, FileType.SIMULATION)
          except Exception as e:
              print(f"An exception was found: {str(e)}.")
              return f"An exception was found trying to write the filenames: {str(e)}."
          try:
              openmmsim = OpenMMSimulation(
                  input, self.path_registry, save, sim_id, pdb_id
              )
              openmmsim.setup_system()
              openmmsim.setup_integrator()
              openmmsim.create_simulation()
  
              print("simulation set!")
              st.markdown("simulation set!", unsafe_allow_html=True)
          except ValueError as e:
              msg = str(e) + f"This were the inputs {input_args}"
              if "No template for" in msg:
                  msg += (
                      "This error is likely due to non standard residues "
                      "in the protein, if you havent done it yet, try "
                      "cleaning the pdb file using the cleaning tool"
                  )
              return msg
          except FileNotFoundError:
              return f"File not found, check File id. This were the inputs {input_args}"
          except OpenMMException as e:
              return f"OpenMM Exception: {str(e)}. This were the inputs {input_args}"
          try:
              openmmsim.run()
          except Exception as e:
              return (
                  f"An exception was found: {str(e)}. Not a problem, thats one "
                  "purpose of this tool: to run a short simulation to check for correct "
                  "initialization. "
                  ""
                  "Try a) with different parameters like "
                  "nonbondedMethod, constraints, etc \n or\n"
                  "b) clean file inputs depending on error "
              )
          try:
              openmmsim.write_standalone_script(filename=file_name)
              self.path_registry.map_path(
                  sim_id,
                  f"files/simulations/{file_name}",
                  f"Basic Simulation of Protein {pdb_id}",
              )
              if save:
                  records = openmmsim.registry_records
                  # move record files to files/records/
                  print(os.listdir("."))
                  if not os.path.exists("files/records"):
                      os.makedirs("files/records")
                  for record in records:
                      os.rename(record[1].split("/")[-1], f"{record[1]}")
                  for record in records:
                      record_list = list(record)
                      record_list[0] = self.path_registry.get_fileid(
                          record_list[1].split("/")[-1], FileType.RECORD
                      )
                      record = tuple(record_list)
                      self.path_registry.map_path(*record)
              return (
                  "Simulation done! \n Summary: \n"
                  "Record files written to files/records/ with IDs and descriptions: "
                  f"{[(record[0],record[2]) for record in records]}\n"
                  "Standalone script written to files/simulations/ with ID: "
                  f"{sim_id}.\n"
                  f"The initial topology file ID is top_{sim_id} saved in files/pdb/"
              )
          except Exception as e:
              print(f"An exception was found: {str(e)}.")
              return f"An exception was found trying to write the filenames: {str(e)}."
  
      def _parse_cutoff(self, cutoff):
          # Check if cutoff is already an OpenMM Quantity (has a unit)
          possible_units = ["nm", "nanometer", "nanometers", "angstrom", "angstroms", "a"]
  
          if isinstance(cutoff, unit.Quantity):
              return cutoff
  
          # Convert to string in case it's not (e.g., int or float)
          cutoff = str(cutoff)
          if cutoff[-1] == "s":
              cutoff = cutoff[:-1]
  
          # Remove spaces and convert to lowercase for easier parsing
          cutoff = cutoff.replace(" ", "").lower()
          if cutoff.endswith("s"):
              cutoff = cutoff[:-1]
          # Check for multiplication symbol and split if necessary
          if "*" in cutoff:
              # Split on the '*' and extract the numerical part and the unit part
              num_part, unit_part = cutoff.split("*")
  
              # Convert the numerical part to a float
              num_value = float(num_part)
          else:
              # If there is no '*', it's either a number or a string like "1nm"
              # Attempt to convert directly to float; if it fails,
              # it must have a unit like "nm" or "angstrom"
              try:
                  num_value = float(cutoff)
                  unit_part = "nm"
              except ValueError:
                  for possible_unit in possible_units:
                      if possible_unit in cutoff:
                          num_value = float(cutoff.replace(possible_unit, ""))
                          unit_part = possible_unit
                          break
                  else:
                      # Use regular expression to extract the
                      # numerical part and the unit part
                      match = re.match(r"([+-]?[0-9]*\.?[0-9]+)([a-zA-Z]*)", cutoff)
                      if match:
                          num_part, unit_part = match.groups()
                          raise ValueError(
                              f"""Unknown unit for nonbondedCutoff
                                           got {unit_part}. Try using nm or angstroms as
                                           value * unit."""
                          )
  
          # Now convert the unit part to an OpenMM unit
          if unit_part in ["nm", "nanometer", "nanometers"]:
              return num_value * unit.nanometers
          elif unit_part in ["angstrom", "angstroms", "a"]:
              return num_value * unit.angstroms
  
          else:
              # If the unit is not recognized, raise an error
              raise ValueError(
                  f"""Unknown unit for nonbondedCutoff
                      got {unit_part}. Try using nm or angstroms as
                      value * unit."""
              )
  
      def _parse_parameter(self, parameter, default_unit, possible_units):
          """
          Parse a parameter and return it as an OpenMM Quantity with the correct unit.
  
          Args:
              parameter (float, str, or unit.Quantity): The input parameter value
              default_unit (unit.Unit): The default unit to use if none is provided
              possible_units (dict): A mapping of strings to their respective unit objects
  
          Returns:
              unit.Quantity: The parameter as an OpenMM Quantity with the correct unit.
          """
          error_msg = ""
          if isinstance(parameter, unit.Quantity):
              return parameter, error_msg
  
          # Convert to string in case it's not (e.g., int or float)
          parameter_str = str(parameter)
  
          # Remove spaces and convert to lowercase for easier parsing
          parameter_str = parameter_str.replace(" ", "").lower()
  
          # Check for multiplication symbol and split if necessary
          # e.g. "1*kelvin" or "1*ps^-1"
          if "*" in parameter_str:
              num_part, unit_part = parameter_str.split("*")
              num_value = float(num_part)
          elif "poundforce/inch^2" in parameter_str:
              num_value = float(parameter_str.replace("poundforce/inch^2", ""))
              unit_part = "poundforce/inch^2"
          # Check for division symbol and split if necessary
          # e.g. "1/ps" or "1/ps^-1"
          elif "/" in parameter_str:
              num_part, unit_part = parameter_str.split("/")
              num_value = float(num_part)
              unit_part = "/" + unit_part
          elif "^-1" in parameter_str:
              parameter_str = parameter_str.replace("^-1", "")
              match = re.match(r"^(\d+(?:\.\d+)?)([a-zA-Z]+)$", parameter_str)
              num_value = float(match.group(1))
              unit_part = "/" + match.group(2)
          else:
              # Attempt to convert directly to float; if it fails,
              # it must have a unit like "K", "ps", etc.
              try:
                  num_value = float(parameter_str)
                  unit_part = default_unit
              except ValueError:
                  match = re.match(r"([+-]?[0-9]*\.?[0-9]+)([a-zA-Z]*)", parameter_str)
                  if match:
                      num_part, unit_part = match.groups()
                      num_value = float(num_part)
                  else:
                      error_msg += f"Invalid format for parameter: '{parameter_str}'."
  
          # Convert the unit part to an OpenMM unit
          if unit_part.lower() in possible_units:
              return num_value * possible_units[unit_part.lower()], error_msg
          else:
              # If the unit is not recognized, raise an error
              error_msg += f"""Unknown unit '{unit_part}' for parameter.
              Valid units include: {list(possible_units.keys())}."""
  
              return parameter, error_msg
  
      def parse_temperature(self, temperature):
          possible_units = {
              "k": unit.kelvin,
              "kelvin": unit.kelvin,
          }
          return self._parse_parameter(temperature, "k", possible_units)
  
      def parse_friction(self, friction):
          possible_units = {
              "/ps": (1 / unit.picoseconds),
              "/picosecond": (1 / unit.picoseconds),
              "/picoseconds": (1 / unit.picoseconds),
              "picosecond^-1": (1 / unit.picoseconds),
              "picoseconds^-1": (1 / unit.picoseconds),
              "/ps^-1": (1 / unit.picoseconds),
              "ps^-1": (1 / unit.picoseconds),
              "1*ps^-1": (1 / unit.picoseconds),
          }
          return self._parse_parameter(friction, "1/ps", possible_units)
  
      def parse_timestep(self, timestep):
          possible_units = {
              "ps": unit.picoseconds,
              "picosecond": unit.picoseconds,
              "picoseconds": unit.picoseconds,
              "fs": unit.femtoseconds,
              "femtosecond": unit.femtoseconds,
              "femtoseconds": unit.femtoseconds,
              "ns": unit.nanoseconds,
              "nanosecond": unit.nanoseconds,
              "nanoseconds": unit.nanoseconds,
          }
          return self._parse_parameter(timestep, "ps", possible_units)
  
      def parse_pressure(self, pressure):
          possible_units = {
              "bar": unit.bar,
              "atm": unit.atmospheres,
              "atmosphere": unit.atmospheres,
              "pascal": unit.pascals,
              "pascals": unit.pascals,
              "pa": unit.pascals,
              "poundforce/inch^2": unit.psi,
              "psi": unit.psi,
          }
          return self._parse_parameter(pressure, "bar", possible_units)
  
      def _process_parameters(self, user_params, param_type="system_params"):
          """
          Process user provided parameters,
          converting strings to openmm objects if necessary.
          """
          error_msg = ""
          processed_params = {}
          if param_type == "system_params":
              for key, value in user_params.items():
                  if key == "nonbondedMethod" or key == "nonbondedmethod":
                      if value == "NoCutoff":
                          processed_params[key] = NoCutoff
                      elif value == "PME":
                          processed_params[key] = PME
                      elif value == "CutoffPeriodic":
                          processed_params[key] = CutoffPeriodic
                      elif value == "CutoffNonPeriodic":
                          processed_params[key] = CutoffNonPeriodic
                      elif value == "Ewald":
                          processed_params[key] = Ewald
                      else:
                          # Assume it's already an openmm object
                          processed_params[key] = value
                  if key == "nonbondedCutoff" or key == "nonbondedcutoff":
                      try:
                          processed_params[key] = self._parse_cutoff(value)
                      except ValueError as e:
                          error_msg += f"Invalid nonbondedCutoff: {e}. \n"
                  if key == "ewaldErrorTolerance" or key == "ewalderrortolerance":
                      try:
                          processed_params[key] = float(value)
                      except TypeError as e:
                          error_msg += (
                              f"Invalid ewaldErrorTolerance: {e}. "
                              "If you are using null or None, "
                              "just dont include it "
                              "as part of the parameters.\n"
                          )
                  if key == "constraints":
                      try:
                          if type(value) == str:
                              if value == "None":
                                  processed_params[key] = None
                              elif value == "HBonds":
                                  processed_params[key] = HBonds
                              elif value == "AllBonds":
                                  processed_params[key] = AllBonds
                              elif value == "HAngles":
                                  processed_params[key] = HAngles
                              else:
                                  error_msg += (
                                      f"Invalid constraints: Got {value}. "
                                      "Try using None, HBonds, AllBonds or "
                                      "HAngles\n"
                                  )
                          else:
                              processed_params[key] = value
                      except TypeError as e:
                          error_msg += (
                              f"Invalid constraints: {e}. If you are using "
                              "null or None, just dont include as "
                              "part of the parameters.\n"
                          )
                  if key == "rigidWater" or key == "rigidwater":
                      if type(value) == bool:
                          processed_params[key] = value
                      elif value == "True":
                          processed_params[key] = True
                      elif value == "False":
                          processed_params[key] = False
                      else:
                          error_msg += (
                              f"Invalid rigidWater: got {value}. "
                              "Try using True or False.\n"
                          )
                  if key == "constraintTolerance" or key == "constrainttolerance":
                      try:
                          processed_params[key] = float(value)
                      except ValueError as e:
                          error_msg += f"Invalid constraintTolerance: {e}."
                      except TypeError as e:
                          error_msg += (
                              f"Invalid constraintTolerance: {e}. If "
                              "constraintTolerance is null or None, "
                              "just dont include as part of "
                              "the parameters.\n"
                          )
                  if key == "solvate":
                      try:
                          if type(value) == bool:
                              processed_params[key] = value
                          elif value == "True":
                              processed_params[key] = True
                          elif value == "False":
                              processed_params[key] = False
                          else:
                              error_msg += (
                                  f"Invalid solvate: got {value}. "
                                  "Use either True or False.\n"
                              )
                      except TypeError as e:
                          error_msg += (
                              f"Invalid solvate: {e}. If solvate is null or "
                              "None, just dont include as part of "
                              "the parameters.\n"
                          )
  
              return processed_params, error_msg
          if param_type == "integrator_params":
              for key, value in user_params.items():
                  if key == "integrator_type" or key == "integratortype":
                      if value == "LangevinMiddle" or value == LangevinMiddleIntegrator:
                          processed_params[key] = "LangevinMiddle"
                      elif value == "Langevin" or value == LangevinIntegrator:
                          processed_params[key] = "Langevin"
                      elif value == "Verlet" or value == VerletIntegrator:
                          processed_params[key] = "Verlet"
                      elif value == "Brownian" or value == BrownianIntegrator:
                          processed_params[key] = "Brownian"
                      else:
                          error_msg += (
                              f"Invalid integrator_type: got {value}. "
                              "Try using LangevinMiddle, Langevin, "
                              "Verlet, or Brownian.\n"
                          )
                  if key == "Temperature" or key == "temperature":
                      temperature, msg = self.parse_temperature(value)
                      processed_params[key] = temperature
                      error_msg += msg
                  if key == "Friction" or key == "friction":
                      friction, msg = self.parse_friction(value)
                      processed_params[key] = friction
                      error_msg += msg
                  if key == "Timestep" or key == "timestep":
                      timestep, msg = self.parse_timestep(value)
                      processed_params[key] = timestep
                      error_msg += msg
                  if key == "Pressure" or key == "pressure":
                      pressure, msg = self.parse_pressure(value)
                      processed_params[key] = pressure
                      error_msg += msg
  
              return processed_params, error_msg
          if param_type == "simulation_params":
              for key, value in user_params.items():
                  if key == "Ensemble" or key == "ensemble":
                      if value == "NPT":
                          processed_params[key] = "NPT"
                      elif value == "NVT":
                          processed_params[key] = "NVT"
                      elif value == "NVE":
                          processed_params[key] = "NVE"
                      else:
                          error_msg += (
                              f"Invalid Ensemble. got {value}. "
                              "Try using NPT, NVT, or NVE.\n"
                          )
  
                  if key == "Number of Steps" or key == "number of steps":
                      processed_params[key] = int(value)
                  if key == "record_interval_steps" or key == "record interval steps":
                      processed_params[key] = int(value)
                  if key == "record_params" or key == "record params":
                      processed_params[key] = value
              return processed_params, error_msg
  
      def check_system_params(cls, values):
          """Check that the system parameters are valid."""
          # lowercase all keys in the dictionary
          error_msg = ""
          values = {k.lower(): v for k, v in values.items()}
  
          system_params = values.get("system_params")
          if system_params:
              system_params, msg = cls._process_parameters(
                  system_params, param_type="system_params"
              )
              if msg != "":
                  error_msg += msg
          else:
              system_params = {
                  "nonbondedMethod": NoCutoff,
                  "nonbondedCutoff": 1 * nanometers,
                  "ewaldErrorTolerance": None,
                  "constraints": AllBonds,
                  "rigidWater": True,
                  "constraintTolerance": 0.00001,
                  "solvate": False,
              }
          integrator_params = values.get("integrator_params")
          if integrator_params:
              integrator_params, msg = cls._process_parameters(
                  integrator_params, param_type="integrator_params"
              )
              if msg != "":
                  error_msg += msg
          else:
              integrator_params = {
                  "integrator_type": "LangevinMiddle",
                  "Temperature": 300 * kelvin,
                  "Friction": 1.0 / picoseconds,
                  "Timestep": 0.004 * picoseconds,
                  "Pressure": 1.0 * bar,
              }
          simulation_params = values.get("simulation_params")
          if simulation_params is None:
              simulation_params = {
                  "Ensemble": "NVT",
                  "Number of Steps": 10000,
                  "record_interval_steps": 100,
                  "record_params": ["step", "potentialEnergy", "temperature"],
              }
          # lowercase all keys in the dictionary
  
          # system_params = {k.lower(): v for k, v in system_params.items()}
          # integrator_params = {k.lower(): v for k, v in integrator_params.items()}
          # simulation_params = {k.lower(): v for k, v in simulation_params.items()}
  
          nonbondedMethod = system_params.get("nonbondedMethod")
          nonbondedCutoff = system_params.get("nonbondedCutoff")
          ewaldErrorTolerance = system_params.get("ewaldErrorTolerance")
          constraints = system_params.get("constraints")
          rigidWater = system_params.get("rigidWater")
          constraintTolerance = system_params.get("constraintTolerance")
          methods_with_cutoff = {
              "PME",
              "CutoffNonPeriodic",
              "CutoffPeriodic",
              "Ewald",
              PME,
              CutoffNonPeriodic,
              CutoffPeriodic,
              Ewald,
          }
          constraints_with_tolerance = {
              "HBonds",
              "AllBonds",
              "OnlyWater",
              HBonds,
              AllBonds,
          }
  
          if nonbondedMethod in methods_with_cutoff and nonbondedCutoff is None:
              error_msg += """nonbondedCutoff must be specified if
                          nonbondedMethod is not NoCutoff\n"""
          if nonbondedMethod in {"PME", PME} and ewaldErrorTolerance is None:
              error_msg += """ewaldErrorTolerance must be specified when
              nonbondedMethod is PME\n"""
          if constraints in constraints_with_tolerance and constraintTolerance is None:
              error_msg += """constraintTolerance must be specified when
                           constraints is HBonds or AllBonds"""
          if rigidWater and constraintTolerance is None:
              error_msg = "constraintTolerance must be specified if rigidWater is True"
  
          """Checking if the file is in the path"""
          pdb_id = values.get("pdb_id")
          if not pdb_id:
              error_msg += "The pdb id is not present in the inputs"
  
          """Validating the forcefield files and Integrator"""
  
          integrator_type = integrator_params.get("integrator_type")
          if integrator_type not in ["LangevinMiddle", "Verlet", "Brownian"]:
              error_msg += """integrator_type must be one of the following:
                               LangevinMiddle, Verlet, Brownian\n"""
          if integrator_type == "LangevinMiddle":
              friction = integrator_params.get("Friction")
              if friction is None:
                  error_msg += """friction must be specified when
                              integrator_type is LangevinMiddle\n"""
              timestep = integrator_params.get("Timestep")
              if timestep is None:
                  error_msg += """timestep must be specified when
                              integrator_type is LangevinMiddle\n"""
              temp = integrator_params.get("Temperature")
              if temp is None:
                  error_msg += """temperature must be specified when
                  integrator_type is LangevinMiddle\n"""
  
          if integrator_type == "Verlet":
              timestep = integrator_params.get("Timestep")
              if timestep is None:
                  error_msg += """timestep must be specified when
                              integrator_type is Verlet\n"""
          if integrator_type == "Brownian":
              temperature = integrator_params.get("Temperature")
              if temperature is None:
                  error_msg += """temperature must be specified when
                      integrator_type is Brownian\n"""
  
          # forcefield
          forcefield_files = values.get("forcefield_files")
          if forcefield_files is None or forcefield_files is []:
              print("Setting default forcefields")
              st.markdown("Setting default forcefields", unsafe_allow_html=True)
              forcefield_files = ["amber14-all.xml", "amber14/tip3pfb.xml"]
          elif len(forcefield_files) == 0:
              print("Setting default forcefields v2")
              st.markdown("Setting default forcefields", unsafe_allow_html=True)
              forcefield_files = ["amber14-all.xml", "amber14/tip3pfb.xml"]
          else:
              for file in forcefield_files:
                  if file not in FORCEFIELD_LIST:
                      error_msg += "The forcefield file is not present"
  
          save = values.get("save", True)
          if type(save) != bool:
              error_msg += "save must be a boolean value"
  
          if error_msg != "":
              return {
                  "error": error_msg
                  + "\n Correct this and try again. \n Everthing else is fine"
              }
          values = {
              "pdb_id": pdb_id,
              "forcefield_files": forcefield_files,
              "save": save,
              "system_params": system_params,
              "integrator_params": integrator_params,
              "simulation_params": simulation_params,
          }
          # if no error, return the values
          return values
  
      async def _arun(self, query: str) -> str:
          """Use the tool asynchronously."""
          raise NotImplementedError("custom_search does not support async")
  
  
  def create_simulation_input(pdb_path, forcefield_files):
      """
      This function takes a PDB file path and a list of forcefield files.
      It creates and returns a PDBFile and ForceField object.
      The forcefield_files list can contain one or more files.
      If only one file is provided, it assumes that the file includes
      both the forcefield and the water model if needed.
  
      Parameters:
      pdb_path (str): The file path to the PDB file.
      forcefield_files (list of str): A list of file paths to the forcefield XML files.
  
      Returns:
      tuple: A tuple containing the PDBFile and ForceField objects.
      """
  
      # Load the PDB file
  
      pdb_path.split(".")[0]
      end = pdb_path.split(".")[1]
      if end == "pdb":
          pdb = PDBFile(pdb_path)
      elif end == "cif":
          pdb = PDBxFile(pdb_path)
  
      # Clean up forcefield files list and remove any empty strings
      forcefield_files = (
          forcefield_files.replace("(default)", "").replace(" and ", ",").strip()
      )
      Forcefield_files = [file.strip() for file in forcefield_files.split(",")]
      Forcefield = Forcefield_files[0]
      Water_model = Forcefield_files[1]
      # check if they are part of the list
      if Forcefield not in FORCEFIELD_LIST:
          raise Exception("Forcefield not recognized")
      if Water_model not in FORCEFIELD_LIST:
          raise Exception("Water model not recognized")
  
      forcefield = ForceField(Forcefield, Water_model)
      # TODO  Not all forcefields require water model
  
      return pdb, forcefield
  
apis: langchain.prompts.PromptTemplate langchain.chat_models.ChatOpenAI langchain.chains.LLMChain
extract_api: [((20314, 20720), 'pydantic.Field', 'Field', (['(True)'], {'description': '"""Set to \'True\' (default) to save the log files and trajectories of the simulation. If set to \'False\', the simulation is considered as being in a testing or preliminary scripting stage, utilizing default parameters and results are not saved. This second setting is ideal for initial experimentation or basic script development before customizing the script for final use."""'}), '(True, description=\n    "Set to \'True\' (default) to save the log files and trajectories of the simulation. If set to \'False\', the simulation is considered as being in a testing or preliminary scripting stage, utilizing default parameters and results are not saved. This second setting is ideal for initial experimentation or basic script development before customizing the script for final use."\n    )\n', (20319, 20720), False, 'from pydantic import BaseModel, Field\n'), ((20979, 22020), 'pydantic.Field', 'Field', (["{'nonbondedMethod': 'NoCutoff', 'nonbondedCutoff': '1 * nanometers',\n    'ewaldErrorTolerance': None, 'constraints': 'None', 'rigidWater': False,\n    'constraintTolerance': None, 'solvate': False}"], {'description': '"""Parameters for the openmm system. For nonbondedMethod, you can choose from the following:\nNoCutoff, CutoffNonPeriodic, CutoffPeriodic, Ewald, PME. If anything but NoCutoff is chosen,you have to include a nonbondedCutoffand a constrainTolerance.\nIf PME is chosen,you have to include an ewaldErrorTolerance too.For constraints, you can choose from the following:\nNone, HBonds, AllBonds or OnlyWater.For rigidWater, you can choose from the following:\nTrue, False.\nFinally, if you want to solvate the system, before the simulation,you can set solvate to True.\nExample1:\n{\'nonbondedMethod\': \'NoCutoff\',\n\'constraints\': \'None\',\n\'rigidWater\': False}\nExample2:\n{\'nonbondedMethod\': \'CutoffPeriodic\',\n\'nonbondedCutoff\': 1.0,\n\'constraints\': \'HBonds\',\n\'rigidWater\': True,\n\'constraintTolerance\': 0.00001,\n\'solvate\': True} """'}), '({\'nonbondedMethod\': \'NoCutoff\', \'nonbondedCutoff\': \'1 * nanometers\',\n    \'ewaldErrorTolerance\': None, \'constraints\': \'None\', \'rigidWater\': False,\n    \'constraintTolerance\': None, \'solvate\': False}, description=\n    """Parameters for the openmm system. For nonbondedMethod, you can choose from the following:\nNoCutoff, CutoffNonPeriodic, CutoffPeriodic, Ewald, PME. If anything but NoCutoff is chosen,you have to include a nonbondedCutoffand a constrainTolerance.\nIf PME is chosen,you have to include an ewaldErrorTolerance too.For constraints, you can choose from the following:\nNone, HBonds, AllBonds or OnlyWater.For rigidWater, you can choose from the following:\nTrue, False.\nFinally, if you want to solvate the system, before the simulation,you can set solvate to True.\nExample1:\n{\'nonbondedMethod\': \'NoCutoff\',\n\'constraints\': \'None\',\n\'rigidWater\': False}\nExample2:\n{\'nonbondedMethod\': \'CutoffPeriodic\',\n\'nonbondedCutoff\': 1.0,\n\'constraints\': \'HBonds\',\n\'rigidWater\': True,\n\'constraintTolerance\': 0.00001,\n\'solvate\': True} """\n    )\n', (20984, 22020), False, 'from pydantic import BaseModel, Field\n'), ((22557, 22791), 'pydantic.Field', 'Field', (["{'integrator_type': 'LangevinMiddle', 'Temperature': '300 * kelvin',\n    'Friction': '1.0 / picoseconds', 'Timestep': '0.002 * picoseconds',\n    'Pressure': '1.0 * bar'}"], {'description': '"""Parameters for the openmm integrator."""'}), "({'integrator_type': 'LangevinMiddle', 'Temperature': '300 * kelvin',\n    'Friction': '1.0 / picoseconds', 'Timestep': '0.002 * picoseconds',\n    'Pressure': '1.0 * bar'}, description=\n    'Parameters for the openmm integrator.')\n", (22562, 22791), False, 'from pydantic import BaseModel, Field\n'), ((22917, 23653), 'pydantic.Field', 'Field', (["{'Ensemble': 'NVT', 'Number of Steps': 5000, 'record_interval_steps': 100,\n    'record_params': ['step', 'potentialEnergy', 'temperature']}"], {'description': '"""Parameters for the openmm simulation.\n        The ensemble can be NPT, NVT or NVE.\n        The number of steps is the number of steps the simulation will run for.\n        record_interval_steps is the number of steps between each record:\n        hould be the number of steps divided by 100.\n        The record_params is a list of parameters that will\n        be recorded during the simulation The options are:\n        [Step,Time,Speed,Progress,RemainingTime,ElapsedTime,\n        PotentialEnergy,KineticEnergy,TotalEnergy,\n        Temperature,Volume,Density]"""'}), '({\'Ensemble\': \'NVT\', \'Number of Steps\': 5000, \'record_interval_steps\':\n    100, \'record_params\': [\'step\', \'potentialEnergy\', \'temperature\']},\n    description=\n    """Parameters for the openmm simulation.\n        The ensemble can be NPT, NVT or NVE.\n        The number of steps is the number of steps the simulation will run for.\n        record_interval_steps is the number of steps between each record:\n        hould be the number of steps divided by 100.\n        The record_params is a list of parameters that will\n        be recorded during the simulation The options are:\n        [Step,Time,Speed,Progress,RemainingTime,ElapsedTime,\n        PotentialEnergy,KineticEnergy,TotalEnergy,\n        Temperature,Volume,Density]"""\n    )\n', (22922, 23653), False, 'from pydantic import BaseModel, Field\n'), ((74259, 74294), 'openmm.app.ForceField', 'ForceField', (['Forcefield', 'Water_model'], {}), '(Forcefield, Water_model)\n', (74269, 74294), False, 'from openmm.app import PME, AllBonds, CutoffNonPeriodic, CutoffPeriodic, DCDReporter, Ewald, ForceField, HAngles, HBonds, Modeller, NoCutoff, PDBFile, PDBReporter, PDBxFile, Simulation, StateDataReporter\n'), ((2848, 3019), 'langchain.chat_models.ChatOpenAI', 'langchain.chat_models.ChatOpenAI', ([], {'temperature': 'self.temperature', 'model_name': 'self.model_name', 'request_timeout': 'self.request_timeout', 'max_tokens': 'self.request_timeout'}), '(temperature=self.temperature, model_name=\n    self.model_name, request_timeout=self.request_timeout, max_tokens=self.\n    request_timeout)\n', (2880, 3019), False, 'import langchain\n'), ((9952, 10019), 'langchain.prompts.PromptTemplate', 'PromptTemplate', ([], {'template': 'prompt_template', 'input_variables': "['query']"}), "(template=prompt_template, input_variables=['query'])\n", (9966, 10019), False, 'from langchain.prompts import PromptTemplate\n'), ((10040, 10077), 'langchain.chains.LLMChain', 'LLMChain', ([], {'prompt': 'prompt', 'llm': 'self.llm'}), '(prompt=prompt, llm=self.llm)\n', (10048, 10077), False, 'from langchain.chains import LLMChain\n'), ((12416, 12477), 'streamlit.markdown', 'st.markdown', (['"""Setting up forcefields"""'], {'unsafe_allow_html': '(True)'}), "('Setting up forcefields', unsafe_allow_html=True)\n", (12427, 12477), True, 'import streamlit as st\n'), ((12732, 12765), 'mdagent.tools.base_tools.preprocess_tools.CleaningTools', 'CleaningTools', (['self.path_registry'], {}), '(self.path_registry)\n', (12745, 12765), False, 'from mdagent.tools.base_tools.preprocess_tools import CleaningTools\n'), ((13056, 13093), 'openmm.app.Modeller', 'Modeller', (['pdb.topology', 'pdb.positions'], {}), '(pdb.topology, pdb.positions)\n', (13064, 13093), False, 'from openmm.app import PME, AllBonds, CutoffNonPeriodic, CutoffPeriodic, DCDReporter, Ewald, ForceField, HAngles, HBonds, Modeller, NoCutoff, PDBFile, PDBReporter, PDBxFile, Simulation, StateDataReporter\n'), ((15170, 15219), 'openmm.app.Simulation', 'Simulation', (['modeller.topology', 'system', 'integrator'], {}), '(modeller.topology, system, integrator)\n', (15180, 15219), False, 'from openmm.app import PME, AllBonds, CutoffNonPeriodic, CutoffPeriodic, DCDReporter, Ewald, ForceField, HAngles, HBonds, Modeller, NoCutoff, PDBFile, PDBReporter, PDBxFile, Simulation, StateDataReporter\n'), ((15873, 15914), 'ast.literal_eval', 'ast.literal_eval', (["params['record_params']"], {}), "(params['record_params'])\n", (15889, 15914), False, 'import ast\n'), ((17472, 17524), 'os.path.exists', 'os.path.exists', (['"""simulation_parameters_summary.json"""'], {}), "('simulation_parameters_summary.json')\n", (17486, 17524), False, 'import os\n'), ((25493, 25547), 'streamlit.markdown', 'st.markdown', (['"""Building system"""'], {'unsafe_allow_html': '(True)'}), "('Building system', unsafe_allow_html=True)\n", (25504, 25547), True, 'import streamlit as st\n'), ((25680, 25702), 'openmm.app.PDBFile', 'PDBFile', (['self.pdb_path'], {}), '(self.pdb_path)\n', (25687, 25702), False, 'from openmm.app import PME, AllBonds, CutoffNonPeriodic, CutoffPeriodic, DCDReporter, Ewald, ForceField, HAngles, HBonds, Modeller, NoCutoff, PDBFile, PDBReporter, PDBxFile, Simulation, StateDataReporter\n'), ((25729, 25770), 'openmm.app.ForceField', 'ForceField', (['*self.params.forcefield_files'], {}), '(*self.params.forcefield_files)\n', (25739, 25770), False, 'from openmm.app import PME, AllBonds, CutoffNonPeriodic, CutoffPeriodic, DCDReporter, Ewald, ForceField, HAngles, HBonds, Modeller, NoCutoff, PDBFile, PDBReporter, PDBxFile, Simulation, StateDataReporter\n'), ((26409, 26469), 'streamlit.markdown', 'st.markdown', (['"""Setting up integrator"""'], {'unsafe_allow_html': '(True)'}), "('Setting up integrator', unsafe_allow_html=True)\n", (26420, 26469), True, 'import streamlit as st\n'), ((27344, 27402), 'streamlit.markdown', 'st.markdown', (['"""Creating simulation"""'], {'unsafe_allow_html': '(True)'}), "('Creating simulation', unsafe_allow_html=True)\n", (27355, 27402), True, 'import streamlit as st\n'), ((32351, 32388), 'openmm.app.Modeller', 'Modeller', (['pdb.topology', 'pdb.positions'], {}), '(pdb.topology, pdb.positions)\n', (32359, 32388), False, 'from openmm.app import PME, AllBonds, CutoffNonPeriodic, CutoffPeriodic, DCDReporter, Ewald, ForceField, HAngles, HBonds, Modeller, NoCutoff, PDBFile, PDBReporter, PDBxFile, Simulation, StateDataReporter\n'), ((43173, 43248), 'streamlit.markdown', 'st.markdown', (['"""Standalone simulation script written"""'], {'unsafe_allow_html': '(True)'}), "('Standalone simulation script written', unsafe_allow_html=True)\n", (43184, 43248), True, 'import streamlit as st\n'), ((43363, 43432), 'streamlit.markdown', 'st.markdown', (['"""Performing energy minimization"""'], {'unsafe_allow_html': '(True)'}), "('Performing energy minimization', unsafe_allow_html=True)\n", (43374, 43432), True, 'import streamlit as st\n'), ((44047, 44125), 'streamlit.markdown', 'st.markdown', (['"""Minimization complete! Equilibrating..."""'], {'unsafe_allow_html': '(True)'}), "('Minimization complete! Equilibrating...', unsafe_allow_html=True)\n", (44058, 44125), True, 'import streamlit as st\n'), ((44439, 44491), 'streamlit.markdown', 'st.markdown', (['"""Simulating..."""'], {'unsafe_allow_html': '(True)'}), "('Simulating...', unsafe_allow_html=True)\n", (44450, 44491), True, 'import streamlit as st\n'), ((44628, 44672), 'streamlit.markdown', 'st.markdown', (['"""Done!"""'], {'unsafe_allow_html': '(True)'}), "('Done!', unsafe_allow_html=True)\n", (44639, 44672), True, 'import streamlit as st\n'), ((73603, 73620), 'openmm.app.PDBFile', 'PDBFile', (['pdb_path'], {}), '(pdb_path)\n', (73610, 73620), False, 'from openmm.app import PME, AllBonds, CutoffNonPeriodic, CutoffPeriodic, DCDReporter, Ewald, ForceField, HAngles, HBonds, Modeller, NoCutoff, PDBFile, PDBReporter, PDBxFile, Simulation, StateDataReporter\n'), ((5414, 5459), 'openmm.LangevinMiddleIntegrator', 'LangevinMiddleIntegrator', ([], {}), '(**integrator_params)\n', (5438, 5459), False, 'from openmm import AndersenThermostat, BrownianIntegrator, LangevinIntegrator, LangevinMiddleIntegrator, MonteCarloBarostat, OpenMMException, Platform, VerletIntegrator, app, unit\n'), ((10642, 10668), 'json.dump', 'json.dump', (['summary_dict', 'f'], {}), '(summary_dict, f)\n', (10651, 10668), False, 'import json\n'), ((11183, 11195), 'json.load', 'json.load', (['f'], {}), '(f)\n', (11192, 11195), False, 'import json\n'), ((12619, 12654), 'openmm.app.ForceField', 'ForceField', (['Forcefield', 'Water_model'], {}), '(Forcefield, Water_model)\n', (12629, 12654), False, 'from openmm.app import PME, AllBonds, CutoffNonPeriodic, CutoffPeriodic, DCDReporter, Ewald, ForceField, HAngles, HBonds, Modeller, NoCutoff, PDBFile, PDBReporter, PDBxFile, Simulation, StateDataReporter\n'), ((12956, 12972), 'openmm.app.PDBFile', 'PDBFile', (['pdbfile'], {}), '(pdbfile)\n', (12963, 12972), False, 'from openmm.app import PME, AllBonds, CutoffNonPeriodic, CutoffPeriodic, DCDReporter, Ewald, ForceField, HAngles, HBonds, Modeller, NoCutoff, PDBFile, PDBReporter, PDBxFile, Simulation, StateDataReporter\n'), ((13846, 13915), 'streamlit.markdown', 'st.markdown', (['"""Setting up Langevin integrator"""'], {'unsafe_allow_html': '(True)'}), "('Setting up Langevin integrator', unsafe_allow_html=True)\n", (13857, 13915), True, 'import streamlit as st\n'), ((15727, 15759), 'openmm.app.PDBReporter', 'PDBReporter', (['f"""{name}.pdb"""', '(1000)'], {}), "(f'{name}.pdb', 1000)\n", (15738, 15759), False, 'from openmm.app import PME, AllBonds, CutoffNonPeriodic, CutoffPeriodic, DCDReporter, Ewald, ForceField, HAngles, HBonds, Modeller, NoCutoff, PDBFile, PDBReporter, PDBxFile, Simulation, StateDataReporter\n'), ((16625, 16680), 'openmm.app.StateDataReporter', 'StateDataReporter', (['f"""{name}.csv"""', '(1000)'], {}), "(f'{name}.csv', 1000, **reporter_args)\n", (16642, 16680), False, 'from openmm.app import PME, AllBonds, CutoffNonPeriodic, CutoffPeriodic, DCDReporter, Ewald, ForceField, HAngles, HBonds, Modeller, NoCutoff, PDBFile, PDBReporter, PDBxFile, Simulation, StateDataReporter\n'), ((17731, 17746), 'os.listdir', 'os.listdir', (['"""."""'], {}), "('.')\n", (17741, 17746), False, 'import os\n'), ((19272, 19284), 'json.load', 'json.load', (['f'], {}), '(f)\n', (19281, 19284), False, 'import json\n'), ((26664, 26767), 'openmm.LangevinMiddleIntegrator', 'LangevinMiddleIntegrator', (["int_params['Temperature']", "int_params['Friction']", "int_params['Timestep']"], {}), "(int_params['Temperature'], int_params['Friction'],\n    int_params['Timestep'])\n", (26688, 26767), False, 'from openmm import AndersenThermostat, BrownianIntegrator, LangevinIntegrator, LangevinMiddleIntegrator, MonteCarloBarostat, OpenMMException, Platform, VerletIntegrator, app, unit\n'), ((27543, 27576), 'openmm.Platform.getPlatformByName', 'Platform.getPlatformByName', (['"""CPU"""'], {}), "('CPU')\n", (27569, 27576), False, 'from openmm import AndersenThermostat, BrownianIntegrator, LangevinIntegrator, LangevinMiddleIntegrator, MonteCarloBarostat, OpenMMException, Platform, VerletIntegrator, app, unit\n'), ((42922, 42947), 'os.path.exists', 'os.path.exists', (['directory'], {}), '(directory)\n', (42936, 42947), False, 'import os\n'), ((42961, 42983), 'os.makedirs', 'os.makedirs', (['directory'], {}), '(directory)\n', (42972, 42983), False, 'import os\n'), ((44714, 44751), 'os.path.exists', 'os.path.exists', (['"""temp_trajectory.dcd"""'], {}), "('temp_trajectory.dcd')\n", (44728, 44751), False, 'import os\n'), ((44817, 44847), 'os.path.exists', 'os.path.exists', (['"""temp_log.txt"""'], {}), "('temp_log.txt')\n", (44831, 44847), False, 'import os\n'), ((44906, 44943), 'os.path.exists', 'os.path.exists', (['"""temp_checkpoint.chk"""'], {}), "('temp_checkpoint.chk')\n", (44920, 44943), False, 'import os\n'), ((47535, 47589), 'streamlit.markdown', 'st.markdown', (['"""simulation set!"""'], {'unsafe_allow_html': '(True)'}), "('simulation set!', unsafe_allow_html=True)\n", (47546, 47589), True, 'import streamlit as st\n'), ((71482, 71548), 'streamlit.markdown', 'st.markdown', (['"""Setting default forcefields"""'], {'unsafe_allow_html': '(True)'}), "('Setting default forcefields', unsafe_allow_html=True)\n", (71493, 71548), True, 'import streamlit as st\n'), ((73658, 73676), 'openmm.app.PDBxFile', 'PDBxFile', (['pdb_path'], {}), '(pdb_path)\n', (73666, 73676), False, 'from openmm.app import PME, AllBonds, CutoffNonPeriodic, CutoffPeriodic, DCDReporter, Ewald, ForceField, HAngles, HBonds, Modeller, NoCutoff, PDBFile, PDBReporter, PDBxFile, Simulation, StateDataReporter\n'), ((5527, 5564), 'openmm.VerletIntegrator', 'VerletIntegrator', ([], {}), '(**integrator_params)\n', (5543, 5564), False, 'from openmm import AndersenThermostat, BrownianIntegrator, LangevinIntegrator, LangevinMiddleIntegrator, MonteCarloBarostat, OpenMMException, Platform, VerletIntegrator, app, unit\n'), ((13018, 13035), 'openmm.app.PDBxFile', 'PDBxFile', (['pdbfile'], {}), '(pdbfile)\n', (13026, 13035), False, 'from openmm.app import PME, AllBonds, CutoffNonPeriodic, CutoffPeriodic, DCDReporter, Ewald, ForceField, HAngles, HBonds, Modeller, NoCutoff, PDBFile, PDBReporter, PDBxFile, Simulation, StateDataReporter\n'), ((15006, 15073), 'streamlit.markdown', 'st.markdown', (['"""Setting up Verlet integrator"""'], {'unsafe_allow_html': '(True)'}), "('Setting up Verlet integrator', unsafe_allow_html=True)\n", (15017, 15073), True, 'import streamlit as st\n'), ((26911, 27008), 'openmm.LangevinIntegrator', 'LangevinIntegrator', (["int_params['Temperature']", "int_params['Friction']", "int_params['Timestep']"], {}), "(int_params['Temperature'], int_params['Friction'],\n    int_params['Timestep'])\n", (26929, 27008), False, 'from openmm import AndersenThermostat, BrownianIntegrator, LangevinIntegrator, LangevinMiddleIntegrator, MonteCarloBarostat, OpenMMException, Platform, VerletIntegrator, app, unit\n'), ((29047, 29122), 'openmm.app.DCDReporter', 'DCDReporter', (['f"""{trajectory_name}"""', "self.sim_params['record_interval_steps']"], {}), "(f'{trajectory_name}', self.sim_params['record_interval_steps'])\n", (29058, 29122), False, 'from openmm.app import PME, AllBonds, CutoffNonPeriodic, CutoffPeriodic, DCDReporter, Ewald, ForceField, HAngles, HBonds, Modeller, NoCutoff, PDBFile, PDBReporter, PDBxFile, Simulation, StateDataReporter\n'), ((29258, 29331), 'openmm.app.PDBReporter', 'PDBReporter', (['f"""{topology_name}"""', "self.sim_params['record_interval_steps']"], {}), "(f'{topology_name}', self.sim_params['record_interval_steps'])\n", (29269, 29331), False, 'from openmm.app import PME, AllBonds, CutoffNonPeriodic, CutoffPeriodic, DCDReporter, Ewald, ForceField, HAngles, HBonds, Modeller, NoCutoff, PDBFile, PDBReporter, PDBxFile, Simulation, StateDataReporter\n'), ((29467, 29612), 'openmm.app.StateDataReporter', 'StateDataReporter', (['f"""{log_name}"""', "self.sim_params['record_interval_steps']"], {'step': '(True)', 'potentialEnergy': '(True)', 'temperature': '(True)', 'separator': '"""\t"""'}), "(f'{log_name}', self.sim_params['record_interval_steps'],\n    step=True, potentialEnergy=True, temperature=True, separator='\\t')\n", (29484, 29612), False, 'from openmm.app import PME, AllBonds, CutoffNonPeriodic, CutoffPeriodic, DCDReporter, Ewald, ForceField, HAngles, HBonds, Modeller, NoCutoff, PDBFile, PDBReporter, PDBxFile, Simulation, StateDataReporter\n'), ((30145, 30221), 'openmm.app.DCDReporter', 'DCDReporter', (['"""temp_trajectory.dcd"""', "self.sim_params['record_interval_steps']"], {}), "('temp_trajectory.dcd', self.sim_params['record_interval_steps'])\n", (30156, 30221), False, 'from openmm.app import PME, AllBonds, CutoffNonPeriodic, CutoffPeriodic, DCDReporter, Ewald, ForceField, HAngles, HBonds, Modeller, NoCutoff, PDBFile, PDBReporter, PDBxFile, Simulation, StateDataReporter\n'), ((30357, 30431), 'openmm.app.PDBReporter', 'PDBReporter', (['"""temp_topology.pdb"""', "self.sim_params['record_interval_steps']"], {}), "('temp_topology.pdb', self.sim_params['record_interval_steps'])\n", (30368, 30431), False, 'from openmm.app import PME, AllBonds, CutoffNonPeriodic, CutoffPeriodic, DCDReporter, Ewald, ForceField, HAngles, HBonds, Modeller, NoCutoff, PDBFile, PDBReporter, PDBxFile, Simulation, StateDataReporter\n'), ((30567, 30713), 'openmm.app.StateDataReporter', 'StateDataReporter', (['"""temp_log.txt"""', "self.sim_params['record_interval_steps']"], {'step': '(True)', 'potentialEnergy': '(True)', 'temperature': '(True)', 'separator': '"""\t"""'}), "('temp_log.txt', self.sim_params['record_interval_steps'],\n    step=True, potentialEnergy=True, temperature=True, separator='\\t')\n", (30584, 30713), False, 'from openmm.app import PME, AllBonds, CutoffNonPeriodic, CutoffPeriodic, DCDReporter, Ewald, ForceField, HAngles, HBonds, Modeller, NoCutoff, PDBFile, PDBReporter, PDBxFile, Simulation, StateDataReporter\n'), ((33597, 33626), 'openmm.unit.value_in_unit', 'unit.value_in_unit', (['unit.unit'], {}), '(unit.unit)\n', (33615, 33626), False, 'from openmm import AndersenThermostat, BrownianIntegrator, LangevinIntegrator, LangevinMiddleIntegrator, MonteCarloBarostat, OpenMMException, Platform, VerletIntegrator, app, unit\n'), ((33629, 33649), 'openmm.unit.unit.get_name', 'unit.unit.get_name', ([], {}), '()\n', (33647, 33649), False, 'from openmm import AndersenThermostat, BrownianIntegrator, LangevinIntegrator, LangevinMiddleIntegrator, MonteCarloBarostat, OpenMMException, Platform, VerletIntegrator, app, unit\n'), ((42802, 42833), 'textwrap.dedent', 'textwrap.dedent', (['script_content'], {}), '(script_content)\n', (42817, 42833), False, 'import textwrap\n'), ((44769, 44801), 'os.remove', 'os.remove', (['"""temp_trajectory.dcd"""'], {}), "('temp_trajectory.dcd')\n", (44778, 44801), False, 'import os\n'), ((44865, 44890), 'os.remove', 'os.remove', (['"""temp_log.txt"""'], {}), "('temp_log.txt')\n", (44874, 44890), False, 'import os\n'), ((44961, 44993), 'os.remove', 'os.remove', (['"""temp_checkpoint.chk"""'], {}), "('temp_checkpoint.chk')\n", (44970, 44993), False, 'import os\n'), ((71728, 71794), 'streamlit.markdown', 'st.markdown', (['"""Setting default forcefields"""'], {'unsafe_allow_html': '(True)'}), "('Setting default forcefields', unsafe_allow_html=True)\n", (71739, 71794), True, 'import streamlit as st\n'), ((5634, 5673), 'openmm.BrownianIntegrator', 'BrownianIntegrator', ([], {}), '(**integrator_params)\n', (5652, 5673), False, 'from openmm import AndersenThermostat, BrownianIntegrator, LangevinIntegrator, LangevinMiddleIntegrator, MonteCarloBarostat, OpenMMException, Platform, VerletIntegrator, app, unit\n'), ((14061, 14112), 'openmm.MonteCarloBarostat', 'MonteCarloBarostat', (['(_pressure * bar)', '(_temp * kelvin)'], {}), '(_pressure * bar, _temp * kelvin)\n', (14079, 14112), False, 'from openmm import AndersenThermostat, BrownianIntegrator, LangevinIntegrator, LangevinMiddleIntegrator, MonteCarloBarostat, OpenMMException, Platform, VerletIntegrator, app, unit\n'), ((49165, 49180), 'os.listdir', 'os.listdir', (['"""."""'], {}), "('.')\n", (49175, 49180), False, 'import os\n'), ((49205, 49236), 'os.path.exists', 'os.path.exists', (['"""files/records"""'], {}), "('files/records')\n", (49219, 49236), False, 'import os\n'), ((49258, 49286), 'os.makedirs', 'os.makedirs', (['"""files/records"""'], {}), "('files/records')\n", (49269, 49286), False, 'import os\n'), ((14496, 14546), 'openmm.AndersenThermostat', 'AndersenThermostat', (['(_temp * kelvin)', '(1 / picosecond)'], {}), '(_temp * kelvin, 1 / picosecond)\n', (14514, 14546), False, 'from openmm import AndersenThermostat, BrownianIntegrator, LangevinIntegrator, LangevinMiddleIntegrator, MonteCarloBarostat, OpenMMException, Platform, VerletIntegrator, app, unit\n'), ((14580, 14631), 'openmm.MonteCarloBarostat', 'MonteCarloBarostat', (['(_pressure * bar)', '(_temp * kelvin)'], {}), '(_pressure * bar, _temp * kelvin)\n', (14598, 14631), False, 'from openmm import AndersenThermostat, BrownianIntegrator, LangevinIntegrator, LangevinMiddleIntegrator, MonteCarloBarostat, OpenMMException, Platform, VerletIntegrator, app, unit\n'), ((52042, 52096), 're.match', 're.match', (['"""([+-]?[0-9]*\\\\.?[0-9]+)([a-zA-Z]*)"""', 'cutoff'], {}), "('([+-]?[0-9]*\\\\.?[0-9]+)([a-zA-Z]*)', cutoff)\n", (52050, 52096), False, 'import re\n'), ((54705, 54763), 're.match', 're.match', (['"""^(\\\\d+(?:\\\\.\\\\d+)?)([a-zA-Z]+)$"""', 'parameter_str'], {}), "('^(\\\\d+(?:\\\\.\\\\d+)?)([a-zA-Z]+)$', parameter_str)\n", (54713, 54763), False, 'import re\n'), ((55149, 55210), 're.match', 're.match', (['"""([+-]?[0-9]*\\\\.?[0-9]+)([a-zA-Z]*)"""', 'parameter_str'], {}), "('([+-]?[0-9]*\\\\.?[0-9]+)([a-zA-Z]*)', parameter_str)\n", (55157, 55210), False, 'import re\n')]