! H0 XXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXX
! H0 X
! H0 X   libAtoms+QUIP: atomistic simulation library
! H0 X
! H0 X   Portions of this code were written by
! H0 X     Albert Bartok-Partay, Silvia Cereda, Gabor Csanyi, James Kermode,
! H0 X     Ivan Solt, Wojciech Szlachta, Csilla Varnai, Steven Winfield.
! H0 X
! H0 X   Copyright 2006-2010.
! H0 X
! H0 X   These portions of the source code are released under the GNU General
! H0 X   Public License, version 2, http://www.gnu.org/copyleft/gpl.html
! H0 X
! H0 X   If you would like to license the source code under different terms,
! H0 X   please contact Gabor Csanyi, gabor@csanyi.net
! H0 X
! H0 X   Portions of this code were written by Noam Bernstein as part of
! H0 X   his employment for the U.S. Government, and are not subject
! H0 X   to copyright in the USA.
! H0 X
! H0 X
! H0 X   When using this software, please cite the following reference:
! H0 X
! H0 X   http://www.libatoms.org
! H0 X
! H0 X  Additional contributions by
! H0 X    Alessio Comisso, Chiara Gattinoni, and Gianpietro Moras
! H0 X
! H0 XXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXX

module k_means_clustering_module
implicit none
private

   integer, parameter :: dp = 8

   public :: k_means_clustering_pick
   interface k_means_clustering_pick
      module procedure k_means_clustering_pick_1d, k_means_clustering_pick_nd
   end interface

contains

   subroutine k_means_clustering_pick_1d(x, periodicity, cluster_indices)
      real(dp), intent(in) :: x(:)
      real(dp), intent(in) :: periodicity
      integer, intent(out) :: cluster_indices(:)

      integer :: k
      real(dp), allocatable :: k_means(:)
      integer :: n, rv, i, j, closest_i, closest_j
      real(dp) :: rrv
      real(dp) :: closest_r, r, t_sum
      integer, allocatable :: a(:), prev_a(:)

      k = size(cluster_indices)
      n = size(x)
      allocate(k_means(k))

      ! random initial guess
      cluster_indices = 0
      do i=1, k
         call random_number(rrv); rv = 1+floor(rrv*n)
         do while (any(cluster_indices == rv))
            call random_number(rrv); rv = 1+floor(rrv*n)
         end do
         cluster_indices(i) = rv
      end do
      k_means = x(cluster_indices)

      allocate(a(n), prev_a(n))

      ! evaluate initial assignments
      do j=1, n
         closest_i = 0
         closest_r = HUGE(1.0_dp)
         do i=1, k
            r = dist((/x(j)/), (/k_means(i)/),  (/periodicity/))
            if (r < closest_r) then
               closest_r = r
               closest_i = i
            endif
         end do
         a(j) = closest_i
      end do

      prev_a = 0
      do while (any(a /= prev_a))
         prev_a = a

         ! update positions
         do i=1, k
            t_sum = 0.0_dp
            do j=1, n
               if (a(j) == i) t_sum = t_sum + x(j)
            end do
            k_means(i) = t_sum / real(count(a == i),dp)
         end do

         ! update assignments
         do j=1, n
            closest_i = 0
            closest_r = HUGE(1.0_dp)
            do i=1, k
               r = dist((/x(j)/), (/k_means(i)/),  (/periodicity/))
               if (r < closest_r) then
                  closest_r = r
                  closest_i = i
               endif
            end do
            a(j) = closest_i
         end do
      end do

      do i=1, k
         closest_i = 0
         closest_r = HUGE(1.0_dp)
         do j=1, n
            r = dist((/x(j)/), (/k_means(i)/),  (/periodicity/))
            if (r < closest_r) then
               closest_r = r
               closest_j = j
            endif
         end do
         cluster_indices(i) = closest_j
      end do

      deallocate(a, prev_a, k_means)

   end subroutine k_means_clustering_pick_1d

   subroutine k_means_clustering_pick_nd(x, periodicity, cluster_indices)
      real(dp), intent(in) :: x(:,:)
      real(dp), intent(in) :: periodicity(:)
      integer, intent(out) :: cluster_indices(:)

      integer :: k
      real(dp), allocatable :: k_means(:,:)
      integer :: n, rv, i, j, closest_i, closest_j, nd
      real(dp) :: rrv
      real(dp) :: closest_r, r, t_sum(size(x,1))
      integer, allocatable :: a(:), prev_a(:)

      k = size(cluster_indices)
      n = size(x,2)
      nd = size(x,1)
      allocate(k_means(nd,k))

      ! random initial guess
      cluster_indices = 0
      do i=1, k
         call random_number(rrv); rv = 1+floor(rrv*n)
         do while (any(cluster_indices == rv))
            call random_number(rrv); rv = 1+floor(rrv*n)
         end do
         cluster_indices(i) = rv
      end do
      k_means(:,:) = x(:,cluster_indices(:))

      allocate(a(n), prev_a(n))

      ! evaluate initial assignments
      do j=1, n
         closest_i = 0
         closest_r = HUGE(1.0_dp)
         do i=1, k
            r = dist(x(:,j), k_means(:,i), periodicity(:))
            if (r < closest_r) then
               closest_r = r
               closest_i = i
            endif
         end do
         a(j) = closest_i
      end do

      prev_a = 0
      do while (any(a /= prev_a))
         prev_a = a

         ! update positions
         do i=1, k
            t_sum = 0.0_dp
            do j=1, n
               if (a(j) == i) t_sum(:) = t_sum(:) + x(:,j)
            end do
            k_means(:,i) = t_sum(:) / real(count(a == i),dp)
         end do

         ! update assignments
         do j=1, n
            closest_i = 0
            closest_r = HUGE(1.0_dp)
            do i=1, k
               r = dist(x(:,j), k_means(:,i), periodicity(:))
               if (r < closest_r) then
                  closest_r = r
                  closest_i = i
               endif
            end do
            a(j) = closest_i
         end do
      end do

      do i=1, k
         closest_i = 0
         closest_r = HUGE(1.0_dp)
         do j=1, n
            r = dist(x(:,j), k_means(:,i), periodicity(:))
            if (r < closest_r) then
               closest_r = r
               closest_j = j
            endif
         end do
         cluster_indices(i) = closest_j
      end do

      deallocate(a, prev_a, k_means)

   end subroutine k_means_clustering_pick_nd

   ! from BRUSH paper, citing MacKay
   function dist(x0, x1, periodicity)
      real(dp), intent(in) :: x0(:), x1(:)
      real(dp), intent(in) :: periodicity(:)
      real(dp) :: dist ! result

      integer :: i

      if (any(periodicity /= 0.0_dp)) then
         dist = 0.0_dp
         do i=1, size(x0)
            if (periodicity(i) /= 0.0_dp) then
               dist = dist + 4.0_dp*sin((2.0_dp*3.14159265358979_dp/periodicity(i))*(x0(i)-x1(i))/2.0_dp)**2
            else
               dist = dist + (x0(i)-x1(i))**2
            endif
         end do
         dist = sqrt(dist)
      else
         dist = sqrt(sum((x0(:)-x1(:))**2))
      endif

   end function dist

end module k_means_clustering_module
