! Defines the simulation box which is the unit cell of the system
!
! A box is orthorombic and is defined by its 3 lengths
! The box defined in this module has 2D periodic boundary conditions in x and y
module MW_box
   use MW_kinds, only: wp
   implicit none
   private

   ! public types
   public :: MW_box_t

   ! public subroutine
   public :: define_type
   public :: void_type
   public :: print_type
   public :: update_type
   public :: minimum_image_distance_2DPBC
   public :: minimum_image_distance_3DPBC
   public :: minimum_image_displacement_2DPBC
   public :: minimum_image_displacement_3DPBC
   public :: fix_outliers
   public :: unwrap
   !   public :: minimum_image_displacement_int
   !   public :: minimum_image_distance_int

   ! Box datatype
   type MW_box_t
      real(wp) :: length(3)      !< length of each side of the box
      real(wp) :: volume         !< volume of the box
      real(wp) :: area(3)        !< (|bxc|,|cxa|,|axb|)
      real(wp) :: height(3)      !< Perpendicular width
   end type MW_box_t

contains

   !================================================================================
   !> Define a box from 3 lengths
   !
   subroutine define_type(this, a, b, c)
      use MW_errors, only: MW_errors_parameter_error => parameter_error
      implicit none
      ! Parameters
      ! ----------
      type(MW_box_t), intent(inout) :: this

      real(wp), intent(in) :: a !> 1st dimension vector length
      real(wp), intent(in) :: b !> 2nd dimension vector length
      real(wp), intent(in) :: c !> 3rd dimension vector length

      ! Lengths (must be non-zero)
      ! --------------------------
      this%length(1) = a
      if (this%length(1) <= 0.0) then
         call MW_errors_parameter_error("define_type", "box.f90",&
               "length(a)", this%length(1))
      end if

      this%length(2) = b
      if (this%length(2) <= 0.0) then
         call MW_errors_parameter_error("define_type", "box.f90",&
               "length(b)", this%length(2))
      end if

      this%length(3) = c
      if (this%length(3) <= 0.0) then
         call MW_errors_parameter_error("define_type", "box.f90",&
               "length(c)", this%length(3))
      end if

      ! Volume (must be > 0)
      ! ------
      ! V = |a|x|b|x|c|
      this%volume = this%length(1) * this%length(2) * this%length(3)
      ! The 3 vectors may not be coplanar and must be properly oriented
      if (this%volume <= 0.0) then
         call MW_errors_parameter_error("define_type", "box.f90",&
               "volume", this%volume)
      end if

      ! Areas (must be > 0)
      ! -------------------
      ! A(b,c) = | b x c |
      this%area(1) = this%length(2) * this%length(3)

      ! A(c,a) = | c x a |
      this%area(2) = this%length(3) * this%length(1)

      ! A(a,b) = | a x b |
      this%area(3) = this%length(1) * this%length(2)

      ! Perpendicular widths
      ! --------------------
      this%height(1) = this%length(1)
      this%height(2) = this%length(2)
      this%height(3) = this%length(3)

   end subroutine define_type

   !================================================================================
   !> Void the data structure
   subroutine void_type(this)
      implicit none
      ! Parameters
      ! ----------
      type(MW_box_t), intent(inout) :: this

      this%length(:)     = 0.0_wp
      this%volume        = 0.0_wp
      this%area(:)       = 0.0_wp
      this%height(:)     = 0.0_wp
   end subroutine void_type

   !================================================================================
   !> Print the data structure
   subroutine print_type(this, ounit)
      implicit none
      ! Parameters
      ! ----------
      type(MW_box_t), intent(in) :: this
      integer,        intent(in) :: ounit

      write(ounit, '("|box| edge length :       ",3(1x,es12.5))') this%length
      write(ounit, '("|box| side area:          ",3(1x,es12.5))') this%area
      write(ounit, '("|box| side height:        ",3(1x,es12.5))') this%height
      write(ounit, '("|box| volume:             ",1(1x,es12.5))') this%volume

   end subroutine print_type

   !================================================================================
   !> update box parameters during NPT runs
   !
   subroutine update_type(this, volume_factor)
      implicit none
      ! Parameters
      ! ----------
      type(MW_box_t), intent(inout) :: this

      real(wp), intent(in) :: volume_factor   !> Variation of the volume during the step


      ! Local
      real(wp) :: length_factor

      length_factor = volume_factor**(1.0/3.0)

      this%length(1) = this%length(1) * length_factor
      this%length(2) = this%length(2) * length_factor
      this%length(3) = this%length(3) * length_factor

      ! Areas (must be > 0)
      ! -------------------
      ! A(b,c) = | b x c |
      this%area(1) = this%length(2) * this%length(3)

      ! A(c,a) = | c x a |
      this%area(2) = this%length(3) * this%length(1)

      ! A(a,b) = | a x b |
      this%area(3) = this%length(1) * this%length(2)

      ! Perpendicular widths
      ! --------------------
      this%height(1) = this%length(1)
      this%height(2) = this%length(2)
      this%height(3) = this%length(3)

   end subroutine update_type
   

   include 'minimum_image_distance_2DPBC.inc'
   include 'minimum_image_displacement_2DPBC.inc'
   !================================================================================
   include 'minimum_image_distance_3DPBC.inc'
   include 'minimum_image_displacement_3DPBC.inc'

   !================================================================================
   !> Make sure xyz coordinates are inside the box, update xyz if necessary
   !!
   !! 0 <= x < box%length(1)
   !! 0 <= y < box%length(2)
   !! 0 <= z < box%length(3)
   pure subroutine fix_outliers(this, numPBC, x, y, z)
      implicit none
      ! Parameters
      ! ----------
      type(MW_box_t), intent(in) :: this    !< structure
      integer, intent(in) :: numPBC         !< Number of periodic boundary conditions
      real(wp), intent(inout) :: x, y, z      !< xyz atom position

      select case (numPBC)
      case (1)
         x = x - this%length(1)*floor(x/this%length(1))
      case (2)
         x = x - this%length(1)*floor(x/this%length(1))
         y = y - this%length(2)*floor(y/this%length(2))
      case (3)
         x = x - this%length(1)*floor(x/this%length(1))
         y = y - this%length(2)*floor(y/this%length(2))
         z = z - this%length(3)*floor(z/this%length(3))
      case default
         continue
      end select

   end subroutine fix_outliers

   !================================================================================
   !> Unwrap xyz coordinates of a mobile ion.
   !!
   !! Implements Eq. (3) in J. Chem. Phys. 153, 021101 (2020); https://doi.org/10.1063/5.0008316
   !! The equations below are suitable for both NVT and NPT simulations:
   !! 
   !! xu(t) = xw(t) - floor([xw(t) - xu(t-1)]/Lx(t) + 0.5)*Lx(t)
   !! yu(t) = yw(t) - floor([yw(t) - yu(t-1)]/Ly(t) + 0.5)*Ly(t)
   !! zu(t) = zw(t) - floor([zw(t) - zu(t-1)]/Lz(t) + 0.5)*Lz(t)
   pure subroutine unwrap(this, numPBC, &
                          xu, yu, zu, xu_prev, yu_prev, zu_prev, xw, yw, zw)
      implicit none
      ! Parameters
      ! ----------
      type(MW_box_t), intent(in) :: this    !< structure
      integer, intent(in) :: numPBC         !< Number of periodic boundary conditions
      real(wp), intent(inout) :: xu, yu, zu !< unwrapped xyz ion position at the current time step
      real(wp), intent(in) :: xu_prev, yu_prev, zu_prev !< unwrapped xyz ion position at the previous time step
      real(wp), intent(in) :: xw, yw, zw !< wrapped xyz ion position at the current time step

      select case (numPBC)
      case (1)
         xu = xw - this%length(1)*floor((xw - xu_prev)/this%length(1) + 0.5_wp)
      case (2)
         xu = xw - this%length(1)*floor((xw - xu_prev)/this%length(1) + 0.5_wp)
         yu = yw - this%length(2)*floor((yw - yu_prev)/this%length(2) + 0.5_wp)
      case (3)
         xu = xw - this%length(1)*floor((xw - xu_prev)/this%length(1) + 0.5_wp)
         yu = yw - this%length(2)*floor((yw - yu_prev)/this%length(2) + 0.5_wp)
         zu = zw - this%length(3)*floor((zw - zu_prev)/this%length(3) + 0.5_wp)
      case default
         continue
      end select

   end subroutine unwrap

end module MW_box
