module elements
    
    !This module contains the elements data structures, structures needed for building regions
    !and operations that are done on elements
    use parameters
    use functions
    use subroutines
    use box
    implicit none

    !Data structures used to represent the CAC elements. Each index represents an element
    character(len=100), allocatable :: type_ele(:) !Element type 
    integer, allocatable :: size_ele(:), lat_ele(:), sbox_ele(:) !Element size
    real(kind=dp), allocatable :: r_node(:,:,:,:) !Nodal position array

    integer, save :: ele_num !Number of elements
    integer, save :: node_num !Total number of nodes

    !Data structure used to represent atoms 
    integer, allocatable ::  type_atom(:)!atom type 
    integer, allocatable :: sbox_atom(:)
    real(kind =dp),allocatable :: r_atom(:,:) !atom position
    integer :: atom_num=0 !Number of atoms

    !Mapping atom type to provided name
    character(len=2), dimension(10) :: type_to_name
    integer :: atom_types = 0

    !Variables for creating elements based on primitive cells
    real(kind=dp) :: cubic_cell(3,8), fcc_cell(3,8), fcc_mat(3,3), fcc_inv(3,3)
    integer :: cubic_faces(4,6)

    !Below are lattice type arrays which provide information on the general form of the elements.
    !We currently have a limit of 10 lattice types for simplicities sake but this can be easily increased.
    integer :: lattice_types = 0
    integer :: max_ng_node, ng_node(10) !Max number of nodes per element and number of nodes per element for each lattice type
    integer :: max_esize=0 !Maximum number of atoms per side of element

    !These variables contain information on the basis, for simplicities sake we limit 
    !the user to the definition of 10 lattice types with 10 basis atoms at each lattice point.
    !This can be easily increased with no change to efficiency
    integer :: max_basisnum, basisnum(10) !Max basis atom number, number of basis atoms in each lattice type 
    integer :: basis_type(10,10)
    real(kind=dp) :: lapa(10)

    !Additional module level variables we need 
    logical :: wrap_flag
    public 
    contains

    subroutine lattice_init
        !This subroutine just intializes variables needed for building the different finite 
        !element types

        !First initialize the cubic cell
        cubic_cell = reshape((/ 0.0_dp, 0.0_dp, 0.0_dp, &
                              1.0_dp, 0.0_dp, 0.0_dp, &
                              1.0_dp, 1.0_dp, 0.0_dp, &
                              0.0_dp, 1.0_dp, 0.0_dp, &
                              0.0_dp, 0.0_dp, 1.0_dp, &
                              1.0_dp, 0.0_dp, 1.0_dp, &
                              1.0_dp, 1.0_dp, 1.0_dp, &
                              0.0_dp, 1.0_dp, 1.0_dp /), &
                             shape(fcc_cell))

        !Now we create a list containing the list of vertices needed to describe the 6 cube faces
        cubic_faces(:,1) = (/ 1, 2, 3, 4 /) 
        cubic_faces(:,2) = (/ 1, 2, 6, 5 /)
        cubic_faces(:,3) = (/ 2, 3, 7, 6 /)
        cubic_faces(:,4) = (/ 3, 4, 8, 7 /)
        cubic_faces(:,5) = (/ 1, 4, 8, 5 /)
        cubic_faces(:,6) = (/ 5, 6, 7, 8 /)

        !!Now initialize the fcc primitive cell
        fcc_cell = reshape((/ 0.0_dp, 0.0_dp, 0.0_dp, &
                              0.5_dp, 0.5_dp, 0.0_dp, &
                             0.5_dp, 1.0_dp, 0.5_dp, &
                             0.0_dp, 0.5_dp, 0.5_dp, &
                             0.5_dp, 0.0_dp, 0.5_dp, &
                             1.0_dp, 0.5_dp, 0.5_dp, &
                             1.0_dp, 1.0_dp, 1.0_dp, &
                             0.5_dp, 0.5_dp, 1.0_dp /), &
                             shape(fcc_cell))

        fcc_mat = reshape((/ 0.5_dp, 0.5_dp, 0.0_dp, &
                             0.0_dp, 0.5_dp, 0.5_dp, &
                             0.5_dp, 0.0_dp, 0.5_dp /), &
                             shape(fcc_mat))
        call matrix_inverse(fcc_mat,3,fcc_inv)

        max_basisnum = 0
        basisnum(:) = 0
        ng_node(:) = 0
        node_num = 0
        ele_num = 0
        atom_num = 0
    end subroutine lattice_init

    subroutine cell_init(lapa,esize,ele_type, orient_mat, cell_mat)
        !This subroutine uses the user provided information to transform the finite element cell to the correct
        !size, orientation, and dimensions using the esize, lattice parameter, element_type, and orientation provided 
        !by the user 

        integer, intent(in) :: esize
        real(kind=dp), intent(in) :: lapa, orient_mat(3,3)
        character(len=100), intent(in) :: ele_type

        real(kind=dp), dimension(3,max_ng_node), intent(out) :: cell_mat

        integer :: inod, i
        real(kind=dp), dimension(3,max_ng_node) :: adjustVar

        adjustVar(:,:) = 0.0_dp

        select case(trim(ele_type))

        case('fcc')
            if(lmpcac) then 
                do inod = 1, 8
                    do i = 1,3
                        if(is_equal(cubic_cell(i, inod),0.0_dp)) then 
                            adjustVar(i,inod) = -0.5_dp
                        else 
                            adjustVar(i, inod) = 0.5_dp
                        end if
                    end do
                end do
                adjustVar(:,1:8) = matmul(fcc_mat, adjustVar(:,1:8))
            end if
            cell_mat(:, 1:8) = fcc_cell + adjustVar(:,1:8)
            cell_mat(:,1:8) = lapa * ((esize-1)*matmul(orient_mat, cell_mat(:,1:8)))
        case default
            print *, "Element type ", trim(ele_type), " currently not accepted"
            stop
        end select

    end subroutine cell_init

    subroutine alloc_ele_arrays(n,m)
        !This subroutine used to provide initial allocation for the atom and element arrays

        integer, intent(in) :: n, m !n-size of element arrays, m-size of atom arrays

        integer :: allostat
        
        !Allocate element arrays
        if(n > 0) then 
            allocate(type_ele(n), size_ele(n), lat_ele(n), sbox_ele(n), r_node(3,max_basisnum, max_ng_node,n), &
                       stat=allostat)
            if(allostat > 0) then 
                print *, "Error allocating element arrays in elements.f90 because of: ", allostat
                stop
            end if
        end if

        if(m > 0) then 
            !Allocate atom arrays
            allocate(type_atom(m), sbox_atom(m), r_atom(3,m), stat=allostat)
            if(allostat > 0) then 
                print *, "Error allocating atom arrays in elements.f90 because of: ", allostat
                stop
            end if
        end if    
    end subroutine

    subroutine grow_ele_arrays(n, m)

        integer, intent(in) :: n, m

        integer :: ele_size, atom_size, buffer_size
        integer, allocatable :: temp_int(:)
        real(kind=dp), allocatable :: temp_ele_real(:,:,:,:), temp_real(:,:)
        character(len=100), allocatable :: char_temp(:)

        !The default size we grow the 
        buffer_size = 1024
        !Figure out the size of the atom and element arrays
        ele_size = size(size_ele)
        atom_size = size(type_atom)

        !Check if we need to grow the ele_size, if so grow all the variables
        if ( n+ele_num > size(size_ele)) then 
            
            allocate(temp_int(n+ele_num+buffer_size))
            temp_int(1:ele_size) = lat_ele
            temp_int(ele_size+1:) = 0
            call move_alloc(temp_int, lat_ele)

            allocate(temp_int(n+ele_num+buffer_size))
            temp_int(1:ele_size) = size_ele
            temp_int(ele_size+1:) = 0
            call move_alloc(temp_int, size_ele)

            allocate(temp_int(n+ele_num+buffer_size))
            temp_int(1:ele_size) = lat_ele
            temp_int(ele_size+1:) = 0
            call move_alloc(temp_int, sbox_ele)

            allocate(char_temp(n+ele_num+buffer_size))
            char_temp(1:ele_size) = type_ele
            call move_alloc(char_temp, type_ele)

            allocate(temp_ele_real(3, max_basisnum, max_ng_node, n+ele_num+buffer_size))
            temp_ele_real(:,:,:,1:ele_size) = r_node
            temp_ele_real(:,:,:,ele_size+1:) = 0.0_dp
            call move_alloc(temp_ele_real, r_node)
        end if

        !Now grow atom arrays if needed
        if (m+atom_num > atom_size) then 
            allocate(temp_int(m+atom_num+buffer_size))
            temp_int(1:atom_size) = type_atom
            temp_int(atom_size+1:) = 0
            call move_alloc(temp_int, type_atom)

            allocate(temp_int(m+atom_num+buffer_size))
            temp_int(1:atom_size) = sbox_atom
            temp_int(atom_size+1:) = 0
            call move_alloc(temp_int, sbox_atom)

            allocate(temp_real(3,m+atom_num+buffer_size))
            temp_real(:,1:atom_size) = r_atom
            temp_real(:, atom_size+1:) = 0.0_dp
            call move_alloc(temp_real, r_atom)
        end if
    end subroutine

    subroutine add_element(type, size, lat, sbox, r)
        !Subroutine which adds an element to the element arrays
        integer, intent(in) :: size, lat, sbox
        character(len=100), intent(in) :: type
        real(kind=dp), intent(in) :: r(3, max_basisnum, max_ng_node)

        ele_num = ele_num + 1
        !Check to see if we need to grow the arrays
        call grow_ele_arrays(1,0)
        type_ele(ele_num) = type
        size_ele(ele_num) = size
        lat_ele(ele_num) = lat
        sbox_ele(ele_num) = sbox
       r_node(:,:,:,ele_num) = r(:,:,:)
       node_num = node_num + ng_node(lat)


    end subroutine add_element

    subroutine add_atom(type, sbox, r)
        !Subroutine which adds an atom to the atom arrays
        integer, intent(in) :: type, sbox
        real(kind=dp), intent(in), dimension(3) :: r

        atom_num = atom_num+1
        !Check to see if we need to grow the arrays
        call grow_ele_arrays(0,1)
        type_atom(atom_num) = type
        r_atom(:,atom_num) = r(:)
        sbox_atom(atom_num) = sbox

    end subroutine add_atom

    subroutine add_atom_type(type, inttype)
        !This subroutine adds a new atom type to the module list of atom types
        character(len=2), intent(in) :: type
        integer, intent(out) :: inttype

        integer :: i
        logical :: exists
        
        exists = .false.
        do i=1, 10
            if(type == type_to_name(i)) then 
                exists = .true.
                inttype = i
                exit
            end if
        end do 

        if (exists.eqv..false.) then 
            atom_types = atom_types+1
            if(atom_types > 10) then 
                print *, "Defined atom types are greater than 10 which is currently not supported."
                stop 3
            end if
            type_to_name(atom_types) = type
            inttype = atom_types
        end if
        return
    end subroutine add_atom_type

    subroutine def_ng_node(n, element_types)
        !This subroutine defines the maximum node number among n element types
        integer, intent(in) :: n !Number of element types
        character(len=100), dimension(n) :: element_types !Array of element type strings

        integer :: i

        max_ng_node = 0

        do i=1, n
            select case(trim(adjustl(element_types(i))))
            case('fcc')
                ng_node(i) = 8
            end select

            if(ng_node(i) > max_ng_node) max_ng_node = ng_node(i)
        end do 
    end subroutine def_ng_node

    subroutine set_max_esize
        !This subroutine sets the maximum esize
        max_esize=maxval(size_ele)
    end subroutine

    subroutine interpolate_atoms(type, esize, lat_type, r_in, type_interp, r_interp)
        !This subroutine returns the interpolated atoms from the elements. 

        !Arguments
        character(len=100), intent(in) :: type !The type of element that it is
        integer, intent(in) :: esize !The number of atoms per side
        integer, intent(in) :: lat_type !The integer lattice type of the element
        real(kind=dp), dimension(3,max_basisnum, max_ng_node), intent(in) :: r_in !Nodal positions
        integer, dimension(max_basisnum*max_esize**3), intent(out) :: type_interp !Interpolated atomic positions
        real(kind=dp), dimension(3, max_basisnum*max_esize**3), intent(out) :: r_interp !Interpolated atomic positions

        !Internal variables
        integer :: it, is, ir, ibasis, inod, ia, bnum, lat_type_temp
        real(kind=dp), allocatable :: a_shape(:)
        real(kind=dp) :: t, s, r
        
        !Initialize some variables
        r_interp(:,:) = 0.0_dp
        type_interp(:) = 0
        ia = 0

        !Define bnum based on the input lattice type. If lat_type=0 then we are interpolating lattice points which means
        !the basis is 0,0,0, and the type doesn't matter

        select case(lat_type)
        case(0)
            bnum=1
            lat_type_temp = 1
        case default
            bnum = basisnum(lat_type)
            lat_type_temp = lat_type
        end select

        select case(trim(adjustl(type)))
        case('fcc')
            allocate(a_shape(8))
            !Now loop over all the possible sites
            do it = 1, esize
                t = (1.0_dp*(it-1)-(esize-1)/2)/(1.0_dp*(esize-1)/2)
                do is =1, esize
                    s = (1.0_dp*(is-1)-(esize-1)/2)/(1.0_dp*(esize-1)/2)
                    do ir = 1, esize
                        r = (1.0_dp*(ir-1) - (esize-1)/2)/(1.0_dp*(esize-1)/2)
                        call rhombshape(r,s,t,a_shape)

                        do ibasis = 1, bnum
                            ia = ia + 1
                            do inod = 1, 8
                                type_interp(ia) = basis_type(ibasis,lat_type_temp)
                                r_interp(:,ia) = r_interp(:,ia) + a_shape(inod) * r_in(:,ibasis,inod)

                            end do
                        end do

                    end do
                end do
            end do 
        end select

        if (ia /= bnum*esize**3) then 
            print *, "Incorrect interpolation"
            stop 3
        end if
        return
    end subroutine interpolate_atoms

    subroutine rhombshape(r,s,t, shape_fun)
        !Shape function for rhombohedral elements
        real(kind=8), intent(in) :: r, s, t
        real(kind=8), intent(out) :: shape_fun(8)

        shape_fun(1) =  (1.0-r)*(1.0-s)*(1.0-t)/8.0
        shape_fun(2) =  (1.0+r)*(1.0-s)*(1.0-t)/8.0
        shape_fun(3) =  (1.0+r)*(1.0+s)*(1.0-t)/8.0
        shape_fun(4) =  (1.0-r)*(1.0+s)*(1.0-t)/8.0
        shape_fun(5) =  (1.0-r)*(1.0-s)*(1.0+t)/8.0
        shape_fun(6) =  (1.0+r)*(1.0-s)*(1.0+t)/8.0
        shape_fun(7) =  (1.0+r)*(1.0+s)*(1.0+t)/8.0
        shape_fun(8) =  (1.0-r)*(1.0+s)*(1.0+t)/8.0


        return
    end subroutine rhombshape

    subroutine delete_atoms(num, index)
        !This subroutine deletes atoms from the atom arrays 
        integer, intent(in) :: num
        integer, intent(inout), dimension(num) :: index

        integer :: i

        call heapsort(index)

        !We go from largest index to smallest index just to make sure that we don't miss 
        !accidentally overwrite values which need to be deleted
        do i = num, 1, -1
            if(index(i) == atom_num) then 
                r_atom(:,index(i)) = 0.0_dp
                type_atom(index(i)) = 0
            else
                r_atom(:,index(i)) = r_atom(:, atom_num)
                type_atom(index(i)) = type_atom(atom_num)
            end if            
            atom_num = atom_num - 1
        end do 
    end subroutine

    subroutine delete_elements(num, index)
        !This subroutine deletes elements from the element array 
        integer, intent(in) :: num
        integer, intent(inout), dimension(num) :: index

        integer :: i

        call heapsort(index)

        !We go from largest index to smallest index just to make sure that we don't miss 
        !accidentally overwrite values which need to be deleted
        do i = num, 1, -1
            if(index(i) == ele_num) then 
                r_node(:,:,:,index(i)) = 0.0_dp
                type_ele(index(i)) =''
                size_ele(index(i)) = 0
                lat_ele(index(i)) = 0
                sbox_ele(index(i)) = 0
            else
                node_num = node_num - ng_node(lat_ele(index(i)))
                r_node(:,:,:,index(i)) = r_node(:,:,:,ele_num)
                type_ele(index(i)) = type_ele(ele_num)
                size_ele(index(i)) = size_ele(ele_num)
                lat_ele(index(i)) = lat_ele(ele_num)
                sbox_ele(index(i)) = sbox_ele(ele_num)
            end if            
            ele_num = ele_num - 1
        end do 

    end subroutine delete_elements

    subroutine wrap_atoms
        !This subroutine wraps atoms back into the simulation cell if they have exited for any reason
        integer :: i
        do i = 1, atom_num
            call apply_periodic(r_atom(:,i))
        end do 
    end subroutine wrap_atoms

    subroutine def_new_box
        !This subroutine calculates new box boundaries based on minimum and maximum nodal/atomic positions
        integer :: i, j, inod, ibasis

        real(kind=dp) :: max_bd(3), min_bd(3)

        max_bd(:) = -huge(1.0_dp)
        min_bd(:) = huge(1.0_dp)

        do i = 1, atom_num
            do j = 1, 3
                if (r_atom(j,i) > max_bd(j)) max_bd(j) = r_atom(j,i) + lim_zero
                if (r_atom(j,i) < min_bd(j)) min_bd(j) = r_atom(j,i) - lim_zero
            end do 
        end do 

        do i = 1, ele_num
            do inod = 1, ng_node(lat_ele(i))
                do ibasis = 1, basisnum(lat_ele(i))
                    do j = 1, 3
                        if (r_node(j,ibasis,inod,i) > max_bd(j)) max_bd(j) = r_node(j,ibasis,inod,i) + lim_zero
                        if (r_node(j,ibasis,inod,i) < min_bd(j)) min_bd(j) = r_node(j,ibasis,inod,i) -lim_zero

                    end do 
                end do 
            end do 
        end do 

        do j = 1, 3
            box_bd(2*j) = max_bd(j)
            box_bd(2*j-1) = min_bd(j)
        end do 
    end subroutine

    recursive subroutine parse_pos(i, pos_string, pos)
        !This subroutine parses the pos command allowing for command which include inf
        integer, intent(in) ::  i !The dimension of the position
        character(len=100), intent(in) :: pos_string !The position string
        real(kind=dp), intent(out) :: pos !The output parsed position value

        integer :: iospara, face, ele
        real(kind=dp) :: rand, rone, rtwo, rand_ele_pos(3)
        character(len=100) :: cone, ctwo

        iospara = 0
        
        if(trim(adjustl(pos_string)) == 'inf') then 
            pos=box_bd(2*i)

        else if(trim(adjustl(pos_string)) == '-inf') then
            pos=box_bd(2*i-1)

        else if (trim(adjustl(pos_string)) == 'rand') then 
            call random_number(rand)
            pos = (box_bd(2*i)-box_bd(2*i-1))*rand + box_bd(2*i-1)

        else if (index(pos_string,'rande')>0) then 
            !First select a random element 
            call random_number(rand)
            ele= 1 + floor(ele_num*rand)
            !Now read the rest of the command which specifies the face we need and check 
            !to make sure it's an accepted face number
            cone = pos_string(index(pos_string, '[')+1:index(pos_string, '[')+1)
            read(cone, *) face
            if ((face < 1).or.(face > 6)) stop "Current face number must be 1,2,3,4,5,6. Please check documentation" 
            !Now get the position
            call offset_pos(ele, face, rand_ele_pos)
            pos = rand_ele_pos(i)

        else if (index(pos_string,'rand')>0) then 
            call random_number(rand)
            cone = pos_string(index(pos_string, '[')+1:index(pos_string,':')-1)
            call parse_pos(i, cone, rone)
            ctwo = pos_string(index(pos_string, ':')+1:index(pos_string,']')-1)
            call parse_pos(i, ctwo, rtwo)
            pos = (rtwo - rone)*rand + rone            

        else if ((index(pos_string,'-') > 0).and.(index(pos_string,'inf')>0)) then 
            !Now extract the number we are reducing from infinity
            if(index(pos_string,'inf') < index(pos_string,'-')) then 
                read(pos_string(index(pos_string,'-')+1:), *, iostat=iospara) pos
            else
                read(pos_string(1:index(pos_string,'-')-1), *, iostat=iospara) pos
            end if
            pos = box_bd(2*i) - pos

        else if ((index(pos_string,'+') > 0).and.(index(pos_string,'inf')>0)) then 
            !Now extract the number we are reducing from infinity
           if(index(pos_string,'inf') < index(pos_string,'+')) then 
                read(pos_string(index(pos_string,'+')+1:), *, iostat=iospara) pos
            else
                read(pos_string(1:index(pos_string,'+')-1), *, iostat=iospara) pos
            end if
            pos = box_bd(2*i-1) + pos

        else if ((index(pos_string,'*') > 0).and.(index(pos_string,'inf')>0)) then 
            !Now extract the number we are reducing from infinity
           if(index(pos_string,'inf') < index(pos_string,'*')) then 
                read(pos_string(index(pos_string,'*')+1:), *, iostat=iospara) pos
            else
                read(pos_string(1:index(pos_string,'*')-1), *, iostat=iospara) pos
            end if
            pos = (box_bd(2*i)-box_bd(2*i-1))*pos + box_bd(2*i-1)

        else
            read(pos_string, *, iostat=iospara) pos

        end if

        if (iospara > 0) then 
            print *, "Error reading position argument ", trim(adjustl(pos_string)), ". Please reformat and try again."
        end if
    end subroutine parse_pos

    subroutine offset_pos(ie, iface, pos)
    !This returns a position slightly offset from the center of an element face
    !This is used to return a random position within an element discontinuity

    integer, intent(in) :: ie !Element index
    integer, intent(in) :: iface !Face number between 1 and 6 
    real(kind=dp), dimension(3), intent(out) :: pos !Position vector

    !Other variables we need
    integer :: esize
    real(kind=dp) :: orient(3,3), ori_inv(3,3), lp, r_cubic_node(3,8)

    !First find the offset vector for the face in the untransformed cubic cell
    esize = size_ele(ie)
    select case(iface)
        case(1)
            pos = (/ real(esize-1,dp)/2.0_dp, real(esize-1,dp)/2.0_dp, -10.0_dp**-2.0_dp /)
        case(2)
            pos = (/ real(esize-1,dp)/2.0_dp, -10.0_dp**-2.0_dp, real(esize-1,dp)/2.0_dp /)
        case(3)
            pos = (/ (esize-1)+10.0_dp**-2.0_dp, real(esize-1,dp)/2.0_dp, real(esize-1,dp)/2.0_dp /)
        case(4)
            pos = (/ real(esize-1,dp)/2.0_dp, (esize-1)+10.0_dp**-2.0_dp, real(esize-1,dp)/2.0_dp /)
        case(5)
            pos = (/ -10.0_dp**-2.0_dp, real(esize-1,dp)/2.0_dp, real(esize-1,dp)/2.0_dp /)
        case(6)
            pos = (/ real(esize-1,dp)/2.0_dp, real(esize-1,dp)/2.0_dp, (esize-1)+10.0_dp**-2.0_dp /)
    end select

    !Now transform it to real space and adjust it to the position of the element in the first node.
    select case (trim(adjustl(type_ele(ie))))

    case('fcc')
        !First we have to extract the element lattice parameter
        call matrix_inverse(sub_box_ori(:,:,sbox_ele(ie)),3,ori_inv)
        r_cubic_node = r_node(:,1,:,ie)
        r_cubic_node = matmul(fcc_inv,matmul(ori_inv,r_cubic_node))
        lp = (maxval(r_cubic_node(1,:))-minval(r_cubic_node(1,:)))/(esize-1)
        pos = matmul(sub_box_ori(:,:,sbox_ele(ie)),matmul(fcc_mat,pos))*lp
        pos = pos + r_node(:,1, 1, ie)

    case default
        print *, trim(adjustl(type_ele(ie))), " is not a currently accepted element type for random element positions"
        stop 3

    end select
    end subroutine

end module elements