module mode_merge
    !This module contains the code needed for merging several .mb files together

    use parameters
    use atoms   
    use io
    use subroutines
    use elements
    use neighbors

    implicit none

    character(len=4) :: dim
    integer :: in_num, new_starts(2)
    real(kind=dp) :: shift_vec(3), replace_vec(3) 
    character(len=100) :: replace_str(3)
    logical :: shift_flag, replace_flag
    real(kind=dp), private, save :: rc_off 

    public 
    contains
    subroutine merge(arg_pos)
        
        integer, intent(out) :: arg_pos 
        integer :: i, j 
        real(kind=dp) :: displace(3), temp_box_bd(6)

        print *, '-----------------------Mode Merge---------------------------'
        
        shift_flag = .false.

        shift_vec(:) = 0.0_dp
        temp_box_bd(:) = 0.0_dp
        !First we parse the merge command
        call parse_command(arg_pos)

        !Now loop over all files and stack them 
        do i = 1, in_num
            displace(:) = 0.0_dp
            
            !The new starts variable dictate where in the atom and element array each new 
            !file starts. This is used for additional options that can be applied to solely 
            !these new atoms/elements that are read in.

            new_starts(1) = atom_num + 1
            new_starts(2) = ele_num + 1

            if ((i==1).or.(trim(adjustl(dim)) == 'none')) then 
                call read_in(i, displace, temp_box_bd)
            else
                select case(trim(adjustl(dim)))
                case('x')
                    displace(1) = box_bd(2)
                case('y')
                    displace(2) = box_bd(4)
                case('z')
                    displace(3) = box_bd(6)
                end select

                call read_in(i, displace, temp_box_bd)
            end if 


            if(shift_flag) call shift(new_starts, i)

            if(replace_flag.and.(i>1)) then 
                !Parse the replace vector
                do j = 1, 3
                    call parse_pos(j, replace_str(j), replace_vec(j))
                end do
                call replace(new_starts, temp_box_bd)
            end if

        end do


        !Now reset tags
        do i = 1, atom_num
            tag_atom(i) = i
        end do
        
        do i = 1, ele_num
            tag_ele(i) = i
        end do

        return
    end subroutine merge

    subroutine parse_command(arg_pos)

        integer, intent(out) :: arg_pos

        character(len=100) :: textholder
        integer :: i, stat, arglen

        !Get dimension to concatenate along
        call get_command_argument(2, dim, arglen)
        if (arglen == 0) STOP "Missing dim in mode_merge command"
        select case(trim(adjustl(dim)))
        case('x','y','z','none')
            continue
        case default
            print *, dim, " not accepted as a dimension in mode_merge"
            stop 3
        end select
        !Get number of files to read in 
        call get_command_argument(3, textholder, arglen)
        if (arglen == 0) STOP "Number of files to read missing in mode_merge command"
        read(textholder, *, iostat = stat) in_num
        if (stat > 0) STOP "Error reading number of files in, must be integer"

        !Now loop and pull out all files 
        do i = 1, in_num
            call get_command_argument(3+i, textholder, arglen)
            if (arglen == 0) STOP "Fewer files to read in then specified"
            !Make sure this file is readable
            call get_in_file(textholder)
        end do

        !Set argpos accordingly
        arg_pos = 3+in_num
        !Now options loop
        !Check for optional keywords
        do while(.true.)
            if(arg_pos > command_argument_count()) exit
            !Pull out the next argument which should either be a keyword or an option
            arg_pos=arg_pos+1
            call get_command_argument(arg_pos, textholder)
            textholder=adjustl(textholder)

            !Choose what to based on what the option string is 
            select case(trim(textholder))
            case('shift')
                shift_flag = .true.
                do i = 1,3 
                    arg_pos = arg_pos + 1
                    call get_command_argument(arg_pos, textholder, arglen)
                    if (arglen==0) stop "Missing vector component for shift command"
                    read(textholder, *) shift_vec(i)
                end do
            case('replace')
                replace_flag = .true.
                do i = 1,3
                    arg_pos = arg_pos + 1
                    call get_command_argument(arg_pos, replace_str(i), arglen)
                    if (arglen==0) stop "Missing vector component for shift command"
                end do
                arg_pos = arg_pos+1
                call get_command_argument(arg_pos, textholder, arglen)
                read(textholder,*) rc_off

            case default
                !If it isn't an available option to mode merge then we just exit
                exit 
            end select
        end do

    end subroutine parse_command

    subroutine shift(array_start, filenum)
        !This subroutine applies a shift to newly added atoms and elements.  

        integer, dimension(2), intent(in) :: array_start
        integer, intent(in) :: filenum

        integer :: i, ibasis, inod
        real(kind=dp), dimension(3) :: current_shift
        character(len=3) :: alldims

        alldims = 'xyz'
        !Calculate the current shift which is the filenum-1 multiplied by the user specified shift
        current_shift = (filenum-1)*shift_vec

        print *, "Atoms/elements from file ", trim(adjustl(infiles(filenum))), " are shifted by ", current_shift

        !First shift all the atoms
        do i = array_start(1), atom_num
            r_atom(:,i) = r_atom(:,i) + current_shift
        end do 

        !Now shift all the elements
        do i = array_start(2), ele_num
            do inod = 1, ng_node(lat_ele(i))
                do ibasis = 1, basisnum(lat_ele(i))
                    r_node(:,ibasis, inod, i) = r_node(:,ibasis, inod, i) + current_shift
                end do
            end do
        end do

        !If we don't include the wrap command then we have to increase the size of the box
        if(.not.(wrap_flag)) then 
            do i = 1,3
                if (alldims(i:i) /= trim(adjustl(dim))) then 
                    if (current_shift(i) < -lim_zero) then 
                        box_bd(2*i-1) = box_bd(2*i-1) - current_shift(i)
                    else if (current_shift(i) > lim_zero) then 
                        box_bd(2*i) = box_bd(2*i) + current_shift(i)
                    end if
                else if (alldims(i:i) == trim(adjustl(dim))) then 
                    box_bd(2*i) = box_bd(2*i) + current_shift(i)
                end if
            end do
        end if

    end subroutine shift

    subroutine replace(array_start, rbox_bd)
        integer, intent(in) :: array_start(2)
        real(kind = dp), intent(in) :: rbox_bd(6)

        integer ::  ibasis, inod, del_num, del_index(atom_num), nump_ele, interp_start
        integer ::  j, ie, type_interp(max_basisnum*max_esize**3), add_atom_num, orig_atom_num, m, n, o, esize, &
                   ele(3,8), new_ele_num, vlat(3), added_points
        real(kind=dp) :: r_interp(3, max_basisnum*max_esize**3), rfill(3,max_basisnum,max_ng_node), ravg(3), ratom(3,max_basisnum)
        logical :: in_bd, lat_points(max_esize, max_esize, max_esize)
        real(kind=dp) :: del_bd(6)
        integer :: i, c(3), ci, cj, ck, num_nei, nei, delete_num

        !These are the variables containing the cell list information
        integer, dimension(3) :: cell_num
        integer, allocatable :: num_in_cell(:,:,:), which_cell(:,:)
        integer, allocatable :: cell_list(:,:,:,:)

        !First apply the replace vec to all new nodes and elements
        do i = array_start(1), atom_num
            r_atom(:,i) = r_atom(:, i) + replace_vec
        end do  

        do i = array_start(2), ele_num
            do inod = 1, ng_node(lat_ele(i))
                do ibasis=1, basisnum(lat_ele(i))
                    r_node(:, ibasis,inod, i) = r_node(:, ibasis,inod, i) + replace_vec 
                end do
            end do

        end do

        !Calculate new boundary
        do i = 1, 6
            del_bd(i) = rbox_bd(i) + replace_vec((i-1)/2 + 1)
        end do
        
        del_num = 0
        del_index=0
        interp_start = atom_num +1
        !Now loop over all old elements, 
        do ie = 1, array_start(2)-1
            !If any element points are within the boundary then we run the refine code
            if(ele_in_bounds(del_bd, type_ele(ie), size_ele(ie), lat_ele(ie), r_node(:,:,:,ie))) then 
                added_points=0
                del_num = del_num + 1
                del_index(del_num) = ie

                !Find all possible elements that we can make while making sure they aren't in the group
                lat_points(1:size_ele(ie),1:size_ele(ie),1:size_ele(ie)) = .true.

                !Now add the leftover lattice points as atoms, only if they aren't within the new boundaries
                do o = 1, size_ele(ie)
                    do n = 1, size_ele(ie)
                        do m = 1, size_ele(ie)
                            if(lat_points(m,n,o)) then 
                                call get_interp_pos(m,n,o, ie, ratom(:,:))
                                do ibasis = 1, basisnum(lat_ele(ie))
                                    call apply_periodic(ratom(:,ibasis))
                                    added_points=added_points + 1
                                    call add_atom(0, basis_type(ibasis,lat_ele(ie)), sbox_ele(ie), ratom(:,ibasis))
                                end do
                            end if
                        end do
                    end do
                end do

                if (added_points /= (size_ele(ie)**3)) then 
                    print *, "Element ", ie, " is refined incorrectly in refinefill"
                end if
            end if
        end do

        !Once all atoms are added we delete all of the elements
        call delete_elements(del_num, del_index)

        !Now delete overlapping atoms
        allocate(which_cell(3,atom_num))

        !First pass the atom list and atom num to the algorithm which builds the cell list
        print *, rc_off
        call build_cell_list(atom_num, r_atom, 4*rc_off, cell_num, num_in_cell, cell_list, which_cell)

        !Now loop over every atom and figure out if it has neighbors within the rc_off
        del_num = 0
        atom_loop: do i = 1, atom_num
            
            !c is the position of the cell that the atom belongs to 
            c = which_cell(:,i)

            !Check to make sure it hasn't already been deleted
            if(all(c /= 0)) then 
                !Now 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 
                                !and that the neighbor hasn't already been deleted
                                if((nei /= i).and.(nei/= 0)) then 

                                    !Now check to see if it is in the cutoff radius, if it is add it to the delete code
                                    if (norm2(r_atom(:,nei)-r_atom(:,i)) < rc_off) then 
                                        
                                        del_num = del_num + 1

                                        !Make sure to delete the older value
                                        if( (i < array_start(1)).or.(i > interp_start)) then 
                                            del_index(del_num) = i
                                            which_cell(:,i) = 0
                                            cycle atom_loop
                                        else
                                            del_index(del_num) = nei
                                            which_cell(:,nei) = 0 
                                            cell_list(num_nei,c(1) + ck, c(2) + cj, c(3) + ci) = 0
                                        end if
                                    end if
                                end if
                            end do 
                        end do
                    end do 
                end do 
            end if

        end do atom_loop

        print *, "Replace command deletes ", del_num, " atoms"
        !Now delete all the atoms
        call delete_atoms(del_num, del_index(1:del_num))


        return
    end subroutine replace

end module mode_merge