module opt_group

    !This module contains all code associated with dislocations

    use parameters
    use elements
    use subroutines
    use box 
    implicit none

    integer :: group_ele_num, group_atom_num, remesh_size, remesh_type
    character(len=15) :: type, shape !Type indicates what element type is selected and shape is the group shape
    real(kind=dp) :: block_bd(6), disp_vec(3), remesh_lat_pam
    logical :: displace, delete, max_remesh, refine

    integer, allocatable :: element_index(:), atom_index(:) 

    public 
    contains

    subroutine group(arg_pos)
        !Main calling function for the group option
        integer, intent(inout) :: arg_pos

        print *, '-----------------------Option Group-------------------------'

        group_ele_num = 0
        group_atom_num = 0
        remesh_size=0
        displace=.false.
        delete=.false.
        ! max_remesh=.false.
        refine = .false.

        if(allocated(element_index)) deallocate(element_index)
        if(allocated(atom_index)) deallocate(atom_index)

        call parse_group(arg_pos)

        call get_group

        !Now call the transformation functions for the group
        if(displace) call displace_group

        if(remesh_size > 0) call remesh_group

        if(delete) call delete_group

        if(refine) call refine_group

    end subroutine group

    subroutine parse_group(arg_pos)
        !Parse the group command
        integer, intent(inout) :: arg_pos

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

        !Parse type and shape command
        arg_pos = arg_pos + 1
        call get_command_argument(arg_pos, type, arglen)
        if (arglen==0) STOP "Missing select_type in group command"
        select case(trim(adjustl(type)))
        case('atoms', 'elements', 'both')
            continue
        case default
            print *, "Select_type ", trim(adjustl(type)), " is not an accept group selection criteria. ", &
                     "Please select from atoms, nodes, or both."
        end select

        arg_pos = arg_pos + 1
        call get_command_argument(arg_pos, shape, arglen)
        if (arglen==0) STOP "Missing group_shape in group command"

        !Now parse the arguments required by the user selected shape
        select case(trim(adjustl(shape)))
        case('block')
            do i= 1, 6
                arg_pos  = arg_pos + 1
                call get_command_argument(arg_pos, textholder, arglen)
                if (arglen==0) STOP "Missing block boundary in dislgen command"
                call parse_pos(int((i+1)/2), textholder, block_bd(i))   
            end do 

        case default
            print *, "Group shape ", trim(adjustl(shape)), " is not currently accepted. Please check documentation ", &
                     "for accepted group shapes."
        end select

        !Now parse the additional options which may be present
        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('displace')
                displace = .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, *) disp_vec(i)
                end do
            case('refine')
                refine=.true.
            case('remesh')
                arg_pos = arg_pos + 1
                call get_command_argument(arg_pos, textholder, arglen)
                if (arglen==0) stop "Missing remesh element size in group command"
                read(textholder, *) remesh_size
                arg_pos = arg_pos + 1
                call get_command_argument(arg_pos, textholder, arglen)
                if (arglen==0) stop "Missing remesh lattice parameter in group command"
                read(textholder, *) remesh_lat_pam
                arg_pos = arg_pos + 1
                call get_command_argument(arg_pos, textholder, arglen)
                if (arglen==0) stop "Missing remesh type in group command"
                read(textholder, *) remesh_type
            case('max')
                max_remesh =.true.
            case('delete')
                delete=.true.
            case default
                !If it isn't an available option to opt_disl then we just exit
                exit 
            end select
        end do
    end subroutine parse_group

    subroutine get_group
        !This subroutine finds all elements and/or atoms within the group boundaries
        !specified by the user.
        integer :: i, inod, ibasis
        integer, allocatable :: resize_array(:) 
        real(kind=dp) :: r_center(3)

        select case(trim(adjustl(shape)))
        case('block')

            print *, "Group has block shape with boundaries: ", block_bd

            !Allocate variables to arbitrary size
            allocate(element_index(1024), atom_index(1024))
            !Check the type to see whether we need to find the elements within the group
            select case(trim(adjustl(type)))
            case('elements', 'both')
                do i = 1, ele_num
                    r_center(:) = 0.0_dp
                    do inod = 1, ng_node(lat_ele(i))
                        do ibasis = 1, basisnum(lat_ele(i))
                            r_center = r_center + r_node(:,ibasis,inod,i)/(basisnum(lat_ele(i))*ng_node(lat_ele(i)))
                        end do
                    end do

                    if (in_block_bd(r_center, block_bd)) then 
                        group_ele_num  = group_ele_num + 1
                        if(group_ele_num > size(element_index)) then
                            allocate(resize_array(size(element_index) + 1024))
                            resize_array(1:group_ele_num-1) = element_index
                            resize_array(group_ele_num:) = 0
                            call move_alloc(resize_array, element_index)
                        end if

                        element_index(group_ele_num) = i
                    end if 
                end do 
            end select

            !Check the type to see if we need to find the atoms within the group
            select case(trim(adjustl(type)))
            case('atoms','both')
                do i = 1, atom_num
                    if(in_block_bd(r_atom(:,i),block_bd)) then 
                        group_atom_num = group_atom_num + 1
                        if (group_atom_num > size(atom_index)) then 
                            allocate(resize_array(size(atom_index) + 1024))
                            resize_array(1:group_atom_num -1) = atom_index
                            resize_array(group_atom_num:) = 0
                            call move_alloc(resize_array, atom_index)
                        end if

                        atom_index(group_atom_num) = i
                    end if
                end do
            end select
        end select 

        print *, 'Group contains ', group_ele_num, " elements and ", group_atom_num, " atoms."
    end subroutine get_group 

    subroutine displace_group
        !This subroutine applies a displacement to elements/atoms in the groups

        integer :: i, inod, ibasis

        print *, "Elements/atoms in group displaced by ", disp_vec

        !Displace atoms
        do i = 1, group_atom_num
            r_atom(:,atom_index(i)) = r_atom(:,atom_index(i)) + disp_vec
        end do  

        !Displace elements
        do i = 1, group_ele_num
            do inod = 1, ng_node(lat_ele(element_index(i)))
                do ibasis = 1, basisnum(lat_ele(element_index(i)))
                    r_node(:,ibasis,inod,element_index(i)) = r_node(:,ibasis,inod,element_index(i)) + disp_vec
                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 (disp_vec(i) < -lim_zero) then 
                    box_bd(2*i-1) = box_bd(2*i-1) - disp_vec(i)
                else if (disp_vec(i) > lim_zero) then 
                    box_bd(2*i) = box_bd(2*i) + disp_vec(i)
                end if
            end do
        end if

    end subroutine displace_group

    subroutine refine_group
        !This command is used to remesh the group to a desired element size

        integer :: i, j, ie, type_interp(max_basisnum*max_esize**3), add_atom_num, orig_atom_num
        real(kind=dp) :: r_interp(3, max_basisnum*max_esize**3)

        !Refining to atoms

        if(group_ele_num > 0) then 
            orig_atom_num = atom_num
            !Estimate number of atoms we are adding, this doesn't have to be exact
            add_atom_num = group_ele_num*basisnum(lat_ele(element_index(1)))*size_ele(element_index(1))**3
            call grow_ele_arrays(0,add_atom_num)
            do i = 1, group_ele_num
                ie = element_index(i)
                !Get the interpolated atom positions
                call interpolate_atoms(type_ele(ie), size_ele(ie), lat_ele(ie), r_node(:,:,:,ie), type_interp, r_interp)
                !Loop over all interpolated atoms and add them to the system, we apply periodic boundaries 
                !here as well to make sure they are in the box
                do j = 1, basisnum(lat_ele(ie))*size_ele(ie)**3
                    call apply_periodic(r_interp(:,j))
                    call add_atom(type_interp(j), sbox_ele(ie), r_interp(:,j))
                end do 
            end do
            !Once all atoms are added we delete all of the elements
            call delete_elements(group_ele_num, element_index)
            print *, group_ele_num, " elements of group are refined to ", atom_num -orig_atom_num, " atoms."
        end if                

    end subroutine

    subroutine remesh_group
        !This command is used to remesh the group to a desired element size

        integer :: i, j, ix, iy, iz, inod, ibasis, ie, type_interp(max_basisnum*max_esize**3), add_atom_num, orig_atom_num,   &
                   current_esize, dof, max_lat(3), r_lat(3), ele(3,8), vlat(3), bd_in_lat(6), bd_in_array(3)
        real(kind=dp) :: r_interp(3, max_basisnum*max_esize**3),  ori_inv(3,3),  r(3), &
                         r_new_node(3,max_basisnum, max_ng_node), orient(3,3), group_in_lat(3,8)
        logical, allocatable :: lat_points(:,:,:)
        character(len=100) :: remesh_ele_type


        !Right now we just hardcode only remeshing to elements
        remesh_ele_type = 'fcc'

        !Get the orientations, this assumes that the orientation of the subbox for the first atom is the 
        !same as the rest of the atoms
        !If this assumption is false then the code will break and exit
        orient = sub_box_ori(:, :, sbox_atom(atom_index(1)))
        call matrix_inverse(orient,3,ori_inv)

        !First calculate max position in lattice space to be able to allocate lat_points array, also sum the total numbers of 
        !degrees of freedom which are added
        dof = 0 
        
        select case(trim(adjustl(shape)))
            case('block')
                group_in_lat = reshape((/ block_bd(1),block_bd(3),block_bd(5), &
                                          block_bd(2),block_bd(3),block_bd(5), &
                                          block_bd(2),block_bd(4),block_bd(5), &
                                          block_bd(1),block_bd(4),block_bd(5), &
                                          block_bd(1),block_bd(3),block_bd(6), &
                                          block_bd(2),block_bd(3),block_bd(6), &
                                          block_bd(2),block_bd(4),block_bd(6), &
                                          block_bd(1),block_bd(4),block_bd(6) /), [3,8])

                group_in_lat = matmul(fcc_inv, matmul(ori_inv, group_in_lat/remesh_lat_pam))
                do i = 1, 3
                    bd_in_lat(2*i-1) = nint(minval(group_in_lat(i,:)))
                    bd_in_lat(2*i) = nint(maxval(group_in_lat(i,:)))
                end do

        end select 

        allocate(lat_points(bd_in_lat(2)-bd_in_lat(1)+10, bd_in_lat(4)-bd_in_lat(3)+10, bd_in_lat(6)-bd_in_lat(5)+10))
        lat_points(:,:,:) = .false.

        !Now place all group atoms and group interpolated atoms into lat_points
        do i = 1, group_atom_num
            r = r_atom(:,atom_index(i))/remesh_lat_pam
            r = matmul(fcc_inv,matmul(ori_inv,r))
            do j = 1, 3
                r_lat(j) = nint(r(j))
            end do 
            !Do a check to make sure the code is working and that lattice points aren't being written on top of each other. 
            !This is primarily a debugging statement
            if((r_lat(1)==13).and.(r_lat(2)==13).and.(r_lat(3)==-13)) then 
                print *, i, atom_index(i), r_atom(:,atom_index(i))
            end if 
            if(lat_points(r_lat(1)-bd_in_lat(1)+5,r_lat(2)-bd_in_lat(3)+5,r_lat(3)-bd_in_lat(5)+5)) then 
                stop "Multiple atoms share same position in lat point array, this shouldn't happen"
            else 
                lat_points(r_lat(1)-bd_in_lat(1)+5, r_lat(2)-bd_in_lat(3)+5, r_lat(3)-bd_in_lat(5)+5) = .true.
            end if
        end do 

        !Now place interpolated atoms within lat_points array
        do i =1, group_ele_num
            ie = element_index(i)
            call interpolate_atoms(type_ele(ie), size_ele(ie), lat_ele(ie), r_node(:,:,:,ie), type_interp, r_interp)
            do j = 1, size_ele(ie)**3 * basisnum(lat_ele(ie))
                r_lat = matmul(fcc_inv,matmul(ori_inv, r_interp(:,atom_index(i))/remesh_lat_pam)) + 1
                !Do a check to make sure the code is working and that lattice points aren't being written on top of each other. 
                !This is primarily a debugging statement
                if(lat_points(r_lat(1)-bd_in_lat(1)+5,r_lat(2)-bd_in_lat(3)+5,r_lat(3)-bd_in_lat(5)+5)) then 
                    stop "Multiple atoms/interpolated atoms share same position in lat point array, this shouldn't happen"
                else 
                    lat_points(r_lat(1)-bd_in_lat(1)+5, r_lat(2)-bd_in_lat(3)+5, r_lat(3)-bd_in_lat(5)+5) = .true.
                end if
            end do 
        end do 

        !Delete all elements and atoms to make space for new elements and atoms
        call delete_atoms(group_atom_num, atom_index)
        call delete_elements(group_ele_num, element_index)

        !Now run remeshing algorithm, not the most optimized or efficient but gets the job done

        ele = (remesh_size-1)*cubic_cell

        !Figure out new looping boundaries
        bd_in_array(1) = bd_in_lat(2) - bd_in_lat(1) + 10
        bd_in_array(2) = bd_in_lat(4) - bd_in_lat(3) + 10
        bd_in_array(3) = bd_in_lat(6) - bd_in_lat(5) + 10

        zloop: do iz = 1, bd_in_array(3)
            yloop: do iy = 1, bd_in_array(2)
                xloop: do ix = 1, bd_in_array(1)
                    if (lat_points(ix, iy,iz)) then 
                        
                        !Check to see if the element overshoots the bound
                        if (iz+remesh_size-1 > bd_in_array(3)) then
                            exit zloop
                        else if (iy+remesh_size-1 > bd_in_array(2)) then
                            cycle zloop
                        else if (ix+remesh_size-1 > bd_in_array(1)) then 
                            cycle yloop
                        end if  

                        if (all(lat_points(ix:ix+remesh_size-1,iy:iy+remesh_size-1,iz:iz+remesh_size-1))) then 
                            do inod = 1, 8
                                vlat = ele(:,inod) + (/ix, iy, iz /)
                                do i = 1, 3
                                    vlat = vlat + bd_in_lat(2*i-1)-5
                                end do 
                               r_new_node(:,1,inod) = matmul(orient, matmul(fcc_mat, vlat))*remesh_lat_pam
                            end do 

                            lat_points(ix:ix+remesh_size-1,iy:iy+remesh_size-1,iz:iz+remesh_size-1) = .false.
                            !Add the element, for the sbox we just set it to the same sbox that we get the orientation from.
                            !In this case it is from the sbox of the first atom in the group.
                            call add_element(remesh_ele_type, remesh_size, remesh_type, sbox_atom(atom_index(1)),r_new_node)

                        end if
                    end if
                end do xloop
            end do yloop
        end do zloop

        !Now we have to add any leftover lattice points as atoms
        do iz = 1, bd_in_array(3)
            do iy=1, bd_in_array(2)
                do ix = 1, bd_in_array(1)
                    if(lat_points(ix,iy,iz)) then 
                        vlat = (/ ix, iy, iz /)
                        do i = 1, 3
                            vlat = vlat + bd_in_lat(2*i-1)-5
                        end do 
                       r = matmul(orient, matmul(fcc_mat, vlat))*remesh_lat_pam
                       call add_atom(remesh_type, sbox_atom(atom_index(1)), r)
                    end if
                end do
            end do 
        end do
        
    end subroutine remesh_group

    subroutine delete_group
        !This subroutine deletes all atoms/elements within a group

        print *, "Deleting group containing ", group_atom_num, " atoms and ", group_ele_num, " elements."

        !Delete atoms
        call delete_atoms(group_atom_num, atom_index)

        !Delete elements
        call delete_elements(group_ele_num, element_index)

    end subroutine delete_group
end module opt_group