! ================================================================================
!> Compute the forces due to harmonic angle potential
subroutine forces_@PBC@(localwork, box, molecules, ions, xyz_ions, force,stress_tensor)
   implicit none
   ! Parameters
   ! ----------
   type(MW_localwork_t), intent(in) :: localwork
   type(MW_box_t), intent(in) :: box
   type(MW_molecule_t), intent(in) :: molecules(:)
   type(MW_ion_t), intent(in) :: ions(:)
   real(wp), intent(in) :: xyz_ions(:,:)
   real(wp), intent(inout) :: force(:,:)
   real(wp), intent(inout) :: stress_tensor(:,:)

   ! Locals
   ! ------
   integer :: num_molecule_types, imoltype, imol
   integer :: num_harmonic_angles, iangle
   integer :: iatype, ibtype, ictype, ia, ib, ic
   integer :: iasite, ibsite, icsite
   real(wp) :: a, b, c
   real(wp) :: dxab, dyab, dzab, drnormab2, drnormab
   real(wp) :: dxac, dyac, dzac, drnormac2, drnormac
   real(wp) :: dxbc, dybc, dzbc, drnormbc2, drnormbc
   real(wp) :: k, theta0, theta, fab, fac, fbc
   real(wp) :: cos_theta
   real(wp) :: dudtheta, dacosdz, dzdab, dzdac, dzdbc
   num_molecule_types = size(molecules,1)

   call MW_timers_start(TIMER_INTRA_MELT_FORCES)

   a = box%length(1)
   b = box%length(2)
   c = box%length(3)

   !$acc data &
   !$acc present(localwork, localwork%molecules_istart(:), localwork%molecules_iend(:), &
   !$acc         xyz_ions(:,:), ions(:), force(:,:), stress_tensor(:,:))

   do imoltype = 1, num_molecule_types
      num_harmonic_angles = molecules(imoltype)%num_harmonic_angles
      if (num_harmonic_angles > 0) then
         do iangle = 1, num_harmonic_angles
            ! angle parameters
            k = molecules(imoltype)%harmonic_angles_strength(iangle)
            theta0 = molecules(imoltype)%harmonic_angles_angle(iangle)

            iasite = molecules(imoltype)%harmonic_angles_sites(1,iangle)
            ibsite = molecules(imoltype)%harmonic_angles_sites(2,iangle)
            icsite = molecules(imoltype)%harmonic_angles_sites(3,iangle)

            iatype = molecules(imoltype)%sites(iasite)
            ibtype = molecules(imoltype)%sites(ibsite)
            ictype = molecules(imoltype)%sites(icsite)

            !$acc parallel loop &
            !$acc private(imol, ia, ib, cos_theta, theta, dudtheta, dacosdz, dzdab, dzdac, dzdbc, &
            !$acc         fab, fac, fbc, &
            !$acc         dxab, dyab, dzab, drnormab2, drnormab, &
            !$acc         dxac, dyac, dzac, drnormac2, drnormac, &
            !$acc         dxbc, dybc, dzbc, drnormbc2, drnormbc)
            do imol = localwork%molecules_istart(imoltype), localwork%molecules_iend(imoltype)

               ia = ions(iatype)%offset + imol
               ib = ions(ibtype)%offset + imol
               ic = ions(ictype)%offset + imol

               call minimum_image_displacement_@PBC@(a, b, c, &
                     xyz_ions(ia,1), xyz_ions(ia,2), xyz_ions(ia,3), &
                     xyz_ions(ib,1), xyz_ions(ib,2), xyz_ions(ib,3), &
                     dxab, dyab, dzab, drnormab2)
               drnormab = sqrt(drnormab2)

               call minimum_image_displacement_@PBC@(a, b, c, &
                     xyz_ions(ia,1), xyz_ions(ia,2), xyz_ions(ia,3), &
                     xyz_ions(ic,1), xyz_ions(ic,2), xyz_ions(ic,3), &
                     dxac, dyac, dzac, drnormac2)
               drnormac = sqrt(drnormac2)

               call minimum_image_displacement_@PBC@(a, b, c, &
                     xyz_ions(ib,1), xyz_ions(ib,2), xyz_ions(ib,3), &
                     xyz_ions(ic,1), xyz_ions(ic,2), xyz_ions(ic,3), &
                     dxbc, dybc, dzbc, drnormbc2)
               drnormbc = sqrt(drnormbc2)

               ! bc = ac - ab => |bc|^2 = |ac|^2 + |ab|^2 - 2*|ab|*|ac|*cos(theta)
               cos_theta = (drnormab2 + drnormac2 - drnormbc2) / (2.0_wp * drnormab * drnormac)
               if (cos_theta > 1) then
                  cos_theta = 1
                  theta = acos(1.0_wp)
               else if (cos_theta < -1) then
                  cos_theta = -1
                  theta = acos(-1.0_wp)
               else
                  theta = acos(cos_theta)
               end if

               ! Compute derivative of the potential
               ! du/dtheta = 2 * k * (theta - theta0)
               ! d(acos(z))/dz = -1 / sqrt(1-z^2)
               ! z = cos_theta
               ! dz/d|ab| = (|ab|^2 - |ac|^2 + |bc|^2) / (2*|ab|^2*|ac|)
               ! dz/d|ac| = (|ac|^2 - |ab|^2 + |bc|^2) / (2*|ac|^2*|ab|)
               ! dz/d|bc| = -|bc| / (|ab|*|ac|)
               dudtheta = 2.0_wp * k * (theta - theta0)
               if (cos_theta*cos_theta == 1) then
                  dacosdz = 0
               else
                   dacosdz = -1.0_wp / sqrt(1 - cos_theta*cos_theta)
               end if
               dzdab = (drnormab2 - drnormac2 + drnormbc2) / (2.0_wp*drnormab2*drnormac)
               dzdac = (drnormac2 - drnormab2 + drnormbc2) / (2.0_wp*drnormac2*drnormab)
               dzdbc = -drnormbc / (drnormac*drnormab)

               ! du/dr_ab = k * (theta - theta0) * (- 1 / sqrt(1-cos_theta^2)) * (drab^2 - drac^2 + drbc^2) / (drab^2*drac)
               fab = - (dudtheta*dacosdz*dzdab) / drnormab

               !$acc atomic update
               force(ib,1) = force(ib,1) + dxab * fab
               !$acc atomic update
               force(ib,2) = force(ib,2) + dyab * fab
               !$acc atomic update
               force(ib,3) = force(ib,3) + dzab * fab

               !$acc atomic update
               force(ia,1) = force(ia,1) - dxab * fab
               !$acc atomic update
               force(ia,2) = force(ia,2) - dyab * fab
               !$acc atomic update
               force(ia,3) = force(ia,3) - dzab * fab

               ! du/dr_ac = k * (theta - theta0) * (- 1 / sqrt(1-cos_theta^2)) * (drac^2 - drab^2 + drbc^2) / (drac^2*drab)
               fac = - (dudtheta*dacosdz*dzdac) / drnormac

               !$acc atomic update
               force(ic,1) = force(ic,1) + dxac * fac
               !$acc atomic update
               force(ic,2) = force(ic,2) + dyac * fac
               !$acc atomic update
               force(ic,3) = force(ic,3) + dzac * fac

               !$acc atomic update
               force(ia,1) = force(ia,1) - dxac * fac
               !$acc atomic update
               force(ia,2) = force(ia,2) - dyac * fac
               !$acc atomic update
               force(ia,3) = force(ia,3) - dzac * fac

               ! du/dr_bc = k * (theta - theta0) * (- 1 / sqrt(1-cos_theta^2)) * (drac^2 - drab^2 + drbc^2) / (drac^2*drab)
               fbc = - (dudtheta*dacosdz*dzdbc) / drnormbc

               !$acc atomic update
               force(ic,1) = force(ic,1) + dxbc * fbc
               !$acc atomic update
               force(ic,2) = force(ic,2) + dybc * fbc
               !$acc atomic update
               force(ic,3) = force(ic,3) + dzbc * fbc

               !$acc atomic update
               force(ib,1) = force(ib,1) - dxbc * fbc
               !$acc atomic update
               force(ib,2) = force(ib,2) - dybc * fbc
               !$acc atomic update
               force(ib,3) = force(ib,3) - dzbc * fbc

            end do
            !$acc end parallel loop
         end do
      end if
   end do
   !$acc end data

   call MW_timers_stop(TIMER_INTRA_MELT_FORCES)   
end subroutine forces_@PBC@

! ================================================================================
!> Compute the forces due to harmonic angle potential
subroutine energy_@PBC@(localwork, box, molecules, ions, xyz_ions, h)
   implicit none
   ! Parameters
   ! ----------
   type(MW_localwork_t), intent(in) :: localwork
   type(MW_box_t), intent(in) :: box
   type(MW_molecule_t), intent(in) :: molecules(:)
   type(MW_ion_t), intent(in) :: ions(:)
   real(wp), intent(in) :: xyz_ions(:,:)
   real(wp), intent(inout) :: h

   ! Locals
   ! ------
   integer :: num_molecule_types, imoltype, imol
   integer :: num_harmonic_angles, iangle
   integer :: iatype, ibtype, ictype, ia, ib, ic
   integer :: iasite, ibsite, icsite
   real(wp) :: a, b, c
   real(wp) :: drnormab2, drnormab
   real(wp) :: drnormac2, drnormac
   real(wp) :: drnormbc2, drnormbc
   real(wp) :: k, theta0, Uab
   real(wp) :: cos_theta, theta

   call MW_timers_start(TIMER_INTRA_MELT_POTENTIAL)

   num_molecule_types = size(molecules,1)

   a = box%length(1)
   b = box%length(2)
   c = box%length(3)

   do imoltype = 1, num_molecule_types
      num_harmonic_angles = molecules(imoltype)%num_harmonic_angles
      if (num_harmonic_angles > 0) then
         do imol = localwork%molecules_istart(imoltype), localwork%molecules_iend(imoltype)
            do iangle = 1, num_harmonic_angles
               ! angle parameters
               k = molecules(imoltype)%harmonic_angles_strength(iangle)
               theta0 = molecules(imoltype)%harmonic_angles_angle(iangle)

               iasite = molecules(imoltype)%harmonic_angles_sites(1,iangle)
               ibsite = molecules(imoltype)%harmonic_angles_sites(2,iangle)
               icsite = molecules(imoltype)%harmonic_angles_sites(3,iangle)

               iatype = molecules(imoltype)%sites(iasite)
               ibtype = molecules(imoltype)%sites(ibsite)
               ictype = molecules(imoltype)%sites(icsite)

               ia = ions(iatype)%offset + imol
               ib = ions(ibtype)%offset + imol
               ic = ions(ictype)%offset + imol

               call minimum_image_distance_@PBC@(a, b, c, &
                     xyz_ions(ia,1), xyz_ions(ia,2), xyz_ions(ia,3), &
                     xyz_ions(ib,1), xyz_ions(ib,2), xyz_ions(ib,3), &
                     drnormab2)
               drnormab = sqrt(drnormab2)

               call minimum_image_distance_@PBC@(a, b, c, &
                     xyz_ions(ia,1), xyz_ions(ia,2), xyz_ions(ia,3), &
                     xyz_ions(ic,1), xyz_ions(ic,2), xyz_ions(ic,3), &
                     drnormac2)
               drnormac = sqrt(drnormac2)

               call minimum_image_distance_@PBC@(a, b, c, &
                     xyz_ions(ib,1), xyz_ions(ib,2), xyz_ions(ib,3), &
                     xyz_ions(ic,1), xyz_ions(ic,2), xyz_ions(ic,3), &
                     drnormbc2)
               drnormbc = sqrt(drnormbc2)

               ! bc = ac - bc => |bc|^2 = |ac|^2 + |ab|^2 - 2*|ab|*|ac|*cos(theta)
               cos_theta = (drnormab2 + drnormac2 - drnormbc2) / (2.0_wp * drnormab * drnormac)
               theta = acos(cos_theta)

               Uab = k * (theta - theta0) * (theta - theta0)
               h = h + Uab
            end do
         end do
      end if
   end do

   call MW_timers_stop(TIMER_INTRA_MELT_POTENTIAL)

end subroutine energy_@PBC@


