# Module to run a normal simulation

import numpy as np
import metalwalls

# initialise parallel environment in Python with MPI4PY
def initialize_parallel():
   me = 0
   from mpi4py import MPI
   comm = MPI.COMM_WORLD
   me = comm.Get_rank()
   nprocs = comm.Get_size()
   my_parallel = metalwalls.mw_parallel.MW_parallel_t()
   return my_parallel, comm


# setup metalwalls simulation
# instantiate system, parallel, algorithm
# read inputs
# setup outputs, timers, precalculate Ewald terms...
def setup_simulation(start_parallel=True, my_parallel=None):
   my_system = metalwalls.mw_system.MW_system_t()
   my_algorithms = metalwalls.mw_algorithms.MW_algorithms_t()

   if start_parallel:
      my_parallel, comm = initialize_parallel()
      fcomm = comm.py2f()
      metalwalls.mw_tools.start_parallel_python(my_parallel, fcomm)

   if not(start_parallel) and my_parallel==None:
      print("No parallel type given in setup_simulation, quitting")
      quit()

   do_output = metalwalls.mw_tools.initialize(my_system, my_parallel, my_algorithms)

   step_output_frequency = metalwalls.mw_tools.setup(my_system, my_parallel, my_algorithms, do_output)
   return my_system, my_parallel, my_algorithms, do_output, step_output_frequency


# run *num_steps* of simulation
def run_simulation(num_steps=-1):
   my_system, my_parallel, my_algorithms, do_output, step_output_frequency = setup_simulation()
   # Read optional parameter or the number of steps provided in the runtime
   if num_steps == -1:
      num_steps = my_system.num_steps
   metalwalls.mw_tools.run_step(my_system, my_parallel, my_algorithms, 0, num_steps, do_output, step_output_frequency, True, True)


# finalize the simulation
def finalize_simulation(my_system, my_parallel, my_algorithms):
   metalwalls.mw_tools.finalize(my_system, my_parallel, my_algorithms)
   metalwalls.mw_tools.finalize_parallel()
