module neighbors

    use parameters
    use elements
    use subroutines
    use functions
    use box

    integer, allocatable :: nei_list(:,:), nei_num(:), nn(:), periodvec(:,:,:)
    real(kind=dp), allocatable :: init_vec(:,:,:), output(:), microrotation(:,:)
    public 

    contains
 
    subroutine build_cell_list(numinlist, r_list, rc_off, cell_num, num_in_cell, cell_list, which_cell)
        !This subroutine builds a cell list based on rc_off
        
        !----------------------------------------Input/output variables-------------------------------------------

        integer, intent(in) :: numinlist !The number of points within r_list

        real(kind=dp), dimension(3,numinlist), intent(in) :: r_list !List of points to be used for the construction of 
                                                                    !the cell list.

        real(kind=dp), intent(in) :: rc_off ! Cutoff radius which dictates the size of the cells

        integer, dimension(3), intent(inout) :: cell_num !Number of cells in each dimension.

        integer, allocatable, intent(inout) :: num_in_cell(:,:,:) !Number of points within each cell

        integer, allocatable, intent(inout) :: cell_list(:,:,:,:) !Index of points from r_list within each cell.

        integer, dimension(3,numinlist), intent(out) :: which_cell !The cell index for each point in r_list

        !----------------------------------------Begin Subroutine -------------------------------------------

        integer :: i, j, cell_lim, c(3)
        real(kind=dp) :: box_len(3)
        integer, allocatable :: resize_cell_list(:,:,:,:)

        !First calculate the number of cells that we need in each dimension
        do i = 1,3
            box_len(i) = box_bd(2*i) - box_bd(2*i-1)
            cell_num(i) = int(box_len(i)/(rc_off/2))+1
        end do 

        !Initialize/allocate variables 
        cell_lim = 10
        allocate(num_in_cell(cell_num(1),cell_num(2),cell_num(3)), cell_list(cell_lim, cell_num(1), cell_num(2), cell_num(3)))

        !Now place points within cell
        num_in_cell = 0
        
        do i = 1, numinlist
            !Check to see if the current point is a filler point and if so just skip it
            if(r_list(1,i) < box_bd(1)) cycle

            !c is the position of the cell that the point belongs to 
            do j = 1, 3
                c(j) = int((r_list(j,i)-box_bd(2*j-1))/(rc_off)) + 1
            end do

            !Place the index in the correct position, growing if necessary
            num_in_cell(c(1),c(2),c(3)) = num_in_cell(c(1),c(2),c(3)) + 1
            if (num_in_cell(c(1),c(2),c(3)) > cell_lim) then 
                allocate(resize_cell_list(cell_lim+10,cell_num(1),cell_num(2),cell_num(3)))
                resize_cell_list(1:cell_lim, :, :, :) = cell_list(1:cell_lim, :, :, :)
                resize_cell_list(cell_lim+1:, :, :, :) = 0
                call move_alloc(resize_cell_list, cell_list)
                cell_lim = cell_lim + 10
            end if

            cell_list(num_in_cell(c(1),c(2),c(3)),c(1),c(2),c(3)) = i
            which_cell(:,i) = c
        end do 

        return
    end subroutine build_cell_list

    subroutine calc_neighbor(rc_off, r_list, n)
        !This function populates the neighbor list in this module
        
        real(kind=dp), intent(in) :: rc_off
        integer, intent(in) :: n
        real(kind=dp), dimension(:, :), intent(in) :: r_list

        integer :: i, j, c(3),cn(3), ci, cj, ck, num_nei, nei, v(3), period_dir(3)
        !Variables for cell list code
        integer, dimension(3) ::cell_num
        integer, allocatable :: num_in_cell(:,:,:), cell_list(:,:,:,:)
        integer :: which_cell(3,n)
        real(kind=dp) :: r(3), box_len(3)
        logical :: period_bd(3)
        
        !First reallocate the neighbor list variables
        if (allocated(nei_list)) then
            deallocate(nei_list,nei_num, periodvec)
        end if

        allocate(nei_list(100,n),nei_num(n), periodvec(3,100,n))
        nei_list=0
        periodvec=0
        nei_num=0

        !Now first pass the position list and and point num to the cell algorithm
        call build_cell_list(n, r_list, rc_off, cell_num, num_in_cell, cell_list, which_cell)

        do i=1, 3
            if (box_bc(i:i) == 'p') then 
                period_bd(i) = .true.
            else
                period_bd(i)=.false.
            end if
            box_len(i) = box_bd(2*i) - box_bd(2*i-1)
        end do
        !Now loop over every point and find it's neighbors
        pointloop: do i = 1, n
            
            !First check to see if the point is a filler point, if so then skip it 
            if(r_list(1,i) < box_bd(1)) cycle

            !c is the position of the cell that the point belongs to 
            c = which_cell(:,i)

            !Loop over all neighboring cells
            do ci = -1, 1, 1
                do cj = -1, 1, 1
                    do ck = -1, 1, 1
                        !First try to apply periodic boundaries 
                        v=(/ ck, cj, ci /)
                        cn=0
                        period_dir=0
                        do j=1, 3
                            if ((c(j) + v(j) == 0).and.period_bd(j)) then 
                                cn(j) = cell_num(j)
                                period_dir(j) = 1
                            else if ((c(j) + v(j) > cell_num(j)).and.period_bd(j)) then 
                                cn(j) = 1
                                period_dir(j) = -1

                            else if ((c(j)+v(j) >= 1) .and. (c(j)+v(j) <= cell_num(j))) then
                                cn(j) = c(j) + v(j)
                            end if
                        end do
                        if(any((c + (/ ck, cj, ci /)) == 0)) cycle

                        do num_nei = 1, num_in_cell(cn(1),cn(2),cn(3))
                            nei = cell_list(num_nei,cn(1),cn(2),cn(3))

                            !Check to make sure the atom isn't the same index as the atom we are checking 
                            !and that the neighbor hasn't already been deleted
                            if((nei /= i)) then 

                                !Now check to see if it is in the cutoff radius, if it is add it to the neighbor list for that
                                !atom and calculate the initial neighbor vector
                        
                                r = r_list(:,nei) + period_dir*box_len
                                if (norm2(r-r_list(:,i)) < rc_off) then 
                                    
                                    nei_num(i) = nei_num(i) + 1
                                    nei_list(nei_num(i), i) = nei
                                    periodvec(:,nei_num(i),i) = period_dir

                                end if
                            end if
                        end do 
                    end do
                end do
            end do

        end do pointloop

        return
    end subroutine calc_neighbor

    subroutine calc_NN(n,points, rc_off)

        integer, intent(in) :: n
        real(kind=dp), intent(in) :: points(3,n)
        real(kind=dp), intent(in) :: rc_off

        integer :: i, c(3), ci, cj, ck, nei
        !cell arrays
        integer, dimension(3) ::cell_num
        integer, allocatable :: num_in_cell(:,:,:), cell_list(:,:,:,:)
        integer :: which_cell(3,n)
        real(kind = dp) :: rmin

        !First reallocate the neighbor list codes
        if (allocated(nn)) then
            deallocate(nn)
        end if

        allocate(nn(n))
        nn=0

        !First build the cell lists
        call build_cell_list(n, points, rc_off, cell_num, num_in_cell, cell_list, which_cell)
        
        pointloop: do i = 1, n
            !First check to see if the point is a filler point, if so then skip it
            if(points(1,i) < -Huge(-1.0_dp)+1) cycle

            !Get the positon of the cell 
            c = which_cell(:,i)

            !Initialize the min vec
            rmin=rc_off

            !loop over all neighboring cells
            do ci = -1, 1, 1
                do cj = -1, 1, 1
                    do ck = -1, 1, 1
                        if(any((c + (/ ck, cj, ci /)) == 0)) cycle
                        if( (c(1) + ck > cell_num(1)).or.(c(2) + cj > cell_num(2)).or. (c(3) + ci > cell_num(3))) cycle

                        do num_nei = 1, num_in_cell(c(1) + ck, c(2) + cj, c(3) + ci)
                            nei = cell_list(num_nei,c(1) + ck, c(2) + cj, c(3) + ci)

                            !Check to make sure the atom isn't the same index as the atom we are checking 
                            if((nei /= i)) then 
                                !If it's the minimum position than we add it to the nearest neighbor list and updat e the min vec
                                
                                if (norm2(points(:,nei)-points(:,i)) < rmin) then 
                                    rmin = norm2(points(:, nei) - points(:,i))
                                    nn(i)=(nei)
                                end if
                            end if
                        end do 
                    end do
                end do
            end do
        end do pointloop

        return
    end subroutine calc_NN

end module neighbors
