import numpy as np
import metalwalls
import simulation

# run a simulation at fixed pressure treating the electrodes as pistons
# python interface necessary if the electrodes consist of several types
def run_piston(first_electrode_list, second_electrode_list, num_steps=-1):
   # Setup the simulation
   my_system, my_parallel, my_algorithms, do_output, step_output_frequency = simulation.setup_simulation()

   # Read optional parameter or the number of steps provided in the runtime
   if num_steps == -1:
      num_steps = my_system.num_steps
   else:
      my_system.num_steps = num_steps
   
   # Sum electrode forces in precalculation step
   metalwalls.mw_tools.sum_electrode_forces(my_system, np.array(first_electrode_list))
   metalwalls.mw_tools.sum_electrode_forces(my_system, np.array(second_electrode_list))
  
   # Setup the timers 
   metalwalls.mw_tools.run_step(my_system, my_parallel, my_algorithms, 0, 0, do_output, step_output_frequency, True, False)
   
   # Set the electrode mass to the sum of the individual electrode masses
   tot_mass = 0
   for electrode in first_electrode_list:
      tot_mass += metalwalls.mw_tools.get_electrode_mass(my_system, electrode)
   for electrode in first_electrode_list:
      metalwalls.mw_tools.set_electrode_mass(my_system, electrode, tot_mass)

   tot_mass = 0
   for electrode in second_electrode_list:
      tot_mass += metalwalls.mw_tools.get_electrode_mass(my_system, electrode)
   for electrode in second_electrode_list:
      metalwalls.mw_tools.set_electrode_mass(my_system, electrode, tot_mass)

   # Run steps
   for step in range(1, num_steps+1):
      # First step of MD step: propagate positions, compute charges...
      metalwalls.mw_tools.step_setup(my_system, my_algorithms, my_parallel, step)
   
      # Second step of MD step: compute forces
      metalwalls.mw_tools.step_compute_forces(my_system, my_parallel, step)
   
      # Sum electrode forces
      metalwalls.mw_tools.sum_electrode_forces(my_system, np.array(first_electrode_list))
      metalwalls.mw_tools.sum_electrode_forces(my_system, np.array(second_electrode_list))
   
      # Third step of MD step: update velocities, output info
      metalwalls.mw_tools.step_update_output(my_system, my_algorithms, step, do_output, step_output_frequency)
   
   # Stop timers, output diagnostics
   metalwalls.mw_tools.run_step(my_system, my_parallel, my_algorithms, num_steps, 0, do_output, step_output_frequency, False, True)
   
   simulation.finalize_simulation(my_system, my_parallel, my_algorithms)
