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

    character(len=4) :: dim
    integer :: in_num, new_starts(2)
    real(kind=dp) :: shift_vec(3)
    logical :: shift_flag

    public 
    contains
    subroutine merge(arg_pos)
        
        integer, intent(out) :: arg_pos 
        integer :: i
        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)
        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 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
end module mode_merge