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
    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 
    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)

    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, 4, 8, 5 /)
        cubic_faces(:,2) = (/ 2, 3, 7, 6 /)
        cubic_faces(:,3) = (/ 1, 2, 6, 5 /)
        cubic_faces(:,4) = (/ 3, 4, 8, 7 /)
        cubic_faces(:,5) = (/ 1, 2, 3, 4 /) 
        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), 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_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
        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, r)
        !Subroutine which adds an atom to the atom arrays
        integer, intent(in) :: type
        real(kind=dp), intent(in), dimension(3) :: r

        atom_num = atom_num+1
        type_atom(atom_num) = type
        r_atom(:,atom_num) = r(:)

    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 :: i, 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 /= 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, j

        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
end module elements