module MW_tools

implicit none
   private

   ! public subroutines
   public :: start_parallel
   public :: start_parallel_python
   public :: parse_command_line
   public :: initialize_output
   public :: read_input_files
   public :: initialize
   public :: setup
   public :: run_step
   public :: step_setup
   public :: step_compute_forces
   public :: step_update_output
   public :: finalize
   public :: finalize_parallel
   public :: compute_lj_energy
   public :: compute_ft_energy
   public :: compute_xft_energy
   public :: modify_charge
   public :: get_electrode_mass
   public :: set_electrode_mass
   public :: sum_electrode_forces

contains

   !================================================================================
   subroutine initialize(system, parallel, algorithms, do_output)
      use MW_parallel, only: MW_parallel_t
      use MW_timers
      use MW_stdio, only: MW_stdout
      use MW_command, only: MW_command_t, ARGUMENT_VALUE_LEN
      use MW_algorithms, only: MW_algorithms_t, &
            MW_algorithms_print_type => print_type
      use MW_system, only: MW_system_t, &
            MW_system_print_type => print_type, &
            MW_system_output_setup => output_setup, &
            MW_system_distribute_work => distribute_work, &
            MW_system_setup_sr_interaction => setup_sr_interaction, &
            MW_system_setup_num_dof => setup_num_dof, &
            MW_system_setup_velocities => setup_velocities, &
            MW_system_setup_polarization => setup_polarization
      use MW_four_site_model, only: &
            MW_four_site_model_compute_fourth_site_position => compute_fourth_site_position

      implicit none

      type(MW_parallel_t), intent(inout)   :: parallel
      type(MW_system_t), intent(inout)     :: system
      type(MW_algorithms_t), intent(inout) :: algorithms
      logical, intent(out)                 :: do_output

      ! Local
      type(MW_command_t)            :: parser
      character(ARGUMENT_VALUE_LEN) :: configfile
      character(ARGUMENT_VALUE_LEN) :: datafile
      integer                       :: funit
      logical                       :: output_rank

!      call start_parallel(parallel)
      call MW_timers_initialize()
      call MW_timers_start(TIMER_TOTAL)
      call MW_timers_start(TIMER_INITIALIZATION)
      call parse_command_line(parser, configfile, datafile)
      call initialize_output(funit, parser, parallel, do_output, output_rank)
      call read_input_files(system, algorithms, configfile, datafile)
      call MW_system_distribute_work(system, parallel) ! Distribute computation across parallel processes
      call MW_four_site_model_compute_fourth_site_position(system%num_pbc, &
            system%molecules, system%xyz_ions(:,:,system%xyz_ions_step(1)), &
            system%box, system%ions)
      call MW_system_setup_sr_interaction(system)
      call MW_system_setup_num_dof(system)
      call MW_system_setup_velocities(system)
      call MW_system_setup_polarization(system)


      ! Print system parameters
      if (do_output) then
         call MW_system_print_type(system, MW_stdout)
         call MW_algorithms_print_type(algorithms, MW_stdout, system)
         flush(MW_stdout)
      end if

      ! Setup output files
      if (do_output) then
         if (output_rank) then
            call MW_system_output_setup(system, parallel%comm_rank)
         else
            call MW_system_output_setup(system)
         end if
      end if

      call MW_timers_stop(TIMER_INITIALIZATION)

      ! Initialize data on GPU
      !$acc enter data &
      !$acc copyin(system, system%q_atoms(:,:), system%q_atoms_step(:), system%electrodes(:), system%ions(:), &
      !$acc        system%xyz_atoms(:,:), system%xyz_ions(:,:,:), system%dipoles_ions(:,:,:), &
      !$acc        system%deltas_ions(:,:), system%epsilons_ions(:,:,:), &
      !$acc        system%ewald, system%ewald%zpoint(:), system%ewald%zweight(:), &
      !$acc        system%localwork, system%localwork%molecules_istart(:), system%localwork%molecules_iend(:), &
      !$acc        system%localwork%count_atoms(:), system%localwork%offset_atoms(:), &
      !$acc        algorithms, algorithms%cg, algorithms%matrix_inversion, &
      !$acc        system%thermostat, system%thermostat%v_zeta(:,:)) &
      !$acc create(system%ewald%Sk_cos(:), system%ewald%Sk_sin(:), &
      !$acc        system%ewald%cos_kx_ions(:,0:), system%ewald%sin_kx_ions(:,0:), &
      !$acc        system%ewald%cos_ky_ions(:,0:), system%ewald%sin_ky_ions(:,0:), &
      !$acc        system%ewald%cos_kz_ions(:,0:), system%ewald%sin_kz_ions(:,0:), &
      !$acc        system%ewald%cos_kx_elec(:,0:), system%ewald%sin_kx_elec(:,0:), &
      !$acc        system%ewald%cos_ky_elec(:,0:), system%ewald%sin_ky_elec(:,0:), &
      !$acc        system%ewald%cos_kz_elec(:,0:), system%ewald%sin_kz_elec(:,0:), &
      !$acc        system%coulomb_potential_atoms(:,:), system%forces_ions(:,:,:), &
      !$acc        system%coulomb_forces_ions(:,:,:), system%intermolecular_forces_ions(:,:,:), &
      !$acc        system%intramolecular_forces_ions(:,:,:), system%lj_forces_ions(:,:,:), &
      !$acc        system%ft_forces_ions(:,:,:), system%xft_forces_ions(:,:,:), &
      !$acc        system%daim_forces_ions(:,:,:), system%steele_forces_ions(:), &
      !$acc        system%external_field_forces_ions(:,:,:),  &
      !$acc        system%stress_tensor(:,:), &
      !$acc        system%coulomb_stress_tensor(:,:,:), system%lj_stress_tensor(:,:,:), &      
      !$acc        system%ft_stress_tensor(:,:,:), system%xft_stress_tensor(:,:,:), &      
      !$acc        system%daim_stress_tensor(:,:,:), system%intramolecular_stress_tensor(:,:,:), &      
      !$acc        system%steele_stress_tensor(:,:), &      
      !$acc        algorithms%cg%b(:), algorithms%cg%res(:), algorithms%cg%p(:), &
      !$acc        algorithms%cg%Ad(:), &
      !$acc        algorithms%matrix_inversion%b(:),algorithms%matrix_inversion%A(:,:), &
      !$acc        algorithms%matrix_inversion%Ainv(:,:),algorithms%matrix_inversion%Ax(:), &
      !$acc        algorithms%matrix_inversion%q0(:))
! Manque d et pour cg ?     

   end subroutine initialize

   !================================================================================
   ! Initialization
   ! --------------
   subroutine start_parallel_python(parallel, comm)
      use MW_parallel, only: MW_parallel_t, MW_COMM_WORLD, &
            MW_parallel_initialize_python => initialize_python, &
            MW_parallel_define_type => define_type

      implicit none

      type(MW_parallel_t), intent(inout) :: parallel
      integer, intent(in) :: comm

      call MW_parallel_initialize_python(comm)
      call MW_parallel_define_type(parallel, MW_COMM_WORLD)

   end subroutine start_parallel_python
   !================================================================================
   ! Initialization
   ! --------------
   subroutine start_parallel(parallel)
      use MW_parallel, only: MW_parallel_t, MW_COMM_WORLD, &
            MW_parallel_initialize => initialize, &
            MW_parallel_define_type => define_type

      implicit none

      type(MW_parallel_t), intent(inout) :: parallel

      call MW_parallel_initialize()
      call MW_parallel_define_type(parallel, MW_COMM_WORLD)

   end subroutine start_parallel

   !================================================================================
   ! Parse command line
   ! ------------------
   subroutine parse_command_line(parser, configfile, datafile)
      use MW_command, only: MW_command_t, ARGUMENT_VALUE_LEN, &
            MW_command_setup_parser => setup_parser, &
            MW_command_parse => parse, &
            MW_command_get_option => get_option
      use MW_parallel, only: MW_parallel_finalize => finalize
      use MW_stdio, only: MW_stdout
      use MW_version

      type(MW_command_t), intent(inout) :: parser
      character(ARGUMENT_VALUE_LEN), intent(out) :: configfile
      character(ARGUMENT_VALUE_LEN), intent(out) :: datafile

      ! Local
      logical :: print_help, print_version

      call MW_command_setup_parser(parser)
      call MW_command_parse(parser)

      call MW_command_get_option(parser, "help", print_help)
      if (print_help) then
         call parser%show_help(MW_stdout)
         call MW_parallel_finalize()
         stop
      end if

      call MW_command_get_option(parser, "version", print_version)
      if (print_version) then
         call MW_version_print(MW_stdout)
         call MW_parallel_finalize()
         stop
      end if

      call MW_command_get_option(parser, "input", configfile)
      call MW_command_get_option(parser, "data", datafile)

   end subroutine parse_command_line

   !================================================================================
   subroutine initialize_output(funit, parser, parallel, do_output, output_rank)
      use MW_command, only: MW_command_t, ARGUMENT_VALUE_LEN, &
            MW_command_get_option => get_option
      use MW_parallel, only: MW_parallel_t, &
            MW_parallel_print_type => print_type
      use MW_fileunit, only: &
            MW_fileunit_get_new_unit => get_new_unit
      use MW_stdio, only: MW_stdout, &
            MW_stdio_set_stdout => set_stdout
      use MW_errors, only: &
            MW_errors_open_error => open_error
      use MW_version
      use MW_banner

      implicit none

      type(MW_parallel_t), intent(inout) :: parallel
      type(MW_command_t), intent(inout) :: parser
      integer, intent(out) :: funit
      logical, intent(out) :: do_output, output_rank

      ! Local
      character(ARGUMENT_VALUE_LEN) :: outputfile
      character(ARGUMENT_VALUE_LEN) :: outputrank
      integer :: irank_start, icomma, rank_id
      integer :: ierr

      ! Setup do_output flag for ranks which will output something
      do_output = .false.
      output_rank = .false.
      outputfile = ""
      call MW_command_get_option(parser, "output-rank", outputrank)
      select case(outputrank)
      case("root")
         ! Only rank 0 does output
         if (parallel%comm_rank == 0) do_output = .true.
         output_rank = .false.
      case("all")
         ! All ranks do output
         do_output = .true.
         output_rank = .true.
      case default
         ! Comma separated list of ranks
         irank_start = 1
         icomma = index(outputrank, ",")
         do while (icomma > 0)
            read(outputrank(irank_start:irank_start+icomma-2),*) rank_id
            if (parallel%comm_rank == rank_id) do_output = .true.
            irank_start = irank_start + icomma
            icomma = index(outputrank(irank_start:), ",")
         end do
         if (len_trim(outputrank(irank_start:)) > 0) then
            read(outputrank(irank_start:),*) rank_id
            if (parallel%comm_rank == rank_id) do_output = .true.
         end if
         output_rank = .true.
      end select

      if (do_output) then
         if (output_rank) then
            write(outputfile, '("run_",i6.6,".out")') parallel%comm_rank
         else
            outputfile = "run.out"
         end if
         call MW_fileunit_get_new_unit(funit)
         open(unit=funit, file=outputfile, access="SEQUENTIAL", action="WRITE", &
               position="REWIND", form="FORMATTED", status="REPLACE", &
               iostat=ierr)
         if (ierr /= 0) then
            call MW_errors_open_error("main", "main.f90", outputfile, ierr)
         end if
         call MW_stdio_set_stdout(funit)
      end if

      if (do_output) then
         call MW_banner_print(MW_stdout)
         call MW_version_print(MW_stdout)
         call MW_parallel_print_type(parallel, MW_stdout)
         flush(MW_stdout)
      end if

   end subroutine initialize_output

   !================================================================================
   subroutine read_input_files(system, algorithms, configfile, datafile)
      use MW_command, only: MW_command_t, ARGUMENT_VALUE_LEN
      use MW_algorithms, only: MW_algorithms_t
      use MW_configuration, only: &
            MW_configuration_read_system_parameters => read_system_parameters
      use MW_system, only: MW_system_t, &
            MW_system_read_data => read_data

      type(MW_system_t), intent(inout) :: system
      type(MW_algorithms_t), intent(inout) :: algorithms
      character(ARGUMENT_VALUE_LEN), intent(in) :: configfile
      character(ARGUMENT_VALUE_LEN), intent(in) :: datafile

      integer :: i, itype

      call MW_configuration_read_system_parameters(system, algorithms, configfile)
      call MW_system_read_data(system, datafile)

      if (algorithms%constant_charge) then
         do itype = 1, system%num_elec
            !$acc loop
            do i = system%electrodes(itype)%offset+1, system%electrodes(itype)%offset + system%electrodes(itype)%count
               system%q_atoms(i, system%q_atoms_step(2)) = system%q_atoms(i, system%q_atoms_step(1))
            end do
            !$acc end loop
         end do
      end if

   end subroutine read_input_files

   !================================================================================
   subroutine setup(system, parallel, algorithms, do_output, step_output_frequency)
      use MW_timers
      use MW_algorithms, only: MW_algorithms_t, &
            MW_algorithms_print_statistics => print_statistics
      use MW_parallel, only: MW_parallel_t
      use MW_stdio, only: MW_stdout
      use MW_system, only: MW_system_t, &
            MW_system_print_step => print_step, &
            MW_system_output_dump => output_dump, &
            MW_system_ewald_setup_cossin_ions => ewald_setup_cossin_ions, &
            MW_system_ewald_setup_cossin_elec => ewald_setup_cossin_elec, &
            MW_system_coulomb_melt_efg => coulomb_melt_efg, &
            MW_system_setup_thermostats_barostats_constraints => setup_thermostats_barostats_constraints, &
            MW_system_ions_forces_compute => ions_forces_compute
      use MW_maze, only: &
            MW_maze_setup_history => setup_history, &
            MW_maze_solve_invert => solve_invert
      use MW_electrode_charge, only: &
            MW_electrode_charge_compute_jacobi_precond => compute_jacobi_precond, &
            MW_electrode_charge_compute => compute
      use MW_dipoles_and_electrodes, only: &
            MW_dipoles_and_electrodes_compute_jacobi_precond => compute_jacobi_precond, &
            MW_dipoles_and_electrodes_compute => compute
      use MW_ions_dipoles, only: &
         MW_ions_dipoles_compute_jacobi_precond => compute_jacobi_precond, &
         MW_ions_dipoles_compute => compute
      use MW_ions_radius, only: &
            MW_ions_radius_compute => compute

#ifdef MW_USE_PLUMED
      use MW_plumed, only : MW_plumed_init => init
#endif

      type(MW_system_t), intent(inout) :: system
      type(MW_parallel_t), intent(inout) :: parallel
      type(MW_algorithms_t), intent(inout) :: algorithms
      integer, intent(out) :: step_output_frequency
      logical, intent(in) :: do_output

      integer :: istep=0

      call MW_timers_start(TIMER_SETUP)
#ifdef MW_USE_PLUMED
      if (system%plumedlog) then
         call MW_plumed_init(system%temperature, system%plumed_file, system%num_ions, &
               system%num_atoms, system%num_ion_types, system%dt, system%ions)
     end if
#endif
      ! pre-compute ewald cos/sin factors
      call MW_system_ewald_setup_cossin_ions(system)
      call MW_system_ewald_setup_cossin_elec(system, istep)

      if (system%pimrun.and.system%elecrun) then ! compute initial dipoles and wall charges
         if (algorithms%maze%calc_dip.and.algorithms%maze%calc_elec) then
            call MW_maze_setup_history(algorithms%maze)
         else
            if (algorithms%cg%preconditioned) then
               call MW_dipoles_and_electrodes_compute_jacobi_precond(system, algorithms)
            end if
            call MW_dipoles_and_electrodes_compute(system, algorithms)
         end if
      else if (system%pimrun) then ! compute initial dipoles

         call MW_timers_start(TIMER_MINMU)
         if (algorithms%maze%calc_dip) then
            call MW_maze_setup_history(algorithms%maze)
            call MW_ions_dipoles_compute(system, algorithms)
         else if (algorithms%matrix_inversion%calc_dip) then
            call MW_ions_dipoles_compute(system, algorithms)
         else
            if (algorithms%cg%preconditioned) then
               call MW_ions_dipoles_compute_jacobi_precond(system, algorithms)
            end if
            call MW_ions_dipoles_compute(system, algorithms)
         end if
         call MW_timers_stop(TIMER_MINMU)

         ! compute initial radius
         if (system%aimrun) then
            call MW_timers_start(TIMER_MINAIM)
            call MW_ions_radius_compute(system, algorithms)
            call MW_timers_stop(TIMER_MINAIM)
         end if
      else if (system%elecrun) then ! compute initial wall charges
         call MW_timers_start(TIMER_MINQ)
         if (algorithms%maze%calc_elec) then
            if (.not.algorithms%maze%CG_nostore) call MW_maze_setup_history(algorithms%maze) !Only sets up history for gradient of constraints
            call MW_electrode_charge_compute(system, algorithms)
         else if (algorithms%matrix_inversion%calc_elec) then
            call MW_electrode_charge_compute(system, algorithms)
         else
            if (algorithms%cg%preconditioned) then
               call MW_electrode_charge_compute_jacobi_precond(system, algorithms)
            end if
            call MW_electrode_charge_compute(system, algorithms)
         end if
         call MW_timers_stop(TIMER_MINQ)
      end if

      ! compute efg
      if (system%efgrun) then
         call MW_system_coulomb_melt_efg(system, system%step)
      end if

      ! compute forces
      call MW_system_ions_forces_compute(system, istep)

      call MW_system_setup_thermostats_barostats_constraints(system)

      call MW_system_print_step(system, istep, MW_stdout, do_output)
      call MW_algorithms_print_statistics(algorithms, MW_stdout, do_output, system)
      call MW_system_output_dump(system, istep, do_output)
      !if (do_output) flush(MW_stdout)

      step_output_frequency = system%output%step_dump_frequency
      call MW_timers_stop(TIMER_SETUP)

   end subroutine setup

   !================================================================================

   subroutine run_step(system, parallel, algorithms, istep, num_steps, do_output, step_output_frequency, pre, post)
      use MW_timers
      use MW_stdio, only: MW_stdout
      use MW_algorithms, only: MW_algorithms_t
      use MW_parallel, only: MW_parallel_t
      use MW_system, only: MW_system_t

      type(MW_system_t), intent(inout) :: system
      type(MW_parallel_t), intent(inout) :: parallel
      type(MW_algorithms_t), intent(inout) :: algorithms
      integer, intent(in) :: step_output_frequency
      logical, intent(in) :: do_output
      integer, intent(in) :: num_steps
      logical, intent(in) :: pre, post
      integer, intent(inout) :: istep

      integer :: step

      if (pre) then
         call MW_timers_start(TIMER_MAIN_LOOP)
      end if

      do step = istep+1, istep+num_steps
         call step_setup(system, algorithms, parallel, step)
         call step_compute_forces(system, parallel, step)
         call step_update_output(system, algorithms, step, do_output, step_output_frequency)
      end do

      if (post) then
         call MW_timers_stop(TIMER_MAIN_LOOP)
         call MW_timers_stop(TIMER_TOTAL)

         if (do_output) then
            call MW_timers_print(MW_stdout)
            !flush(MW_stdout)
         end if
      end if

   end subroutine run_step

   !================================================================================

   subroutine step_setup(system, algorithms, parallel, step)
      use MW_timers
      use MW_algorithms, only: MW_algorithms_t, &
            MW_algorithms_print_statistics => print_statistics
      use MW_system, only: MW_system_t, &
            MW_system_ewald_setup_cossin_ions => ewald_setup_cossin_ions, &
            MW_system_ewald_setup_cossin_elec => ewald_setup_cossin_elec, &
            MW_system_coulomb_melt_efg => coulomb_melt_efg, &
            MW_system_update_history => update_history, &
            MW_system_ions_positions_update => ions_positions_update, &
            MW_system_piston_positions_update => piston_positions_update
      use MW_dipoles_and_electrodes, only: &
            MW_dipoles_and_electrodes_compute => compute
      use MW_parallel, only: MW_parallel_t
      use MW_electrode_charge, only: &
            MW_electrode_charge_compute => compute
      use MW_ions_dipoles, only: &
            MW_ions_dipoles_compute => compute
      use MW_ions_radius, only: &
            MW_ions_radius_compute => compute

      type(MW_system_t), intent(inout) :: system
      type(MW_algorithms_t), intent(inout) :: algorithms
      type(MW_parallel_t), intent(inout) :: parallel
      integer, intent(in) :: step

      call MW_timers_start(TIMER_STEP)
      call MW_system_update_history(system)
      call MW_system_ions_positions_update(system,parallel)
      call MW_system_piston_positions_update(system)
      call MW_system_ewald_setup_cossin_ions(system)
      if (system%pimrun.and.system%elecrun) then
         call MW_dipoles_and_electrodes_compute(system, algorithms)
      else if (system%elecrun) then
         call MW_timers_start(TIMER_MINQ)
         call MW_system_ewald_setup_cossin_elec(system, step)
         call MW_electrode_charge_compute(system, algorithms)
         call MW_timers_stop(TIMER_MINQ)
      else if (system%pimrun) then
         call MW_timers_start(TIMER_MINMU)
         call MW_ions_dipoles_compute(system, algorithms)
         call MW_timers_stop(TIMER_MINMU)
         if (system%aimrun) then
            call MW_timers_start(TIMER_MINAIM)
            call MW_ions_radius_compute(system, algorithms)
            call MW_timers_stop(TIMER_MINAIM)
         end if
      end if
      if (system%efgrun) then
         call MW_system_coulomb_melt_efg(system, step)
      end if


   end subroutine step_setup

   !================================================================================

   subroutine step_compute_forces(system, parallel, step)
      use MW_parallel, only: MW_parallel_t
      use MW_system, only: MW_system_t, &
            MW_system_ions_forces_compute => ions_forces_compute

      type(MW_system_t), intent(inout) :: system
      type(MW_parallel_t), intent(inout) :: parallel
      integer, intent(in) :: step

      call MW_system_ions_forces_compute(system, step)

   end subroutine step_compute_forces

   !================================================================================

   subroutine step_update_output(system, algorithms, step, do_output, step_output_frequency)
      use MW_timers
      use MW_stdio, only: MW_stdout
      use MW_algorithms, only: MW_algorithms_t, &
            MW_algorithms_print_statistics => print_statistics
      use MW_system, only: MW_system_t, &
            MW_system_print_step => print_step, &
            MW_system_ions_velocities_update => ions_velocities_update, &
            MW_system_piston_velocities_update => piston_velocities_update, &
            MW_system_output_dump => output_dump, &
            MW_system_compute_diagnostics=> compute_diagnostics

      type(MW_system_t), intent(inout) :: system
      type(MW_algorithms_t), intent(inout) :: algorithms
      integer, intent(in) :: step_output_frequency
      logical, intent(in) :: do_output
      integer, intent(in) :: step

      call MW_system_ions_velocities_update(system, step)
      call MW_system_piston_velocities_update(system)
      call MW_timers_stop(TIMER_STEP)

      if ((step_output_frequency > 0) .and. (mod(step, step_output_frequency) == 0)) then
         call MW_system_print_step(system, step, MW_stdout, do_output)
         if ((system%elecrun).or.system%pimrun.or.system%aimrun) then
            call MW_algorithms_print_statistics(algorithms, MW_stdout, do_output, system)
         end if
      end if

      call MW_system_compute_diagnostics(system)
      call MW_system_output_dump(system, step, do_output)
      !if (do_output) flush(MW_stdout)

   end subroutine step_update_output

   !================================================================================

   subroutine finalize(system, parallel, algorithms)
      use MW_algorithms, only: MW_algorithms_t, &
            MW_algorithms_void_type => void_type
      use MW_parallel, only: MW_parallel_t, &
            MW_parallel_finalize => finalize
      use MW_system, only: MW_system_t, &
            MW_system_deallocate_data_arrays => deallocate_data_arrays, &
            MW_system_void_parameters => void_parameters
      use MW_errors, only: MW_errors_close_error => close_error
      use MW_stdio, only: MW_stdout
#ifdef MW_USE_PLUMED
      use MW_plumed, only : MW_plumed_terminate => terminate
#endif

      type(MW_system_t), intent(inout) :: system
      type(MW_parallel_t), intent(inout) :: parallel
      type(MW_algorithms_t), intent(inout) :: algorithms

      integer :: ierr
      logical :: lopen

#ifdef MW_USE_PLUMED
      if (system%plumedlog) then
         call MW_plumed_terminate()
      end if
#endif
      call MW_algorithms_void_type(algorithms)
      call MW_system_deallocate_data_arrays(system)
      call MW_system_void_parameters(system)

      ! Run.out
      if (MW_stdout > 0) then
         inquire(unit=MW_stdout, opened=lopen)
         if (lopen) then
            close(MW_stdout, iostat=ierr)
            if (ierr /= 0) then
               call MW_errors_close_error("MW_system_close", "system.f90", &
                     MW_stdout, ierr)
            end if
         end if
      end if

   end subroutine finalize

   !================================================================================

   subroutine finalize_parallel()
      use MW_parallel, only: MW_parallel_finalize => finalize

      call MW_parallel_finalize()

   end subroutine finalize_parallel

   !================================================================================
   subroutine compute_lj_energy(this, lj, h)
      use MW_kinds, only: wp
      use MW_system, only: MW_system_t
      use MW_lennard_jones, only: MW_lennard_jones_t, &
            MW_lennard_jones_energy => energy, &
            MW_lennard_jones_fix_molecule_energy => fix_molecule_energy

      implicit none
      type(MW_system_t), intent(in) :: this
      type(MW_lennard_jones_t), intent(in) :: lj

      real(wp), intent(out) :: h(4)       !< Potential energy due to L-J interactions

      ! Local
      integer :: xyz_now

      h(:) = 0.0_wp
      xyz_now = this%xyz_ions_step(1)

      if (this%lj%num_species > 0) then
         call MW_lennard_jones_energy(this%num_pbc, this%localwork, lj, this%box, this%ions, &
               this%xyz_ions(:,:,xyz_now), this%electrodes, this%xyz_atoms, h(1:2), .False.)

         if (this%num_molecule_types > 0) then
            call MW_lennard_jones_fix_molecule_energy(this%num_PBC, this%localwork, lj, &
                  this%molecules, this%box,this%ions, this%xyz_ions(:,:,xyz_now), h(3:4), .False.)
         end if
      end if

   end subroutine compute_lj_energy

   !================================================================================
   subroutine compute_ft_energy(this, ft, h)
      use MW_kinds, only: wp
      use MW_system, only: MW_system_t
      use MW_fumi_tosi, only: MW_fumi_tosi_t, &
            MW_fumi_tosi_energy => energy, &
            MW_fumi_tosi_fix_molecule_energy => fix_molecule_energy

      implicit none
      type(MW_system_t), intent(in) :: this
      type(MW_fumi_tosi_t), intent(in) :: ft

      real(wp), intent(out) :: h(2)       !< Potential energy due to L-J interactions

      ! Local
      integer :: xyz_now

      h(:) = 0.0_wp
      xyz_now = this%xyz_ions_step(1)

      if (this%ft%num_species > 0) then
         call MW_fumi_tosi_energy(this%num_PBC, this%localwork, ft, this%box, &
               this%ions, this%xyz_ions(:,:,xyz_now), &
               this%electrodes, this%xyz_atoms, h(1))

         if (this%num_molecule_types > 0) then
            call MW_fumi_tosi_fix_molecule_energy(this%num_PBC, this%localwork, ft, this%molecules, &
                  this%box,this%ions, this%xyz_ions(:,:,xyz_now), h(2))
         end if
      end if

   end subroutine compute_ft_energy

   !================================================================================
   subroutine compute_xft_energy(this, xft, h)
      use MW_kinds, only: wp
      use MW_system, only: MW_system_t
      use MW_xft, only: MW_xft_t, &
            MW_xft_energy => energy, &
            MW_xft_fix_molecule_energy => fix_molecule_energy

      implicit none
      type(MW_system_t), intent(in) :: this
      type(MW_xft_t), intent(in) :: xft

      real(wp), intent(out) :: h(2)       !< Potential energy due to XFT interactions

      ! Local
      integer :: xyz_now

      h(:) = 0.0_wp
      xyz_now = this%xyz_ions_step(1)

      if (this%xft%num_species > 0) then
         call MW_xft_energy(this%num_PBC, this%localwork, xft, this%box, &
               this%ions, this%xyz_ions(:,:,xyz_now), &
               this%electrodes, this%xyz_atoms, h(1))

         if (this%num_molecule_types > 0) then
            call MW_xft_fix_molecule_energy(this%num_PBC, this%localwork, xft, this%molecules, &
                  this%box,this%ions, this%xyz_ions(:,:,xyz_now), h(2))
         end if
      end if

   end subroutine compute_xft_energy

   !================================================================================
   subroutine modify_charge(this, itype, charge)
      use MW_kinds, only: wp
      use MW_system, only: MW_system_t

      implicit none
      type(MW_system_t), intent(inout) :: this
      integer, intent(in) :: itype
      real(wp), intent(in) :: charge

      this%ions(itype)%charge = charge

   end subroutine

   !================================================================================
   subroutine get_electrode_mass(this, itype, mass)
      use MW_kinds, only: wp
      use MW_system, only: MW_system_t

      implicit none
      type(MW_system_t), intent(in) :: this
      integer, intent(in) :: itype
      real(wp), intent(out) :: mass

      mass = this%electrodes(itype)%mass

   end subroutine

   !================================================================================
   subroutine set_electrode_mass(this, itype, mass)
      use MW_kinds, only: wp
      use MW_system, only: MW_system_t

      implicit none
      type(MW_system_t), intent(inout) :: this
      integer, intent(in) :: itype
      real(wp), intent(in) :: mass

      this%electrodes(itype)%mass = mass

   end subroutine

   !================================================================================
   subroutine sum_electrode_forces(this, itypes)
      use MW_kinds, only: wp
      use MW_system, only: MW_system_t
      use MW_errors, only: MW_errors_allocate_error => allocate_error

      implicit none
      type(MW_system_t), intent(inout) :: this
      integer, intent(in) :: itypes(:)

      !Local
      real(wp), allocatable :: forces(:, :)
      integer :: num_types, elec, itype, ierr

      num_types = size(itypes)
      allocate(forces(3, size(this%electrodes(1)%force_ions,2)), stat=ierr)
      if (ierr /= 0) then
         call MW_errors_allocate_error("allocate_data_arrays", "electrode.f90", ierr)
      end if
      forces(:,:) = 0.0_wp

      do elec = 1, num_types
         itype = itypes(elec)
         forces(:, :) = forces(:, :) + this%electrodes(itype)%force_ions(:, :)
      end do

      do elec = 1, num_types
         itype = itypes(elec)
         this%electrodes(itype)%force_ions(:, :) = forces(:, :)
      end do

   end subroutine

end module MW_tools
