module MW_four_site_model
#ifndef MW_SERIAL
   use MPI
#endif
   use MW_kinds, only: wp
   use MW_localwork, only: MW_localwork_t
   use MW_molecule, only: MW_molecule_t
   use MW_ion, only: MW_ion_t
   use MW_box, only: MW_box_t
   use MW_parallel, only: MW_COMM_WORLD
   
   ! Public subroutines
   ! ------------------
   public :: compute_fourth_site_position

contains

   ! ========================================================
   subroutine compute_fourth_site_position(num_pbc, molecules, xyz_ions, box, ions)

      implicit none

      ! Parameters
      ! ----------
      integer, intent(in) :: num_pbc
      type(MW_molecule_t), intent(in) :: molecules(:) !< molecules parameters
      type(MW_box_t), intent(in) :: box !< box parameters
      type(MW_ion_t), intent(in) :: ions(:) !< ions parameters
      real(wp), intent(inout) :: xyz_ions(:,:)  !< ions xyz positions

      select case(num_pbc)
      case(2)
         call compute_fourth_site_position_2DPBC(molecules, xyz_ions, box, ions)
      case(3)
         call compute_fourth_site_position_3DPBC(molecules, xyz_ions, box, ions)
      end select

   end subroutine

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

   subroutine redistribute_force(num_pbc, localwork, molecules, xyz_ions, forces_ions, box, ions, &
         stress_tensor)
      use Mw_errors, only: MW_errors_allocate_error => allocate_error
      implicit none

      ! Parameters
      ! ----------
      integer, intent(in) :: num_pbc
      type(MW_localwork_t), intent(in) :: localwork !< Local work assignment
      type(MW_molecule_t), intent(in) :: molecules(:) !< molecules parameters
      type(MW_box_t), intent(in) :: box !< box parameters
      type(MW_ion_t), intent(in) :: ions(:) !< ions parameters
      real(wp), intent(inout) :: xyz_ions(:,:)  !< ions xyz positions
      real(wp), intent(inout) :: forces_ions(:,:)  !< ions forces
      real(wp), intent(inout) :: stress_tensor(:,:)  !< stress tensor

      ! Local
      ! -----
      real(wp), allocatable        :: force_to_distribute(:,:)
      integer :: n, ixyz, iion, count, ierr

      n = size(forces_ions, 1)

      allocate(force_to_distribute(n, 3), stat=ierr)
      if (ierr /= 0) then
         call MW_errors_allocate_error("allocate_data_arrays", "four_site_model.f90", ierr)
      end if
      force_to_distribute(:,:) = 0.0_wp

      select case(num_pbc)
      case(2)
         call redistribute_force_2DPBC(localwork, molecules, xyz_ions, forces_ions, box, ions, force_to_distribute, &
               stress_tensor)
      case(3)
         call redistribute_force_3DPBC(localwork, molecules, xyz_ions, forces_ions, box, ions, force_to_distribute, &
               stress_tensor)
      end select

#ifndef MW_SERIAL
      count = 3*n
      call MPI_Allreduce(MPI_IN_PLACE, force_to_distribute, count, MPI_DOUBLE_PRECISION, MPI_SUM, MW_COMM_WORLD, ierr)
      count = 3*3
      call MPI_Allreduce(MPI_IN_PLACE, stress_tensor, count, MPI_DOUBLE_PRECISION, MPI_SUM, MW_COMM_WORLD, ierr)
#endif

      do ixyz = 1, 3
         do iion = 1, n
            forces_ions(iion,ixyz) = forces_ions(iion,ixyz) + force_to_distribute(iion,ixyz)  
         end do
      end do

   end subroutine redistribute_force

   ! ================================================================================
   include 'four_site_model_2DPBC.inc'
   include 'four_site_model_3DPBC.inc'
   ! ================================================================================
   include 'minimum_image_displacement_2DPBC.inc'
   include 'minimum_image_displacement_3DPBC.inc'

end module MW_four_site_model
