!  Copyright (C) 2017 The authors of Py-ChemShell
!
!  This file is part of Py-ChemShell.
!
!  Py-ChemShell is free software: you can redistribute it and/or modify
!  it under the terms of the GNU Lesser General Public License as
!  published by the Free Software Foundation, either version 3 of the
!  License, or (at your option) any later version.
!
!  Py-ChemShell is distributed in the hope that it will be useful,
!  but WITHOUT ANY WARRANTY; without even the implied warranty of
!  MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
!  GNU Lesser General Public License for more details.
!
!  You should have received a copy of the GNU Lesser General Public
!  License along with Py-ChemShell.  If not, see
!  <http://www.gnu.org/licenses/>.

subroutine cluster_adjust_charge(nccore, ncbq, ncshell, &
                                 npcore, npbq, npshell, &
                                 ierr)
  !
  ! Adjust charges at the cluster surface taking coordination into account
  !

  use ChemShellClusterModule, only : PyCell, PyOptions, debug

  implicit none

  integer, intent(in) :: nccore, ncbq, ncshell
  integer, intent(in) :: npcore, npbq, npshell
  integer, intent(out) :: ierr

  integer :: i, j, k, l, m, n, p, ci
  double precision :: tot_charge

  ! Size of connectivity table entries in fragment
  integer, parameter :: max_coord = 14

  ! Arrays for crystal information
  double precision, dimension(npcore+npshell) :: crystal_charges
  double precision, dimension(npcore) :: charge_core
  double precision, dimension(npshell) :: charge_shel
  integer, dimension(npcore) :: znum_core
  integer, dimension(npshell) :: point_core
  integer, dimension(npcore) :: point_shel  
  integer, dimension(npcore) :: coordination
  integer, dimension(max_coord, npcore) :: point_neighbour
  double precision, dimension(3, max_coord, npcore) :: v_neighbour

  ! Arrays for cluster information
  double precision, dimension(nccore+ncshell) :: cluster_charges
  double precision, dimension(nccore) :: charge_cluster_core
  double precision, dimension(ncshell) :: charge_cluster_shel  
  integer, dimension(ncshell) :: point_cluster_core
  integer, dimension(nccore) :: point_cluster_shel
  integer, dimension(nccore+ncshell) :: xtal_map
  integer, dimension(nccore) :: point_uc_c
  integer, dimension(ncshell) :: point_uc_s
  integer, dimension(nccore) :: coordination_in_cluster
  integer, dimension(max_coord, nccore) :: point_cluster_neighbour

  ! Species table/valence model
  integer :: no_species
  integer, dimension(npcore) :: species_core
  double precision, dimension (npcore) :: species_charge
  integer, dimension(npcore) :: species_coord
  integer, dimension(npcore) :: species_znum
  double precision, dimension (max_coord, npcore) :: species_neighbour_charge
  integer, dimension(max_coord, npcore) :: species_neighbour_coord  
  integer, dimension(max_coord) :: coord_neighbour
  double precision, dimension(max_coord) :: q_neighbour
  double precision, parameter :: tol_q = 1.0d-6
  logical :: new_species, neighbour_charge_ok
  integer, dimension(npcore) :: core_species
  integer, dimension(npcore, npcore) :: species_species
  double precision, dimension(npcore, npcore) :: valence

  double precision, pointer   :: twodimdblbuff(:,:)
  
  logical :: shell_model

  integer, parameter :: periodic_frag = 1, cluster_frag = 2
  double precision, parameter :: d_min = 0.1d0

  character(len=:), pointer :: adjust_mode
  integer :: ndim

  ierr = 0

  call PyOptions%get('adjust_charge', adjust_mode)
  call PyCell%get('ndimensions', ndim)
  
  if (debug > 0) then
     write(6,'(a)')    "--------------------------------------"
     write(6,'(a,i4)') "cluster_adjust_charge debug level:", debug
     write(6,'(a)')    "--------------------------------------"
  end if

  ! no action required
  if (trim(adjust_mode).eq.'') then
     write(6,*)
     write(6,*) " >>> ChemShell CLUSTER: charge adjustment not requested"
     return
  end if
  ! Check there are no point charges
  if (ncbq .ne. 0 .or. npbq .ne. 0) then
     write(6,*) " >>> ChemShell CLUSTER ERROR: number of point charges must be zero for charge adjustment"
     ierr = 1
     return
  end if
  ! valence model is for bulk systems only
  if (ndim /= 3 .and.trim(adjust_mode).eq.'valence_model') then
     write(6,*) " >>> ChemShell CLUSTER ERROR: valence model scheme only valid for bulk systems"
     ierr = 1
     return     
  end if

  if (npshell .eq. 0) then
     shell_model = .false.
  else
     shell_model = .true.
  end if

  ! Get crystal fragment information

  ! Get charges of cartesian sites from crystal fragment
  ! Split into core and shell arrays
!  call get_cart_charges(periodic_frag, crystal_charges, ierr)
  call get_cart_charges(periodic_frag, crystal_charges, npcore+npshell,ierr)
  if (ierr .ne. 0) return
  charge_core = crystal_charges(1:npcore)
  if (shell_model) charge_shel = crystal_charges(npcore+1:npcore+npshell)

  ! Get pointers from shells to cores
  ! Then setup pointers in the other direction
  if (shell_model) then
     call get_point_core(periodic_frag, point_core, npshell, ierr)
     if (ierr .ne. 0) return
     point_shel = 0
     do i = 1, npshell
        point_shel(point_core(i)) = i
     end do
  end if

  ! Determine coordination number and neighbours for all atoms in the unit cell
  coordination(:) = 0
  call periodic_neighbour_list(npcore, coordination, &
       point_neighbour, v_neighbour, ierr)
  if (ierr .ne. 0) return
  if (debug > 0) then
     write(6,'(/a)') "Crystal neighbour list"
     do i = 1, npcore
        write(6,'(i4,a,i4)') i, ":", coordination(i)
        do j = 1, coordination(i)
           write(6,'(2x,i4,2x,3f12.7)') point_neighbour(j,i), v_neighbour(1:3,j,i)
        end do
        write(6,*)
     end do
     write(6,*)
  end if

  ! Determine species table consisting of groups of atoms which 
  ! share the same atomic number, charge, coordination and neighbour charges
  if(trim(adjust_mode).eq.'valence_model'.or.trim(adjust_mode).eq.'print_species') then
     ! Initialisation
     call get_znum(periodic_frag, znum_core, npcore, ierr)
     no_species = 0
     ! Build species table
     do i = 1, npcore
        ! Sort charges of neighbouring ions by value
        ci = coordination(i)
        do j = 1, ci
           k = point_neighbour(j, i)
           coord_neighbour(j) = coordination(k)
           q_neighbour(j) = charge_core(k)
           if (shell_model .and. point_shel(k) /= 0) then
              q_neighbour(j) = q_neighbour(j) + charge_shel(point_shel(k))
           end if
        end do
        call sortvec(ci, q_neighbour, coord_neighbour)
        ! Check if the species exists and add to the list of species if not
        new_species = .true.
        tot_charge = charge_core(i)
        if (shell_model .and. point_shel(i) /= 0) then
           tot_charge = tot_charge + charge_shel(point_shel(i))
        end if
        do j = 1, no_species
           if (dabs(tot_charge - species_charge(j)) > tol_q) cycle
           if (znum_core(i) /= species_znum(j)) cycle
           if (ci /= species_coord(j)) cycle
           neighbour_charge_ok = .true.
           do k = 1, ci
              if (dabs(q_neighbour(k) - species_neighbour_charge(k, j)) > tol_q) then
                 neighbour_charge_ok = .false.
                 exit
              end if
           end do
           if (.not. neighbour_charge_ok) cycle
           ! Species already exists
           species_core(i) = j
           new_species = .false.
           exit
        end do
        if (new_species) then
           no_species = no_species + 1
           species_core(i) = no_species 
           species_charge(no_species) = tot_charge
           species_coord(no_species) = ci
           species_znum(no_species) = znum_core(i)
           do j = 1, ci
              species_neighbour_charge(j, no_species) = q_neighbour(j)
              species_neighbour_coord(j, no_species) = coord_neighbour(j)
           end do
        end if
     end do ! loop over cores
     ! Determine coordination between species to help the user
     ! construct the valence matrix
     do i = 1, no_species
        do j = 1, npcore
           if (species_core(j) == i) then
              core_species(i) = j
              exit
           end if
        end do
     end do
     species_species(:,:) = 0
     do i = 1, no_species
        j = core_species(i)
        l = species_coord(i)
        do k = 1, l
           p = species_core(point_neighbour(k,j))
           species_species(i,p) = species_species(i,p) + 1
        end do
     end do
     ! Print out species information
     write(6,'(a)') 'Crystal species list'
     write(6,'(a,4x,a,6x,a,8x,a)') &
          '    No    Z ','q','Coord','Charges on neighbours'
     do i = 1, no_species
        write(6, '(/3x,i3,2x,i3,f10.4,2x,i3,2x,5f10.4/(28x,5f10.4))') &
             i, species_znum(i), species_charge(i), &
             species_coord(i), &
             (species_neighbour_charge(j,i), j=1,species_coord(i))
     enddo
     write(6,'(/a)') 'Table of species coordination'
     do i = 1, no_species
        write(6,'(3x,10i5)') (species_species(i,j), j=1, no_species)
     end do
  end if

  ! If called by periodic_species, no further calculations required
  if(trim(adjust_mode).eq.'print_species') return

  ! Get cluster fragment information

  ! Get pointers from shells to cores
  ! Then setup pointers in the other direction
  if (shell_model) then
     call get_point_core(cluster_frag, point_cluster_core, ncshell, ierr)
     if (ierr .ne. 0) return
     point_cluster_shel = 0
     do i = 1, ncshell
        point_cluster_shel(point_cluster_core(i)) = i
     end do
  end if

  ! Get coordination of atoms in cluster
  call get_cluster_connectivity(nccore, point_cluster_neighbour, ierr)
  if (ierr .ne. 0) return
  do i = 1, nccore
     coordination_in_cluster(i) = 0
     do j = 1, max_coord
        if (point_cluster_neighbour(j,i) /= 0) then
           coordination_in_cluster(i) = coordination_in_cluster(i) + 1
        end if
     end do
  end do
  if (debug > 0) then
     write(6,'(/a)') "Coordination in cluster"
     do i = 1, nccore
        write(6,'(i4,a,i2,a)', advance='no') i, ": (", coordination_in_cluster(i), ")  "
        do j = 1, max_coord
           if (point_cluster_neighbour(j,i) /= 0) then
              write(6,'(i4)', advance='no') point_cluster_neighbour(j,i)
           end if
        end do
        write(6,*)
     end do
     write(6,*)
  end if
  
  ! Get pointer from cores/shells to the unit cell (xtal map) 
  call get_xtal_map(cluster_frag, xtal_map, nccore, ncshell, ierr)
  if (ierr .ne. 0) return
  point_uc_c = xtal_map(1:nccore)
  if (shell_model) point_uc_s = xtal_map(nccore+1:nccore+ncshell)
  if (debug > 0) then
     write(6,'(/a)') "XTAL MAP"
     write(6,'(/a)') "CORES"
     do i = 1, nccore
        write(6,'(2i6)') i, point_uc_c(i)
     end do
     if (shell_model) then
        write(6,'(/a)') "SHELLS"
        do i = 1, ncshell
           write(6,'(2i6)') i, point_uc_s(i)
        end do
     end if
  end if  

  ! Set cluster charges
  charge_cluster_core(1:nccore) = 0
  charge_cluster_shel(1:ncshell) = 0
  select case(trim(adjust_mode))
  case ('coordination_scaled')
     ! This is the standard scheme in the surface subroutine of construct
     ! But it seems equally applicable to the bulk case
     do i = 1, nccore
        j = point_uc_c(i)
        if (coordination(j) /= 0) then
           charge_cluster_core(i) = charge_core(j) * &
                coordination_in_cluster(i) / coordination(j)
        else
           charge_cluster_core(i) = charge_core(j)
        end if
     end do
     if (shell_model) then
        do i = 1, ncshell
           j = point_uc_s(i)
           k = point_cluster_core(i)
           l = point_uc_c(k)
           if (coordination(l) /= 0) then
              charge_cluster_shel(i) = charge_shel(j) * &
                   coordination_in_cluster(k) / coordination(l)
           else
              charge_cluster_shel(i) = charge_shel(j)
           end if
        end do
     end if
  case('valence_model')
     ! The valence model scheme is only applicable to the bulk case

!     call get_valence_matrix(no_species, valence(1:no_species, 1:no_species), ierr)
     allocate(twodimdblbuff(no_species,no_species))
     call PyOptions%get('valence_matrix', twodimdblbuff)

     valence(1:no_species,1:no_species) = twodimdblbuff
         
     if (ierr .ne. 0) return
     do i = 1, nccore
        tot_charge = 0.0d0
        j = point_uc_c(i)
        m = species_core(j)
        ! Cations are fully coordinated in the valence model, while
        ! anions should have their charge adjusted according to their
        ! coordination
        if (species_charge(m) < 0.0d0) then
           ! Loop over neighbours, adding charge contributions
           ! according to the valence matrix
           do l = 1, max_coord
              p = point_cluster_neighbour(l,i)
              if (p /= 0) then
                 n = species_core(point_uc_c(p))
                 ! In construct, valence.dat is input as a column-major list
                 ! But in chemshell, matrices are input as row-major lists
                 ! To avoid the user having to transpose the input, 
                 ! we instead transpose the internal valence matrrix
                 !  - so rows are core i, and columns are neighbouring cores
                 tot_charge = tot_charge + valence(m,n)
              end if
           end do
           charge_cluster_core(i) = charge_core(j) * &
                tot_charge / species_charge(m)
        else
           charge_cluster_core(i) = charge_core(j)
        end if
        if (shell_model) then
           k = point_cluster_shel(i)
           if (k /= 0) then
              l = point_uc_s(k)
              charge_cluster_shel(k) = charge_shel(l) * &
                   charge_cluster_core(i) / charge_core(j)
           end if
        end if
     end do
  case default
     write(6,*) " >>> cluster_adjust_charge ERROR: unknown adjust_charge ", trim(adjust_mode)
     ierr = 1
     return 
  end select
  if (debug > 3) then
     write(6,'(/a)') "Adjusted cluster charges"
     write(6,'(/a)') "CORES"
     do i = 1, nccore
        write(6,'(i6, 3f12.7)') i, charge_cluster_core(i)
     end do
     write(6,'(/a)') "SHELLS"
     do i = 1, ncshell
        write(6,'(i6, 3f12.7)') i, charge_cluster_shel(i)
     end do
     write(6,*)    
  end if
  
  ! Write the new charges to the cluster fragment
  cluster_charges(1:nccore) = charge_cluster_core
  if (shell_model) cluster_charges(nccore+1:nccore+ncshell) = charge_cluster_shel

  call set_cart_charges(cluster_frag, cluster_charges, nccore+ncshell, ierr)
  if (ierr .ne. 0) return

end subroutine cluster_adjust_charge


subroutine sortvec(ne,q,l)
  implicit none
  ! Sort vector q in ascending order
  integer :: ne,m,n,ium,l(ne)
  double precision :: q(ne), dum
  do m=2,ne
     dum=q(1)
     ium=l(1)
     do n=2,ne-m+2
        if(dum.gt.q(n)) then
           q(n-1)=q(n)
           q(n)=dum
           l(n-1)=l(n)
           l(n)=ium
        else
           dum=q(n)
           ium=l(n)
        endif
     enddo
  enddo
  return
end subroutine sortvec
