program test_multipole
  use f_precisions, only: dp => f_double
  use multipoles
  use dictionaries
  use yaml_output
  implicit none

  call f_lib_initialize()

  call yaml_sequence_open("Tests")

  call run("monopole", test_monopole)
  call run("dipole", test_dipole)
  call run("quadrupole", test_quadrupole)
  call run("setter", test_set)
  call run("conversions", test_conversions)
  call run("monopole_at", test_monopole_at)
  call run("monopole_gausian", test_monopole_gaussian_at)
  call run("dipole_at", test_dipole_at)
  call run("dipole_gausian", test_dipole_gaussian_at)
  call run("quadpole_at", test_quadpole_at)
  call run("quadpole_gausian", test_quadpole_gaussian_at)
  call run("gaussian_nbox", test_gaussian_nbox)
  
  call yaml_sequence_close()
  
  call f_lib_finalize()

contains
  subroutine run(test, func)
    implicit none
    character(len = *), intent(in) :: test
    interface
       subroutine func()
       end subroutine func
    end interface

    type(dictionary), pointer ::errors

    nullify(errors)
    call f_err_open_try()
    call func()
    call f_err_close_try(errors)
    call yaml_sequence(advance = "no")
    if (associated(errors)) then
       call yaml_mapping_open(test)
       call yaml_map("status", "failed")
       call yaml_map("error", errors)
       call yaml_mapping_close()
       call dict_free(errors)
    else
       call yaml_map(test, "succeed")
    end if
  end subroutine run

  subroutine test_monopole()
    type(monopole) :: mono

    mono = pole(1._dp, 2._dp)
    if (mono%l /= 0) call f_err_throw("Wrong monopole l value.")
    if (mono%qlm /= 1._dp) call f_err_throw("Wrong monopole qlm value.")
    if (mono%sigma /= 2._dp) call f_err_throw("Wrong monopole sigma value.")
  end subroutine test_monopole
  
  subroutine test_dipole()
    type(dipole) :: di

    di = pole(1._dp, 2._dp, 3._dp, 4._dp)
    if (di%l /= 1) call f_err_throw("Wrong dipole l value.")
    if (any(di%qlm /= (/ 1._dp, 2._dp, 3._dp /))) &
         & call f_err_throw("Wrong dipole qlm value.")
    if (di%sigma /= 4._dp) call f_err_throw("Wrong dipole sigma value.")
  end subroutine test_dipole
  
  subroutine test_quadrupole()
    type(quadrupole) :: quad

    quad = pole(1._dp, 2._dp, 3._dp, 4._dp, 5._dp, 6._dp)
    if (quad%l /= 2) call f_err_throw("Wrong quadrupole l value.")
    if (any(quad%qlm /= (/ 1._dp, 2._dp, 3._dp, 4._dp, 5._dp /))) &
         & call f_err_throw("Wrong quadrupole qlm value.")
    if (quad%sigma /= 6._dp) call f_err_throw("Wrong quadrupole sigma value.")
  end subroutine test_quadrupole
  
  subroutine test_set()
    use f_enums
    type(multipole_center) :: center
    type(monopole) :: mono
    type(dipole) :: di
    type(quadrupole) :: quad
    real(dp), dimension(3), parameter :: rxyz = (/ 1._dp, 2._dp, 3._dp /)

    mono = pole(1._dp, 2._dp)
    di = pole(1._dp, 2._dp, 3._dp, 4._dp)
    quad = pole(1._dp, 2._dp, 3._dp, 4._dp, 5._dp, 6._dp)
    call multipole_center_set(center, rxyz, "H", mono, di, quad, CHARACTER_NET_ENUM)

    if (any(center%rxyz /= rxyz)) call f_err_throw("Wrong rxyz value.")
    if (center%label /= "H") call f_err_throw("Wrong label value.")
    if (center%character /= CHARACTER_NET_ENUM) call f_err_throw("Wrong character value.")
    if (center%mono /= mono) call f_err_throw("Wrong monopole value.")
    if (center%di /= di) call f_err_throw("Wrong dipole value.")
    if (center%quad /= quad) call f_err_throw("Wrong quadrupole value.")
  end subroutine test_set
  
  subroutine test_conversions()
    use f_utils
    use dictionaries
    use yaml_output
    use yaml_parse

    type(multipole_center), dimension(3) :: H2O
    type(multipole_center), dimension(:), pointer :: centers
    type(dictionary), pointer :: dict
    integer :: iunit

    call multipole_center_set(H2O(1), &
         & (/ 100.86_dp, 97.39_dp, 101.89_dp /), "H", &
         & pole(0.30_dp, 0.2_dp), &
         & pole(0.013_dp, -0.018_dp, 0.058_dp, 0.2_dp), &
         & pole(0.038_dp, -0.003_dp, -0.045_dp, -0.038_dp, 0.071_dp, 0.2_dp), &
         & CHARACTER_NET_ENUM)
    
    call multipole_center_set(H2O(2), &
         & (/ 101.99_dp, 97.13_dp, 103.04_dp /), "H", &
         & pole(0.30_dp, 0.2_dp), &
         & pole(0.025_dp, -0.051_dp, 0.023_dp, 0.2_dp), &
         & pole(0.014_dp, -0.061_dp, 0.080_dp, -0.035_dp, 0.006_dp, 0.2_dp), &
         & CHARACTER_NET_ENUM)
    
    call multipole_center_set(H2O(3), &
         & (/ 101.83_dp, 97.44_dp, 102.11_dp /), "O", &
         & pole(-0.60_dp, 0.245_dp), &
         & pole(-0.015_dp, 0.088_dp, -0.109_dp, 0.245_dp), &
         & pole(-0.137_dp, -0.154_dp, 0.232_dp, 0.495_dp, 0.401_dp, 0.245_dp), &
         & CHARACTER_NET_ENUM)

    nullify(dict)
    call multipole_centers_to_dict(dict, H2O, "external_potential", "angstroem")
    iunit = f_get_free_unit(9)
    call yaml_set_stream(unit = iunit, filename = "water.yaml", &
         & record_length = 92, tabbing = 0, setdefault = .false.)
    call yaml_new_document(unit = iunit)
    call yaml_map("dft", dict, unit = iunit)
    call yaml_close_stream(iunit)
    call dict_free(dict)

    nullify(centers)
    nullify(dict)
    call yaml_parse_from_file(dict, "water.yaml")
    call multipole_centers_from_dict(centers, dict // 0 // "dft" // "external_potential")
    call dict_free(dict)

    if (any(centers /= H2O)) call f_err_throw("Differences between after re-read.")

    deallocate(centers)
  end subroutine test_conversions

  subroutine compare(a, b)
    use yaml_strings
    implicit none
    real(dp), intent(in) :: a, b
    real(dp), parameter :: eps = 1.d-12
    if (abs(a - b) > eps) call f_err_throw("expected:" // yaml_toa(b) // ", actual:" // yaml_toa(a))
  end subroutine compare

  function test_iter()
    use at_domain
    use box
    implicit none
    type(domain), save :: dom
    type(cell), save :: mesh
    type(box_iterator) :: test_iter
    integer, dimension(2, 3) :: nbox

    dom = domain_new(ATOMIC_UNITS, (/ FREE_BC, FREE_BC, FREE_BC /), &
         & acell = (/ 4._dp, 5._dp, 6._dp /))
    mesh = cell_new(dom, (/ 4, 5, 6 /), (/ 1._dp, 1._dp, 1._dp /))

    nbox(:, 1) = (/ 0, 2 /)
    nbox(:, 2) = (/ 0, 2 /)
    nbox(:, 3) = (/ 0, 2 /)
    test_iter = box_iter(mesh, nbox = nbox)
  end function test_iter

  subroutine test_monopole_at()
    use box
    implicit none
    type(multipole_center) :: center
    type(box_iterator) :: iter

    integer :: i
    real(dp), dimension(8), parameter :: expected = &
         & (/ -0.53452248382484879_dp, -0.66666666666666663_dp, &
         &    -0.60302268915552726_dp, -0.81649658092772615_dp, &
         &    -0.55470019622522915_dp, -0.70710678118654746_dp, &
         &    -0.63245553203367588_dp, -0.89442719099991586_dp /)

    center%rxyz = (/ 3._dp, 2._dp, 1._dp /)
    center%mono%qlm = 2._dp

    iter = test_iter()
    i = 0
    do while(box_next_point(iter))
       i = i + 1
       call compare(multipole_at(center, iter), expected(i))
       call compare(multipole_at(center, iter%mesh%dom, iter%rxyz), expected(i))
    end do

  end subroutine test_monopole_at

  subroutine test_monopole_gaussian_at()
    use box
    implicit none
    type(multipole_center) :: center
    type(box_iterator) :: iter

    integer :: i
    real(dp), dimension(8), parameter :: expected = &
         & (/ -2.1607863711204283E-3_dp, -2.8526545844999216E-3_dp, &
         &    -2.5526674795012896E-3_dp, -3.3700131977079050E-3_dp, &
         &    -2.2842272234111934E-3_dp, -3.0156203074923516E-3_dp, &
         &    -2.6984956157279950E-3_dp, -3.5625344515051905E-3_dp /)

    center%rxyz = (/ 3._dp, 2._dp, 1._dp /)
    center%mono%qlm = 2._dp
    center%mono%sigma = 3._dp

    iter = test_iter()
    i = 0
    do while(box_next_point(iter))
       i = i + 1
       call compare(multipole_gaussian_density_at(center, iter), expected(i))
       call compare(multipole_gaussian_density_at(center, iter%mesh%dom, iter%rxyz), expected(i))
    end do

  end subroutine test_monopole_gaussian_at

  subroutine test_dipole_at()
    use box
    implicit none
    type(multipole_center) :: center
    type(box_iterator) :: iter

    integer :: i
    real(dp), dimension(8), parameter :: expected = &
         & (/ 0.2481711532043941_dp, 0.3703703703703703_dp, &
         &    0.3289214668121058_dp, 0.6123724356957947_dp, &
         &    0.2346808522491354_dp, 0.3535533905932737_dp, &
         &    0.3162277660168379_dp, 0.6260990336999410_dp /)

    center%rxyz = (/ 3._dp, 2._dp, 1._dp /)
    center%di%qlm = (/ 1._dp, 2._dp, 3._dp /)

    iter = test_iter()
    i = 0
    do while(box_next_point(iter))
       i = i + 1
       call compare(multipole_at(center, iter), expected(i))
       call compare(multipole_at(center, iter%mesh%dom, iter%rxyz), expected(i))
    end do

  end subroutine test_dipole_at

  subroutine test_dipole_gaussian_at()
    use box
    implicit none
    type(multipole_center) :: center
    type(box_iterator) :: iter

    integer :: i
    real(dp), dimension(8), parameter :: expected = &
         & (/ 4.2135334236848357E-2_dp, 4.2789818767498829E-2_dp, &
         &    4.5948014631023229E-2_dp, 4.5495178169056724E-2_dp, &
         &    3.7689749186284698E-2_dp, 3.6187443689908221E-2_dp, &
         &    4.0477434235919943E-2_dp, 3.7406611740804507E-2_dp /)

    center%rxyz = (/ 3._dp, 2._dp, 1._dp /)
    center%di%qlm = (/ 1._dp, 2._dp, 3._dp /)
    center%di%sigma = 3._dp

    iter = test_iter()
    i = 0
    do while(box_next_point(iter))
       i = i + 1
       call compare(multipole_gaussian_density_at(center, iter), expected(i))
       call compare(multipole_gaussian_density_at(center, iter%mesh%dom, iter%rxyz), expected(i))
    end do

  end subroutine test_dipole_gaussian_at

  subroutine test_quadpole_at()
    use box
    implicit none
    type(multipole_center) :: center
    type(box_iterator) :: iter

    integer :: i
    real(dp), dimension(8), parameter :: expected = &
         & (/ -3.9301033063618429E-2_dp, -7.7007460580666806E-2_dp, &
         &    -7.2242788743151917E-2_dp, -0.2337758616699322_dp, &
         &     3.1030664032001567E-3_dp,  2.8017983505251670E-2_dp, &
         &    -5.5123489896403690E-3_dp, -5.2633218134796279E-3_dp /)

    center%rxyz = (/ 3._dp, 2._dp, 1._dp /)
    center%quad%qlm = (/ 1._dp, 2._dp, 3._dp, 4._dp, 5._dp /)

    iter = test_iter()
    i = 0
    do while(box_next_point(iter))
       i = i + 1
       call compare(multipole_at(center, iter), expected(i))
       call compare(multipole_at(center, iter%mesh%dom, iter%rxyz), expected(i))
    end do

  end subroutine test_quadpole_at

  subroutine test_quadpole_gaussian_at()
    use box
    implicit none
    type(multipole_center) :: center
    type(box_iterator) :: iter

    integer :: i
    real(dp), dimension(8), parameter :: expected = &
         & (/ -0.2336661031371878_dp, -0.1334529789206778_dp, &
         &    -0.3323948292775257_dp, -0.2466427629554140_dp, &
         &    -7.1627312181537270E-2_dp, 3.8236533335638215E-2_dp, &
         &    -0.1675574762606583_dp, -7.9751632535163611E-2_dp /)

    center%rxyz = (/ 3._dp, 2._dp, 1._dp /)
    center%quad%qlm = (/ 1._dp, 2._dp, 3._dp, 4._dp, 5._dp /)
    center%quad%sigma = 3._dp

    iter = test_iter()
    i = 0
    do while(box_next_point(iter))
       i = i + 1
       call compare(multipole_gaussian_density_at(center, iter), expected(i))
       call compare(multipole_gaussian_density_at(center, iter%mesh%dom, iter%rxyz), expected(i))
    end do

  end subroutine test_quadpole_gaussian_at

  subroutine test_gaussian_nbox()
    use box
    implicit none
    type(multipole_center) :: center
    type(box_iterator) :: iter

    type(subbox) :: expected
    
    ! The cutoff will be 30., hgrids = 1., so 30 indices in each directions
    center%rxyz = (/ 3._dp, 2._dp, 1._dp /)
    expected%nbox(:, 1) = (/ -27, 33 /)
    expected%nbox(:, 2) = (/ -28, 32 /)
    expected%nbox(:, 3) = (/ -29, 31 /)

    iter = test_iter()

    center%mono%sigma = 3._dp
    center%di%sigma = 2._dp
    center%quad%sigma = 1._dp
    if (multipole_to_gaussian_subbox(center, iter) /= expected) call f_err_throw("Wrong nbox from mono")

    center%mono%sigma = 1._dp
    center%di%sigma = 3._dp
    center%quad%sigma = 2._dp
    if (multipole_to_gaussian_subbox(center, iter) /= expected) call f_err_throw("Wrong nbox from di")

    center%mono%sigma = 2._dp
    center%di%sigma = 1._dp
    center%quad%sigma = 3._dp
    if (multipole_to_gaussian_subbox(center, iter) /= expected) call f_err_throw("Wrong nbox from quad")

  end subroutine test_gaussian_nbox
end program test_multipole
