module multipoles
  use f_precisions, only: dp => f_double
  use f_enums
  implicit none
  private

  type(f_enumerator), public :: CHARACTER_NET_ENUM   = f_enumerator('net', 1, null())
  type(f_enumerator), public :: CHARACTER_GROSS_ENUM = f_enumerator('gross', 2, null())

  type, public :: monopole
     integer :: l = 0
     real(dp) :: sigma = 0._dp
     real(dp) :: qlm = 0._dp
  end type monopole

  type, public :: dipole
     integer :: l = 1
     real(dp) :: sigma = 0._dp
     real(dp), dimension(-1:1) :: qlm = (/ 0._dp, 0._dp, 0._dp /)
  end type dipole

  type, public :: quadrupole
     integer :: l = 2
     real(dp) :: sigma = 0._dp
     real(dp), dimension(-2:2) :: qlm = (/ 0._dp, 0._dp, 0._dp, 0._dp, 0._dp /)
  end type quadrupole

  type,public :: multipole_center
     character(len=20) :: label
     real(dp), dimension(3) :: rxyz
     type(monopole) :: mono
     type(dipole) :: di
     type(quadrupole) :: quad
     type(f_enumerator) :: character
  end type multipole_center

  interface operator(==)
     module procedure monopole_is_monopole, dipole_is_dipole, quadrupole_is_quadrupole
  end interface operator(==)

  interface operator(/=)
     module procedure monopole_isnot_monopole, dipole_isnot_dipole, quadrupole_isnot_quadrupole
  end interface operator(/=)

  interface operator(/=)
     module procedure center_isnot_center
  end interface operator(/=)

  interface operator(==)
     module procedure center_is_center
  end interface operator(==)

  interface pole_at
     module procedure monopole_at, dipole_at, quadpole_at
  end interface pole_at

  interface multipole_gaussian_density_at
     module procedure multipole_gaussian_density_at_bit, &
          & multipole_gaussian_density_at_rxyz, multipole_gaussian_density_at_delta
  end interface multipole_gaussian_density_at

  interface multipole_at
     module procedure multipole_at_bit, multipole_at_rxyz
  end interface multipole_at

  interface pole
     module procedure monopole_new, dipole_new, quadrupole_new
  end interface pole

  public :: pole
  public :: multipole_center_set
  public :: multipole_centers_to_dict, multipole_centers_from_dict
  public :: multipole_centers_to_gaussian_density
  public :: operator(/=), operator(==)
  public :: multipole_to_gaussian_subbox, multipole_gaussian_density_at
  public :: multipole_at

contains

  pure function has_monopole(center) result(ok)
    implicit none
    type(multipole_center), intent(in) :: center
    logical :: ok

    ok = (center%mono%qlm /= 0._dp)
  end function has_monopole

  pure function has_dipole(center) result(ok)
    implicit none
    type(multipole_center), intent(in) :: center
    logical :: ok

    ok = all(center%di%qlm /= 0._dp)
  end function has_dipole

  pure function has_quadrupole(center) result(ok)
    implicit none
    type(multipole_center), intent(in) :: center
    logical :: ok

    ok = all(center%quad%qlm /= 0._dp)
  end function has_quadrupole

  elemental pure function monopole_is_monopole(a, b) result(ok)
    implicit none
    type(monopole), intent(in) :: a, b
    logical :: ok
    real(dp), parameter :: eps = 1.d-12

    ok = (abs(a%qlm - b%qlm) < eps) .and. (abs(a%sigma - b%sigma) < eps)
  end function monopole_is_monopole

  elemental pure function dipole_is_dipole(a, b) result(ok)
    implicit none
    type(dipole), intent(in) :: a, b
    logical :: ok
    real(dp), parameter :: eps = 1.d-12

    ok = (all(abs(a%qlm - b%qlm) < eps)) .and. (abs(a%sigma - b%sigma) < eps)
  end function dipole_is_dipole

  elemental pure function quadrupole_is_quadrupole(a, b) result(ok)
    implicit none
    type(quadrupole), intent(in) :: a, b
    logical :: ok
    real(dp), parameter :: eps = 1.d-12

    ok = (all(abs(a%qlm - b%qlm) < eps)) .and. (abs(a%sigma - b%sigma) < eps)
  end function quadrupole_is_quadrupole

  elemental pure function monopole_isnot_monopole(a, b) result(ok)
    implicit none
    type(monopole), intent(in) :: a, b
    logical :: ok

    ok = .not. (a == b)
  end function monopole_isnot_monopole

  elemental pure function dipole_isnot_dipole(a, b) result(ok)
    implicit none
    type(dipole), intent(in) :: a, b
    logical :: ok

    ok = .not. (a == b)
  end function dipole_isnot_dipole

  elemental pure function quadrupole_isnot_quadrupole(a, b) result(ok)
    implicit none
    type(quadrupole), intent(in) :: a, b
    logical :: ok

    ok = .not. (a == b)
  end function quadrupole_isnot_quadrupole

  elemental pure function center_is_center(a, b) result(ok)
    implicit none
    type(multipole_center), intent(in) :: a, b
    logical :: ok

    real(dp), parameter :: eps = 1.d-12
    
    ok = trim(a%label) == trim(b%label)
    ok = ok .and. (a%character == b%character)
    ok = ok .and. (all(abs(a%rxyz - b%rxyz) < eps))
    ok = ok .and. (a%mono == b%mono)
    ok = ok .and. (a%di == b%di)
    ok = ok .and. (a%quad == b%quad)
  end function center_is_center

  elemental pure function center_isnot_center(a, b) result(ok)
    implicit none
    type(multipole_center), intent(in) :: a, b
    logical :: ok
    ok = .not. center_is_center(a, b)
  end function center_isnot_center

  pure function monopole_new(qlm, sigma) result(out)
    implicit none
    real(dp), intent(in) :: qlm, sigma
    type(monopole) :: out

    out%qlm = qlm
    out%sigma = sigma
  end function monopole_new

  pure function dipole_new(q1, q2, q3, sigma) result(out)
    implicit none
    real(dp), intent(in) :: q1, q2, q3, sigma
    type(dipole) :: out

    out%qlm = (/ q1, q2, q3 /)
    out%sigma = sigma
  end function dipole_new

  pure function quadrupole_new(q1, q2, q3, q4, q5, sigma) result(out)
    implicit none
    real(dp), intent(in) :: q1, q2, q3, q4, q5, sigma
    type(quadrupole) :: out

    out%qlm = (/ q1, q2, q3, q4, q5 /)
    out%sigma = sigma
  end function quadrupole_new

  subroutine multipole_center_set(center, rxyz, label, mono, di, quad, charac)
    implicit none
    type(multipole_center), intent(inout) :: center
    real(dp), dimension(3), intent(in) :: rxyz
    character(len = *), intent(in), optional :: label
    type(monopole), intent(in), optional :: mono
    type(dipole), intent(in), optional :: di
    type(quadrupole), intent(in), optional :: quad
    type(f_enumerator), intent(in), optional :: charac

    center%rxyz = rxyz
    if (present(label)) center%label = label
    if (present(mono)) center%mono = mono
    if (present(di)) center%di = di
    if (present(quad)) center%quad = quad
    if (present(charac)) center%character = charac
  end subroutine multipole_center_set

  subroutine center_to_dict(dict, center, factor)
    use dictionaries, dict_set => set, dict_add => add
    implicit none
    type(dictionary), pointer, intent(inout) :: dict
    type(multipole_center), intent(in) :: center
    real(dp), intent(in) :: factor

    type(dictionary), pointer :: sigmas

    if (.not. associated(dict)) call dict_init(dict)
    call dict_set(dict // "sym", center%label)
    call dict_set(dict // "multipole character", toa(center%character))
    call dict_set(dict // "r", center%rxyz * factor)
    if (has_monopole(center))   call dict_set(dict // "q0", (/ center%mono%qlm /))
    if (has_dipole(center))     call dict_set(dict // "q1", center%di%qlm)
    if (has_quadrupole(center)) call dict_set(dict // "q2", center%quad%qlm)
    call dict_init(sigmas)
    if (has_monopole(center))   call dict_add(sigmas, center%mono%sigma)
    if (has_dipole(center))     call dict_add(sigmas, center%di%sigma)
    if (has_quadrupole(center)) call dict_add(sigmas, center%quad%sigma)
    call dict_set(dict // "sigma", sigmas)
  end subroutine center_to_dict
  
  subroutine multipole_centers_to_dict(dict, centers, label, units)
    use dictionaries, dict_set => set, dict_add => add
    use yaml_strings, only: f_strcpy
    use numerics, only: Bohr_Ang
    implicit none
    type(dictionary), pointer, intent(inout) :: dict
    type(multipole_center), dimension(:), intent(in) :: centers
    character(len = *), intent(in) :: label
    character(len = *), intent(in), optional :: units

    character(len = max_field_length) :: myUnits
    real(dp) :: factor
    type(dictionary), pointer :: out, myCenters, oneCenter
    integer :: i
    
    if (.not. associated(dict)) call dict_init(dict)

    if (present(units)) then
       call f_strcpy(src = units, dest = myUnits)
    else
       myUnits = "atomic"
    end if
    select case (trim(myUnits))
    case ('angstroem', 'angstroemd0')
       factor = Bohr_Ang
    case ('atomic', 'atomicd0', 'bohr', 'bohrd0')
       factor = 1._dp
    case default
       factor = 1._dp
       call f_err_throw('Invalid units ('//trim(myUnits)//') for external multipoles.')
    end select

    call dict_init(out)
    
    call dict_set(out // "units", myUnits)

    call dict_init(myCenters)
    do i = 1, size(centers)
       call dict_init(oneCenter)
       call center_to_dict(oneCenter, centers(i), factor)
       call dict_add(myCenters, oneCenter)
    end do
    call dict_set(out // "values", myCenters)

    call dict_set(dict // label, out)
  end subroutine multipole_centers_to_dict

  subroutine center_from_dict(center, dict, factor)
    use dictionaries
    implicit none
    type(multipole_center), intent(out) :: center
    type(dictionary), pointer :: dict
    real(dp), intent(in) :: factor

    character(len = max_field_length) :: character
    integer :: i
    real(dp), dimension(1) :: q0
    
    if ("sym" .notin. dict) call f_err_throw("Missing 'sym' definition.")
    center%label = dict // "sym"

    if ("r" .notin. dict) call f_err_throw("Missing 'r' definition.")
    center%rxyz = dict // "r"
    center%rxyz = center%rxyz * factor

    if ("multipole character" .notin. dict) call f_err_throw("Missing 'multipole character' definition.")
    character = dict // "multipole character"
    if (CHARACTER_NET_ENUM == character) then
       center%character = CHARACTER_NET_ENUM
    else if (CHARACTER_GROSS_ENUM == character) then
       center%character = CHARACTER_GROSS_ENUM
    else
       call f_err_throw("Invalid multipole character.")
    end if

    i = 0
    if ("q0" .in. dict) then
       if (dict_islist(dict // "q0")) then
          q0 = dict // "q0"
          center%mono%qlm = q0(1)
       else
          center%mono%qlm = dict // "q0"
       end if
       if ("sigma" .notin. dict) call f_err_throw("Missing 'sigma' definition.")
       center%mono%sigma = dict // "sigma" // i
       i = i + 1
    end if
    if ("q1" .in. dict) then
       center%di%qlm = dict // "q1"
       if ("sigma" .notin. dict) call f_err_throw("Missing 'sigma' definition.")
       center%di%sigma = dict // "sigma" // i
       i = i + 1
    end if
    if ("q2" .in. dict) then
       center%quad%qlm = dict // "q2"
       if ("sigma" .notin. dict) call f_err_throw("Missing 'sigma' definition.")
       center%quad%sigma = dict // "sigma" // i
       i = i + 1
    end if
  end subroutine center_from_dict
  
  subroutine multipole_centers_from_dict(centers, dict)
    use dictionaries
    use numerics, only: Bohr_Ang
    implicit none
    type(multipole_center), dimension(:), pointer :: centers
    type(dictionary), pointer :: dict

    integer :: i
    real(dp) :: factor
    character(len = max_field_length) :: myUnits
    type(dictionary), pointer :: iter
    
    if ("values" .notin. dict) call f_err_throw("Missing 'values' definition.")
    if (.not. dict_islist(dict // "values")) call f_err_throw("'values' data are not a list.")
    if (associated(centers)) deallocate(centers)
    allocate(centers(dict_len(dict // "values")))

    factor = 1._dp
    if ("units" .in. dict) then
       myUnits = dict // "units"
       select case (trim(myUnits))
       case ('angstroem', 'angstroemd0')
          factor = 1._dp / Bohr_Ang
       case ('atomic', 'atomicd0', 'bohr', 'bohrd0')
          factor = 1._dp
       case default
          call f_err_throw('Invalid units ('//trim(myUnits)//') for external multipoles.')
       end select
    end if

    i = 0
    iter => dict_iter(dict // "values")
    do while (associated(iter))
       i = i + 1
       call center_from_dict(centers(i), iter, factor)
       iter => dict_next(iter)
    end do
  end subroutine multipole_centers_from_dict

  pure function multipole_to_gaussian_subbox(center, bit) result(sub)
    use box
    implicit none
    type(multipole_center), intent(in) :: center
    type(box_iterator), intent(in) :: bit
    type(subbox) :: sub
    real(dp) :: cutoff

    cutoff = 10._dp * max(center%mono%sigma, center%di%sigma, center%quad%sigma)
    sub%nbox = box_nbox_from_cutoff(bit%mesh, center%rxyz + bit%oxyz, cutoff)
  end function multipole_to_gaussian_subbox

  pure function gaussian(sigma, r2) result(g)
    use numerics
    implicit none
    real(dp), intent(in) :: sigma, r2
    real(dp) :: tt, g

    ! Only calculate the Gaussian if the result will be larger than 10^-30
    tt = r2 / (2._dp * sigma**2)
    if (tt <= 69.07755279_dp) then
       g = safe_exp(-tt)
       g = g / sqrt(2._dp * pi * sigma**2)**3
    else
       g = 0._dp
    end if
    !g = g/(sigma**3*sqrt(2.d0*pi)**3)

  end function gaussian
  
  pure function monopole_at(mono, rxyz, r) result(val)
    use at_domain
    use numerics, only: pi
    use f_harmonics, only: solid_harmonic
    implicit none
    type(monopole), intent(in) :: mono
    real(dp), dimension(3), intent(in) :: rxyz
    real(dp), intent(in) :: r
    real(dp) :: val

    real(dp), parameter :: factor = 2._dp * sqrt(pi)
    real(dp), dimension(3), parameter :: O = (/ 0._dp, 0._dp, 0._dp /)
    real(dp) :: gg

    gg = gaussian(mono%sigma, r * r)
    if (gg > 0._dp) then
       val = mono%qlm * solid_harmonic(0, mono%l, 0, rxyz(1), rxyz(2), rxyz(3))
       val = factor * gg * val
    else
       val = 0._dp
    end if
  end function monopole_at

  pure function dipole_at(di, rxyz, r) result(val)
    use at_domain
    use numerics, only: pi
    use f_harmonics, only: solid_harmonic
    implicit none
    type(dipole), intent(in) :: di
    real(dp), dimension(3), intent(in) :: rxyz
    real(dp), intent(in) :: r
    real(dp) :: val

    real(dp), parameter :: factor = 2._dp * sqrt(3._dp * pi)
    real(dp), dimension(3), parameter :: O = (/ 0._dp, 0._dp, 0._dp /)
    real(dp) :: gg
    integer :: m

    val = 0._dp
    gg = gaussian(di%sigma, r * r)
    if (gg == 0._dp) return
    do m = -di%l, di%l
       val = val + di%qlm(m) * solid_harmonic(0, di%l, m, rxyz(1), rxyz(2), rxyz(3))
    end do
    val = factor * gg * val
  end function dipole_at

  pure function quadpole_at(quad, rxyz, r) result(val)
    use numerics, only: pi
    use f_harmonics, only: solid_harmonic
    implicit none
    type(quadrupole), intent(in) :: quad
    real(dp), dimension(3), intent(in) :: rxyz
    real(dp), intent(in) :: r
    real(dp) :: val

    real(dp), parameter :: factor = 2._dp * sqrt(5._dp * pi)
    real(dp), dimension(3), parameter :: O = (/ 0._dp, 0._dp, 0._dp /)
    real(dp) :: gg
    integer :: m

    val = 0._dp
    gg = gaussian(quad%sigma, r * r)
    if (gg == 0._dp) return
    do m = -quad%l, quad%l
       val = val + quad%qlm(m) * solid_harmonic(0, quad%l, m, rxyz(1), rxyz(2), rxyz(3))
    end do
    val = factor * gg * val
  end function quadpole_at

  pure function multipole_gaussian_density_at_bit(center, bit) result(val)
    use box
    use f_harmonics, only: solid_harmonic
    use numerics, only: pi
    implicit none
    type(multipole_center), intent(in) :: center
    type(box_iterator), intent(in) :: bit
    real(dp) :: val

    val = multipole_gaussian_density_at_rxyz(center, bit%mesh%dom, bit%rxyz)
  end function multipole_gaussian_density_at_bit

  pure function multipole_gaussian_density_at_rxyz(center, dom, at) result(val)
    use at_domain
    use f_harmonics, only: solid_harmonic
    use numerics, only: pi
    implicit none
    type(multipole_center), intent(in) :: center
    type(domain), intent(in) :: dom
    real(dp), dimension(3), intent(in) :: at
    real(dp) :: val

    val = multipole_gaussian_density_at_delta(center, at - center%rxyz, &
         & distance(dom, at, center%rxyz))
  end function multipole_gaussian_density_at_rxyz

  pure function multipole_gaussian_density_at_delta(center, dxyz, r) result(val)
    use f_harmonics, only: solid_harmonic
    use numerics, only: pi
    implicit none
    type(multipole_center), intent(in) :: center
    real(dp), dimension(3), intent(in) :: dxyz
    real(dp), intent(in) :: r
    real(dp) :: val

    val = 0._dp
    if (has_monopole(center)) val = val - pole_at(center%mono, dxyz, r)
    if (has_dipole(center)) val = val - pole_at(center%di, dxyz, r)
    if (has_quadrupole(center)) val = val - pole_at(center%quad, dxyz, r)
  end function multipole_gaussian_density_at_delta

  pure function multipole_at_bit(center, bit) result(val)
    use box
    use f_harmonics, only: solid_harmonic
    use numerics, only: pi
    implicit none
    type(multipole_center), intent(in) :: center
    type(box_iterator), intent(in) :: bit
    real(dp) :: val

    val = multipole_at_rxyz(center, bit%mesh%dom, bit%rxyz)
  end function multipole_at_bit

  pure function multipole_at_rxyz(center, dom, at) result(val)
    use at_domain
    use f_harmonics, only: solid_harmonic
    use numerics, only: pi
    implicit none
    type(multipole_center), intent(in) :: center
    type(domain), intent(in) :: dom
    real(dp), dimension(3), intent(in) :: at
    real(dp) :: val

    real(dp) :: r, tt
    real(dp), dimension(3, 3) :: qq
    real(dp), dimension(3) :: rxyz, Qr
    real(dp), parameter :: fac = sqrt(3._dp)
    real(dp), parameter :: inv = 1._dp / sqrt(3._dp)

    r = distance(dom, at, center%rxyz)
    rxyz = at - center%rxyz
    
    val = 0._dp
    if (has_monopole(center)) then
       val = val - center%mono%qlm / r
    end if
    if (has_dipole(center)) then
       tt =      center%di%qlm( 1) * rxyz(1)
       tt = tt + center%di%qlm(-1) * rxyz(2)
       tt = tt + center%di%qlm( 0) * rxyz(3)
       val = val - tt / r ** 3
    end if
    if (has_quadrupole(center)) then
       qq(1, 1) = inv * center%quad%qlm(2) - center%quad%qlm(0)
       qq(2, 1) = fac * center%quad%qlm(-2)
       qq(3, 1) = fac * center%quad%qlm(+1)
       qq(1, 2) = qq(2, 1)
       qq(2, 2) = -inv * center%quad%qlm(2) - center%quad%qlm(0)
       qq(3, 2) = fac * center%quad%qlm(-1)
       qq(1, 3) = qq(3, 1)
       qq(2, 3) = qq(3, 2)
       qq(3, 3) = 2._dp * center%quad%qlm(0)

       ! Calculate r^T*Q*r
       Qr(1) = qq(1, 1) * rxyz(1) + qq(2, 1) * rxyz(2) + qq(3, 1) * rxyz(3)
       Qr(2) = qq(1, 2) * rxyz(1) + qq(2, 2) * rxyz(2) + qq(3, 2) * rxyz(3)
       Qr(3) = qq(1, 3) * rxyz(1) + qq(2, 3) * rxyz(2) + qq(3, 3) * rxyz(3)
       tt = rxyz(1) * Qr(1) + rxyz(2) * Qr(2) + rxyz(3) * Qr(3)

       val = val - 0.5_dp * tt / r ** 5
    end if
  end function multipole_at_rxyz

  subroutine multipole_centers_to_gaussian_density(centers, bit, density)
    use box
    implicit none
    type(multipole_center), dimension(:), intent(in) :: centers
    type(box_iterator) :: bit
    real(dp), dimension(*), intent(inout) :: density

    integer :: ithread,nthread
    integer :: i
    type(subbox), dimension(:), allocatable :: nboxes
    type(subbox) :: nbox
    real(dp), dimension(3) :: dxyz
    real(dp) :: r
    !$ integer, external :: omp_get_thread_num,omp_get_num_threads

    nbox = subbox_null()
    allocate(nboxes(size(centers)))
    do i = 1, size(centers)
       nboxes(i) = multipole_to_gaussian_subbox(centers(i), bit)
       nbox = nbox .or. nboxes(i)
    end do

    call box_iter_set_nbox(bit, nbox = nbox%nbox)

    ithread=0
    nthread=1
    !$omp parallel default(shared) &
    !$omp private(ithread) &
    !$omp firstprivate(bit)
    !$ ithread = omp_get_thread_num()
    !$ nthread = omp_get_num_threads()
    call box_iter_split(bit, nthread, ithread)
    do while(box_next_point(bit))
       do i = 1, size(centers)
          if (box_iter_inside(bit, nboxes(i))) then
             dxyz = bit%rxyz - centers(i)%rxyz
             r = box_iter_distance(bit, centers(i)%rxyz)
             density(bit%ind) = density(bit%ind) + &
                  & multipole_gaussian_density_at(centers(i), dxyz, r)
          end if
       end do
    end do
    call box_iter_merge(bit)
    !$omp end parallel

    call box_iter_expand_nbox(bit)

    deallocate(nboxes)
  end subroutine multipole_centers_to_gaussian_density
end module multipoles
