! Module  : vgrid_m
! Purpose : Define vgrid_t type, providing support for virtual N-dimensional
!           interpolating grids
!
! Copyright 2021-2023 Rich Townsend & The MSG Team
!
! This file is part of MSG. MSG is free software: you can redistribute
! it and/or modify it under the terms of the GNU General Public
! License as published by the Free Software Foundation, version 3.
!
! MSG is distributed in the hope that it will be useful, but WITHOUT
! ANY WARRANTY; without even the implied warranty of MERCHANTABILITY
! or FITNESS FOR A PARTICULAR PURPOSE.  See the GNU General Public
! License for more details.
!
! You should have received a copy of the GNU General Public License
! along with this program.  If not, see <http://www.gnu.org/licenses/>.

#:include 'forum.inc'

module vgrid_m

   ! Uses
   
   use forum_m

   use axis_m
   use file_m
   use ninterp_m
   use stat_m

   #:if OMP is not None
   use omp_lib
   #:endif

   use ISO_FORTRAN_ENV

   ! No implicit typing

   implicit none

   ! Derived-type definitions

   type vgrid_t
      private
      type(axis_t), allocatable :: axes(:)
      logical, allocatable      :: v_lin_mask(:)
      logical, allocatable      :: c_lin_mask(:)
      integer, allocatable      :: v_seq_lin(:)
      integer, allocatable      :: v_lin_seq(:)
      integer, allocatable      :: shape(:)
      integer                   :: rank
      integer                   :: n_v_lin
      integer                   :: n_c_lin
      integer                   :: n_seq
      integer                   :: n_stn
   contains
      private
      procedure, public :: get_n_seq
      procedure, public :: get_rank
      procedure, public :: get_shape
      procedure, public :: get_axis
      procedure, public :: map
      procedure         :: interp_0_
      procedure         :: interp_1_
      generic, public   :: interp => interp_0_, interp_1_
      procedure, public :: adjust_x_vec
      procedure         :: locate_
      procedure         :: c_lin_vec_
      procedure         :: c_vec_lin_
      procedure         :: v_lin_vec_
      procedure         :: v_vec_lin_
      procedure         :: v_lin_cnr_
      procedure         :: v_lin_stn_
      procedure         :: setup_interp_
      procedure, public :: read
      procedure, public :: write
   end type vgrid_t

   ! Interfaces

   interface vgrid_t
      module procedure vgrid_t_
   end interface vgrid_t

   ! Access specifiers

   private

   public :: vgrid_t
   public :: optimize

   ! Procedures

contains

   function vgrid_t_ (x_vecs, labels) result(vgrid)

      real(RD), intent(in)     :: x_vecs(:,:)
      character(*), intent(in) :: labels(:)
      type(vgrid_t)            :: vgrid

      integer                   :: rank
      integer                   :: n_seq
      real(RD), allocatable     :: x(:)
      type(axis_t), allocatable :: axes(:)
      integer                   :: r
      integer, allocatable      :: v_lin_seq(:)
      integer                   :: v_seq
      integer                   :: p
      integer                   :: i
      integer                   :: n

      @:CHECK_BOUNDS(SIZE(labels), SIZE(x_vecs, 1))

      ! Construct vgrid from the supplied list of vertex coordinates
      ! and axis labels

      ! First, build the axes

      rank = SIZE(labels)
      n_seq = SIZE(x_vecs, 2)

      allocate(x(n_seq))
      allocate(axes(rank))
      
      do r = 1, rank

         x = x_vecs(r,:)

         associate (u => unique_indices(x))
           axes(r) = axis_t(x(u), labels(r))
         end associate

      end do

      ! Create the vertex sequence-to-linear index map

      allocate(v_lin_seq(n_seq))

      do v_seq = 1, n_seq

         ! Determine the linear index of the vertex

         v_lin_seq(v_seq) = 1
         p = 1

         do r = 1, rank
            call axes(r)%get_n(n)
            call axes(r)%locate(x_vecs(r,v_seq), i)
            v_lin_seq(v_seq) = v_lin_seq(v_seq) + (i-1)*p
            p = p*n
         end do

      end do

      ! Construct vgrid

      vgrid = vgrid_t_map_(axes, v_lin_seq)
      
      ! Finish

      return

   end function vgrid_t_

   !****

   function vgrid_t_mask_ (axes, mask) result(vgrid)

      type(axis_t), intent(in) :: axes(:)
      logical, intent(in)      :: mask(:)
      type(vgrid_t)            :: vgrid

      integer              :: n_v_lin
      integer              :: n_seq
      integer, allocatable :: v_lin_seq(:)
      integer              :: v_seq
      integer              :: v_lin

      ! Construct vgrid from the axes and occupancy mask. This is a
      ! compatibility constructor for grids read in the "old" (rev <=
      ! 2) format

      ! Set up the vertex sequence-to-linear mask

      n_v_lin = SIZE(mask)
      n_seq = COUNT(mask)

      allocate(v_lin_seq(n_seq))

      v_seq = 0

      do v_lin = 1, n_v_lin

         if (mask(v_lin)) then
            v_seq = v_seq + 1
            v_lin_seq(v_seq) = v_lin
         end if

      end do

      ! Construct vgrid

      vgrid = vgrid_t_map_(axes, v_lin_seq)

      ! Finish

      return

   end function vgrid_t_mask_

   !****

   function vgrid_t_map_ (axes, v_lin_seq) result(vgrid)

      type(axis_t), intent(in) :: axes(:)
      integer, intent(in)      :: v_lin_seq(:)
      type(vgrid_t)            :: vgrid

      integer                   :: rank
      integer                   :: r
      integer                   :: shape(SIZE(axes))
      integer                   :: n_v_lin
      integer                   :: n_seq
      logical                   :: mask(SIZE(axes))
      type(axis_t), allocatable :: axes_nd(:)
      integer                   :: rank_nd
      integer, allocatable      :: shape_nd(:)
      integer                   :: n_c_lin
      logical, allocatable      :: v_lin_mask(:)
      integer                   :: v_seq
      integer, allocatable      :: v_seq_lin(:)
      logical, allocatable      :: c_lin_mask(:)
      integer, allocatable      :: c_vec(:)
      integer                   :: c_lin
      integer                   :: i
      integer                   :: v_lin

      ! Construct vgrid from the axes and sequence-to-linear index map

      ! First, determine the shape

      rank = SIZE(axes)

      do r = 1, rank
         call axes(r)%get_n(shape(r))
      end do

      n_v_lin = PRODUCT(shape)
      n_seq = SIZE(v_lin_seq)

      @:ASSERT_DEBUG(MINVAL(v_lin_seq) >= 1, 'Invalid v_lin_seq')
      @:ASSERT_DEBUG(MAXVAL(v_lin_seq) <= n_v_lin, 'Invalid v_lin_seq')

      ! Strip out degenerate (length-1) axes

      mask = shape > 1

      axes_nd = PACK(axes, mask)

      rank_nd = COUNT(mask)
      shape_nd = PACK(shape, mask)

      n_c_lin = PRODUCT(shape_nd-1)

      ! Set up the linear vertex mask and the linear-to-sequence
      ! index map

      allocate(v_lin_mask(n_v_lin))
      v_lin_mask = .FALSE.

      allocate(v_seq_lin(n_v_lin))
      v_seq_lin = 0

      do v_seq = 1, n_seq

         v_lin = v_lin_seq(v_seq)

         if (.NOT. v_lin_mask(v_lin)) then
            v_lin_mask(v_lin) = .TRUE.
            v_seq_lin(v_lin) = v_seq
         end if

      end do

      ! Initialize vgrid

      vgrid%axes = axes_nd

      call MOVE_ALLOC(v_lin_mask, vgrid%v_lin_mask)

      vgrid%v_lin_seq = v_lin_seq
      call MOVE_ALLOC(v_seq_lin, vgrid%v_seq_lin)

      vgrid%shape = shape_nd
      vgrid%rank = rank_nd

      vgrid%n_v_lin = n_v_lin
      vgrid%n_c_lin = n_c_lin
      vgrid%n_seq = n_seq
      vgrid%n_stn = 4**rank_nd

      ! Initialize c_lin_mask last, to take advantage of indexing
      ! routines

      allocate(c_lin_mask(n_c_lin))
      allocate(c_vec(rank_nd))

      cell_loop: do c_lin = 1, n_c_lin

         c_vec = vgrid%c_vec_lin_(c_lin)

         ! Iterate over corner vertices of the cell

         do i = 1, 2**rank_nd

            v_lin = vgrid%v_lin_cnr_(i, c_vec)

            if (.NOT. vgrid%v_lin_mask(v_lin)) then
               c_lin_mask(c_lin) = .FALSE.
               cycle cell_loop
            end if

         end do

         c_lin_mask(c_lin) = .TRUE.

      end do cell_loop

      call MOVE_ALLOC(c_lin_mask, vgrid%c_lin_mask)
      
      ! Finish

      return

   end function vgrid_t_map_

   !****

   #:for name, type in (('n_seq', 'integer'),  &
                        ('rank', 'integer'))

      subroutine get_${name}$(self, ${name}$)

         class(vgrid_t), intent(in) :: self
         ${type}$, intent(out)      :: ${name}$

         ! Get ${name}$

         ${name}$ = self%${name}$

         ! Finish

         return

      end subroutine get_${name}$

   #:endfor
      
   !****

   subroutine get_shape(self, shape)

      class(vgrid_t), intent(in) :: self
      integer, intent(out)       :: shape(:)

      @:CHECK_BOUNDS(SIZE(shape), self%rank)

      ! Get the shape

      shape = self%shape

      ! Finish

      return

   end subroutine get_shape

   !****
      
   subroutine get_axis(self, i, axis)

      class(vgrid_t), intent(in) :: self
      integer, intent(in)        :: i
      type(axis_t), intent(out)  :: axis

      ! Get the i'th axis

      axis = self%axes(i)

      ! Finish

      return

   end subroutine get_axis

   !****

   subroutine map(self, map_proc, stat)

      class(vgrid_t), intent(in) :: self
      interface
         subroutine map_proc(v_seq, stat)
            integer, intent(in)            :: v_seq
            integer, intent(out), optional :: stat
         end subroutine map_proc
      end interface
      integer, intent(out), optional :: stat

      integer :: v_lin
      integer :: v_seq

      ! Apply the map_proc function to each unmasked vertex in the
      ! grid

      map_loop : do v_lin = 1, self%n_v_lin

         if (self%v_lin_mask(v_lin)) then

            v_seq = self%v_seq_lin(v_lin)

            call map_proc(v_seq, stat)
            if (PRESENT(stat)) then
               if (stat /= STAT_OK) return
            end if

         end if

      end do map_loop

      ! Finish

      return

   end subroutine map

   !****

   subroutine interp_0_(self, data_proc, x_vec, f_int, stat, vderiv, order)

      class(vgrid_t), intent(in) :: self
      interface
         subroutine data_proc(v_seq, data, stat)
            use forum_m
            integer, intent(in)            :: v_seq
            real(RD), intent(out)          :: data
            integer, intent(out), optional :: stat
         end subroutine data_proc
      end interface
      real(RD), intent(in)           :: x_vec(:)
      real(RD), intent(out)          :: f_int
      integer, intent(out), optional :: stat
      logical, intent(in), optional  :: vderiv(:)
      integer, intent(in), optional  :: order

      integer  :: c_vec(self%rank)
      logical  :: mask(self%n_stn)
      real(RD) :: vdelta(3,self%rank)
      real(RD) :: vu(self%rank)
      real(RD) :: K_D(self%n_stn,self%n_stn)
      real(RD) :: K_I(self%n_stn)
      real(RD) :: K(self%n_stn)
      real(RD) :: F
      integer  :: i
      integer  :: v_lin
      integer  :: v_seq

      @:CHECK_BOUNDS(SIZE(x_vec), self%rank)

      ! Interpolate a scalar value

      ! Set up the interpolation

      call self%setup_interp_(x_vec, c_vec, mask, vdelta, vu, stat)
      if (PRESENT(stat)) then
         if (stat /= STAT_OK) return
      end if

      ! Evaluate the kernels

      call eval_deriv_kernel(self%rank, mask, vdelta, K_D, stat, order)
      if (PRESENT(stat)) then
         if (stat /= STAT_OK) return
      end if

      call eval_interp_kernel(self%rank, vu, vdelta, K_I, vderiv)

      K = MATMUL(K_I, K_D)

      ! Initialize f_int

      f_int = 0._RD

      ! Now perform the interpolation

      do i = 1, self%n_stn

         if (K(i) /= 0._RD) then

            ! Evaluate data for this vertex

            v_lin = self%v_lin_stn_(i, c_vec)
            v_seq = self%v_seq_lin(v_lin)

            call data_proc(v_seq, F, stat)
            if (PRESENT(stat)) then
               if (stat /= STAT_OK) return
            endif

            ! Add the contribution to f_int

            f_int = f_int + K(i)*F

         end if

      end do

      ! Finish

      if (PRESENT(stat)) stat = STAT_OK

      return

   end subroutine interp_0_

   !****

   subroutine interp_1_(self, data_proc, x_vec, f_int, stat, vderiv, order)

      class(vgrid_t), intent(in) :: self
      interface
         subroutine data_proc(v_seq, data, stat)
            use forum_m
            integer, intent(in)            :: v_seq
            real(RD), intent(out)          :: data(:)
            integer, intent(out), optional :: stat
         end subroutine data_proc
      end interface
      real(RD), intent(in)           :: x_vec(:)
      real(RD), intent(out)          :: f_int(:)
      integer, intent(out), optional :: stat
      logical, intent(in), optional  :: vderiv(:)
      integer, intent(in), optional  :: order

      integer  :: c_vec(self%rank)
      logical  :: mask(self%n_stn)
      real(RD) :: vdelta(3,self%rank)
      real(RD) :: vu(self%rank)
      real(RD) :: K_D(self%n_stn,self%n_stn)
      real(RD) :: K_I(self%n_stn)
      real(RD) :: K(self%n_stn)
      integer  :: n_m
      integer  :: i
      integer  :: v_lin
      integer  :: v_seq
      real(RD) :: F(SIZE(f_int))
      integer  :: m

      #:if OMP is not None
         integer :: num_threads
         integer :: thread_num
         integer :: block_size
         integer :: m_a
         integer :: m_b
      #:endif

      @:CHECK_BOUNDS(SIZE(x_vec), self%rank)

      ! Interpolate an array value

      ! Set up the interpolation

      call self%setup_interp_(x_vec, c_vec, mask, vdelta, vu, stat)
      if (PRESENT(stat)) then
         if (stat /= STAT_OK) return
      end if

      ! Evaluate the kernels

      call eval_deriv_kernel(self%rank, mask, vdelta, K_D, stat, order)
      if (PRESENT(stat)) then
         if (stat /= STAT_OK) return
      end if

      call eval_interp_kernel(self%rank, vu, vdelta, K_I, vderiv)

      K = MATMUL(K_I, K_D)

      ! Initialize f_int

      n_m = SIZE(f_int)

      !$OMP PARALLEL DO
      do m = 1, n_m
         f_int(m) = 0._RD
      end do

      ! Now perform the interpolation

      do i = 1, self%n_stn

         if (K(i) /= 0._RD) then

            ! Evaluate data for this node

            v_lin = self%v_lin_stn_(i, c_vec)
            v_seq = self%v_seq_lin(v_lin)

            call data_proc(v_seq, F, stat)
            if (PRESENT(stat)) then
               if (stat /= STAT_OK) return
            endif

            ! Add the contribution to f_int

            #:if OMP is not None

               !$OMP PARALLEL PRIVATE(num_threads, block_size, thread_num, m_a, m_b)

               num_threads = omp_get_num_threads()

               block_size = n_m/num_threads
               if (block_size*num_threads < n_m) block_size = block_size+1

               thread_num = omp_get_thread_num()

               m_a = thread_num*block_size + 1
               m_b = thread_num*block_size + block_size

               m_b = MIN(m_b, n_m)

               f_int(m_a:m_b) = f_int(m_a:m_b) + K(i)*F(m_a:m_b)

               !$OMP END PARALLEL

            #:else

               f_int = f_int + K(i)*F

            #:endif

         end if

      end do
               
      ! Finish

      if (PRESENT(stat)) stat = STAT_OK

      return

   end subroutine interp_1_

   !****

   subroutine adjust_x_vec(self, x_vec, dx_vec, x_adj, stat)

      class(vgrid_t), intent(in)     :: self
      real(RD), intent(in)           :: x_vec(:)
      real(RD), intent(in)           :: dx_vec(:)
      real(RD), intent(out)          :: x_adj(:)
      integer, intent(out), optional :: stat

      integer              :: c_vec(self%rank)
      integer              :: r
      character(LABEL_LEN) :: label
      real(RD)             :: x
      integer              :: c_lin
      real(RD)             :: x_bdy(2,self%rank)
      real(RD)             :: t(2,self%rank)
      integer              :: i(2)

      @:CHECK_BOUNDS(SIZE(x_vec), self%rank)
      @:CHECK_BOUNDS(SIZE(dx_vec), self%rank)
      @:CHECK_BOUNDS(SIZE(x_adj), self%rank)

      ! Adjust x_vec in the direction dx_vec, until it falls within a
      ! valid cell (one with c_lin_mask set)

      ! Check the direction vector

      if (ALL(dx_vec == 0._RD)) then

         if (PRESENT(stat)) then
            stat = STAT_INVALID_ARGUMENT
            return
         else
            @:ABORT('dx_vec cannot be null')
         end if

      end if

      ! Determine the starting cell

      call self%locate_(x_vec, c_vec, stat)
      if (PRESENT(stat)) then
         if (stat /= STAT_OK) return
      end if

      ! If on a cell boundary, adjust c_vec so that dx_vec points into
      ! the cell

      do r = 1, self%rank

         call self%axes(r)%fetch(c_vec(r), x)

         if (x_vec(r) == x) then
            if (dx_vec(r) < 0._RD .OR. (dx_vec(r) == 0._RD .AND. c_vec(r) == self%shape(r))) then
               c_vec(r) = c_vec(r) - 1
            end if
         end if

         if (c_vec(r) < 1) then
            if (PRESENT(stat)) then
               stat = STAT_OUT_OF_BOUNDS_AXIS_LO
               return
            else
               call self%axes(r)%get_label(label)
               @:ABORT('out-of-bounds (lo) on '//label//' axis')
            end if
         elseif (c_vec(r) >= self%shape(r)) then
            if (PRESENT(stat)) then
               stat = STAT_OUT_OF_BOUNDS_AXIS_HI
               return
            else
               call self%axes(r)%get_label(label)
               @:ABORT('out-of-bounds (hi) on '//label//' axis')
            end if
         end if

      end do

      ! Move from cell to cell within the grid, updating x_adj and
      ! c_vec until a valid cell is found

      x_adj = x_vec

      do

         c_lin = self%c_lin_vec_(c_vec)

         ! Finish if the cell is valid

         if (self%c_lin_mask(c_lin)) exit

         ! Get the bounds of the current cell

         do r = 1, self%rank
            call self%axes(r)%fetch(c_vec(r), x_bdy(1,r))
            call self%axes(r)%fetch(c_vec(r)+1, x_bdy(2,r))
         end do

         !  Find the intersections of the ray x = x_adj + t*dx_vec with
         !  these bounds

         where (dx_vec /= 0._RD)
            t(1,:) = (x_bdy(1,:) - x_adj)/dx_vec
            t(2,:) = (x_bdy(2,:) - x_adj)/dx_vec
         elsewhere
            t(1,:) = SIGN(HUGE(0._RD), x_bdy(1,:)-x_vec)
            t(2,:) = SIGN(HUGE(0._RD), x_bdy(2,:)-x_vec)
         end where

         ! Choose the first intersection along the ray (t > 0)

         i = MINLOC(t, t > 0._RD)

         @:ASSERT(ALL(i /= 0), 'No intersection found')

         ! Update x_adj and c_vec, being careful with rounding

         x_adj = x_adj + t(i(1),i(2))*dx_vec

         do r = 1, self%rank

            if ((i(1) == 1 .AND. i(2) == r) .OR. x_adj(r) < x_bdy(1,r)) then

               x_adj(r) = x_bdy(1,r)
               c_vec(r) = c_vec(r) - 1

            elseif ((i(1) == 2 .AND. i(2) == r) .OR. x_adj(r) > x_bdy(2,r)) then

               x_adj(r) = x_bdy(2,r)
               c_vec(r) = c_vec(r) + 1
               
            end if

         end do

         ! Check if we've gone off the grid

         do r = 1, self%rank

            if (c_vec(r) < 1) then
               if (PRESENT(stat)) then
                  stat = STAT_OUT_OF_BOUNDS_AXIS_LO
                  return
               else
                  call self%axes(r)%get_label(label)
                  @:ABORT('out-of-bounds (lo) on '//label//' axis')
               endif
            endif

            if (c_vec(r) >= self%shape(r)) then
               if (PRESENT(stat)) then
                  stat = STAT_OUT_OF_BOUNDS_AXIS_HI
                  return
               else
                  call self%axes(r)%get_label(label)
                  @:ABORT('out-of-bounds (hi) on '//label//' axis')
               endif
            endif

         end do

      end do

      ! Finish

      if (PRESENT(stat)) stat = STAT_OK

      return

   end subroutine adjust_x_vec
         
   !****

   subroutine locate_(self, x_vec, c_vec, stat)

      class(vgrid_t), intent(in)     :: self
      real(RD), intent(in)           :: x_vec(:)
      integer, intent(out)           :: c_vec(:)
      integer, intent(out), optional :: stat

      integer              :: r
      character(LABEL_LEN) :: label

      @:CHECK_BOUNDS(SIZE(x_vec), self%rank)
      @:CHECK_BOUNDS(SIZE(c_vec), self%rank)

      ! Locate x_vec in the grid, returning the vector cell index

      do r = 1, self%rank

         call self%axes(r)%locate(x_vec(r), c_vec(r))

         if (c_vec(r) < 1) then
            if (PRESENT(stat)) then
               stat = STAT_OUT_OF_BOUNDS_AXIS_LO
               return
            else
               call self%axes(r)%get_label(label)
               @:ABORT('out-of-bounds (lo) on '//label//' axis')
            endif
         endif

         if (c_vec(r) > self%shape(r)) then
            if (PRESENT(stat)) then
               stat = STAT_OUT_OF_BOUNDS_AXIS_HI
               return
            else
               call self%axes(r)%get_label(label)
               @:ABORT('out-of-bounds (hi) on '//label//' axis')
            endif
         endif

      end do

      ! Finish

      if (PRESENT(stat)) stat = STAT_OK

      return

   end subroutine locate_

   !****

   function c_lin_vec_(self, c_vec) result(c_lin)

      class(vgrid_t), intent(in) :: self
      integer, intent(in)        :: c_vec(:)
      integer                    :: c_lin

      integer :: p
      integer :: r

      @:CHECK_BOUNDS(SIZE(c_vec), self%rank)

      ! Calculate a linear cell index from a vector cell index

      c_lin = 1
      p = 1

      do r = 1, self%rank
         c_lin = c_lin + (c_vec(r)-1)*p
         p = p*(self%shape(r)-1)
      end do

      ! Finish

      return

   end function c_lin_vec_

   !****

   function c_vec_lin_(self, c_lin) result(c_vec)

      class(vgrid_t), intent(in) :: self
      integer, intent(in)        :: c_lin
      integer                    :: c_vec(self%rank)

      integer :: p
      integer :: q
      integer :: r

      ! Calculate a vector cell index from a linear cell index

      p = PRODUCT(self%shape-1)
      q = c_lin - 1

      do r = self%rank, 1, -1
         p = p/(self%shape(r)-1)
         c_vec(r) = q/p + 1
         q = q - (c_vec(r)-1)*p
      end do

      ! Finish

      return

   end function c_vec_lin_

   !****

   function v_lin_vec_(self, v_vec) result(v_lin)

      class(vgrid_t), intent(in) :: self
      integer, intent(in)        :: v_vec(:)
      integer                    :: v_lin

      integer :: p
      integer :: r

      @:CHECK_BOUNDS(SIZE(v_vec), self%rank)

      ! Calculate a linear vertex index from a vector vertex index

      v_lin = 1
      p = 1

      do r = 1, self%rank
         v_lin = v_lin + (v_vec(r)-1)*p
         p = p*self%shape(r)
      end do

      ! Finish

      return

   end function v_lin_vec_

   !****

   function v_vec_lin_(self, v_lin) result(v_vec)

      class(vgrid_t), intent(in) :: self
      integer, intent(in)        :: v_lin
      integer                    :: v_vec(self%rank)

      integer :: p
      integer :: q
      integer :: r

      ! Calculate a vector vertex index from a linear vertex index

      p = PRODUCT(self%shape)
      q = v_lin - 1

      do r = self%rank, 1, -1
         p = p/self%shape(r)
         v_vec(r) = q/p + 1
         q = q - (v_vec(r)-1)*p
      end do

      ! Finish

      return

   end function v_vec_lin_

   !****

   function v_lin_cnr_(self, i, c_vec) result(v_lin)

      class(vgrid_t), intent(in) :: self
      integer, intent(in)        :: i
      integer, intent(in)        :: c_vec(:)
      integer                    :: v_lin

      integer :: r
      integer :: v_vec(self%rank)

      @:CHECK_BOUNDS(SIZE(c_vec), self%rank)

      ! Calculate a linear vertex index from a corner index
      ! (i=1,...,2**rank) and a vector cell index

      ! First calculate the vector vertex index

      do r = 1, self%rank
         v_vec(r) = c_vec(r) + IBITS(i-1, r-1, 1)
      end do

      ! Convert to a linear index

      v_lin = self%v_lin_vec_(v_vec)

      ! Finish

      return

   end function v_lin_cnr_

   !****

   function v_lin_stn_(self, i, c_vec) result(v_lin)

      class(vgrid_t), intent(in) :: self
      integer, intent(in)        :: i
      integer, intent(in)        :: c_vec(:)
      integer                    :: v_lin

      integer :: r
      integer :: v_vec(self%rank)

      @:CHECK_BOUNDS(SIZE(c_vec), self%rank)

      ! Calculate a linear vertex index from a stencil index
      ! (i=1,...,4**rank) and a vector cell index

      ! First calculate the vector vertex index

      do r = 1, self%rank
         v_vec(r) = c_vec(r) + IBITS(i-1, 2*(r-1), 2) - 1
      end do

      ! Convert to a linear index

      v_lin = self%v_lin_vec_(v_vec)

      ! Finish

      return

   end function v_lin_stn_

   !****
   
   subroutine setup_interp_(self, x_vec, c_vec, mask, vdelta, vu, stat)

      class(vgrid_t), intent(in)     :: self
      real(RD), intent(in)           :: x_vec(:)
      integer, intent(out)           :: c_vec(:)
      logical, intent(out)           :: mask(:)
      real(RD), intent(out)          :: vdelta(:,:)
      real(RD), intent(out)          :: vu(:)
      integer, intent(out), optional :: stat

      integer  :: c_lin
      integer  :: r
      real(RD) :: x_bdy(self%rank)
      logical  :: on_bdy(self%rank)
      integer  :: a
      integer  :: c_vec_adj(self%rank)
      integer  :: c_lin_adj
      integer  :: v_lin
      integer  :: i
      integer  :: j
      real(RD) :: x(4)

      @:CHECK_BOUNDS(SIZE(x_vec), self%rank)
      @:CHECK_BOUNDS(SIZE(c_vec), self%rank)
      @:CHECK_BOUNDS(SIZE(mask), self%n_stn)

      @:CHECK_BOUNDS(SIZE(vdelta, 1), 3)
      @:CHECK_BOUNDS(SIZE(vdelta, 2), self%rank)

      @:CHECK_BOUNDS(SIZE(vu), self%rank)

      ! Locate the initial interpolation cell

      call self%locate_(x_vec, c_vec, stat)
      if (PRESENT(stat)) then
         if (stat /= STAT_OK) return
      end if

      c_vec = MIN(c_vec, self%shape-1)

      ! See if the cell has sufficient data (i.e., all vertices
      ! defined) to do an interpolation

      c_lin = self%c_lin_vec_(c_vec)

      if (.NOT. self%c_lin_mask(c_lin)) then

         ! If not, then if the point is on a boundary of the cell, see
         ! if one of the adjacent (toward-the-origin) cells can be
         ! used

         ! First, determine which boundary(ies) the point is on

         do r = 1, self%rank
            call self%axes(r)%fetch(c_vec(r), x_bdy(r))
         end do

         on_bdy = x_vec == x_bdy

         ! Iterate over adjacent cells

         adjacent_loop : do a = 1, 2**self%rank-1

            ! Set up the vector index for the adjacent cell

            do r = 1, self%rank

               if (BTEST(a, r-1)) then

                  if (.NOT. on_bdy(r)) cycle adjacent_loop
                  if (c_vec(r) == 1) cycle adjacent_loop

                  c_vec_adj(r) = c_vec(r) - 1

               else

                  c_vec_adj(r) = c_vec(r)

               end if

            end do

            ! Check whether the adjacent cell will do

            c_lin_adj = self%c_lin_vec_(c_vec_adj)

            if (self%c_lin_mask(c_lin_adj)) exit adjacent_loop

         end do adjacent_loop

         ! Check whether a suitable adjacent cell was found

         if (a < 2**self%rank) then

            ! Update the cell index

            c_vec = c_vec_adj

         else

            if (PRESENT(stat)) then
               stat = STAT_UNAVAILABLE_DATA
               return
            else
               @:ABORT('unavailable data')
            endif

         end if

      endif

      ! Set up the stencil mask

      do i = 1, self%n_stn

         v_lin = self%v_lin_stn_(i, c_vec)

         mask(i) = v_lin >= 1 .AND. v_lin <= self%n_v_lin

         if (mask(i)) then
            mask(i) = mask(i) .AND. self%v_lin_mask(v_lin)
         end if

      end do

      ! Set up the vdelta and vu values

      do r = 1, self%rank

         do j = 1, 4

            i = c_vec(r) + j - 2

            if (i < 1) then
               x(j) = -HUGE(0._RD)
            elseif (i > self%shape(r)) then
               x(j) = HUGE(0._RD)
            else
               call self%axes(r)%fetch(i, x(j))
            endif

         end do

         vdelta(:,r) = x(2:4) - x(1:3)

         vu(r) = (x_vec(r) - x(2))/vdelta(2,r)

      end do

      ! Finish

      if (PRESENT(stat)) stat = STAT_OK

      return

   end subroutine setup_interp_

   !****

   subroutine read(self, hdf5io, stat)

      class(vgrid_t), intent(out)    :: self
      type(hdf5io_t), intent(inout)  :: hdf5io
      integer, intent(out), optional :: stat

      integer :: revision

      ! Read the vgrid_t

      call check_type(hdf5io, 'vgrid_t', stat)
      if (PRESENT(stat)) then
         if (stat /= STAT_OK) return
      end if

      call hdf5io%read_attr('REVISION', revision)

      select case(revision)
      case(1)
         call read_rev1_(self, hdf5io, stat)
      case(2)
         call read_rev2_(self, hdf5io, stat)
      case(3)
         call read_rev3_(self, hdf5io, stat)
      case default
         if (PRESENT(stat)) then
            stat = STAT_INVALID_GROUP_REVISION
            return
         else
            @:ABORT('invalid group revision')
         end if
      end select

      ! Finish

      return

   contains

      subroutine read_rev1_(self, hdf5io, stat)

         class(vgrid_t), intent(out)    :: self
         type(hdf5io_t), intent(inout)  :: hdf5io
         integer, intent(out), optional :: stat

         integer                   :: rank
         logical, allocatable      :: mask(:)
         type(axis_t), allocatable :: axes(:)
         integer                   :: r
         type(hdf5io_t)            :: hdf5io_axis

         call hdf5io%read_attr('rank', rank)

         allocate(axes(rank))

         axis_loop : do r = 1, rank
            hdf5io_axis = hdf5io_t(hdf5io, axis_group_name_rev1_(r))
            call axes(r)%read(hdf5io_axis, stat)
            call hdf5io_axis%final()
            if (PRESENT(stat)) then
               if (stat /= STAT_OK) return
            endif
         end do axis_loop

         call hdf5io%alloc_read_dset('mask', mask)

         select type(self)
         type is(vgrid_t)
            self = vgrid_t_mask_(axes, mask)
         class default
            @:ABORT('invalid type')
         end select

         if (PRESENT(stat)) stat = STAT_OK

         return

      end subroutine read_rev1_

      !****

      function axis_group_name_rev1_(r) result (name)

         integer, intent(in) :: r
         character(256)      :: name

         write(name, 100) r
100      format('axes(',I0,')')

      end function axis_group_name_rev1_

      !****

      subroutine read_rev2_(self, hdf5io, stat)

         class(vgrid_t), intent(out)    :: self
         type(hdf5io_t), intent(inout)  :: hdf5io
         integer, intent(out), optional :: stat

         integer                   :: rank
         logical, allocatable      :: mask(:)
         type(axis_t), allocatable :: axes(:)
         integer                   :: r
         type(hdf5io_t)            :: hdf5io_axis

         call hdf5io%read_attr('rank', rank)

         allocate(axes(rank))

         axis_loop : do r = 1, rank
            hdf5io_axis = hdf5io_t(hdf5io, axis_group_name_(r))
            call axes(r)%read(hdf5io_axis, stat)
            call hdf5io_axis%final()
            if (PRESENT(stat)) then
               if (stat /= STAT_OK) return
            endif
         end do axis_loop

         call hdf5io%alloc_read_dset('mask', mask)

         select type(self)
         type is(vgrid_t)
            self = vgrid_t_mask_(axes, mask)
         class default
            @:ABORT('invalid type')
         end select

         if (PRESENT(stat)) stat = STAT_OK

         return

      end subroutine read_rev2_

      !****

      subroutine read_rev3_(self, hdf5io, stat)

         class(vgrid_t), intent(out)    :: self
         type(hdf5io_t), intent(inout)  :: hdf5io
         integer, intent(out), optional :: stat

         integer                   :: rank
         integer, allocatable      :: v_lin_seq(:)
         type(axis_t), allocatable :: axes(:)
         integer                   :: r
         type(hdf5io_t)            :: hdf5io_axis

         call hdf5io%read_attr('rank', rank)

         allocate(axes(rank))

         axis_loop : do r = 1, rank
            hdf5io_axis = hdf5io_t(hdf5io, axis_group_name_(r))
            call axes(r)%read(hdf5io_axis, stat)
            call hdf5io_axis%final()
            if (PRESENT(stat)) then
               if (stat /= STAT_OK) return
            endif
         end do axis_loop

         call hdf5io%alloc_read_dset('v_lin_seq', v_lin_seq)

         select type(self)
         type is(vgrid_t)
            self = vgrid_t_map_(axes, v_lin_seq)
         class default
            @:ABORT('invalid type')
         end select

         if (PRESENT(stat)) stat = STAT_OK

         return

      end subroutine read_rev3_

   end subroutine read

   !****

   subroutine write(self, hdf5io, stat)

      class(vgrid_t), intent(in)     :: self
      type(hdf5io_t), intent(inout)  :: hdf5io
      integer, intent(out), optional :: stat

      integer        :: r
      type(hdf5io_t) :: hdf5io_axis

      ! Write the vgrid
      
      call hdf5io%write_attr('TYPE', 'vgrid_t')
      call hdf5io%write_attr('REVISION', 3)

      call hdf5io%write_attr('shape', self%shape)
      call hdf5io%write_attr('rank', self%rank)

      axis_loop: do r = 1, self%rank
         hdf5io_axis = hdf5io_t(hdf5io, axis_group_name_(r))
         call self%axes(r)%write(hdf5io_axis, stat)
         call hdf5io_axis%final()
         if (PRESENT(stat)) then
            if (stat /= STAT_OK) return
         endif
      end do axis_loop

      call hdf5io%write_dset('v_lin_seq', self%v_lin_seq)

      ! Finish

      if (PRESENT(stat)) stat = STAT_OK

      return

   end subroutine write

   !****

   function axis_group_name_(r) result (name)

      integer, intent(in) :: r
      character(256)      :: name

      write(name, 100) r
100   format('axes[',I0,']')

   end function axis_group_name_

   !****

   subroutine optimize(x_vecs, mask)

      real(RD), intent(in) :: x_vecs(:,:)
      logical, intent(out) :: mask(:)

      integer                   :: rank
      character(1), allocatable :: labels(:)
      type(vgrid_t)             :: vgrid
      integer                   :: c_lin
      integer, allocatable      :: c_vec(:)
      integer                   :: i
      integer                   :: v_lin
      integer                   :: v_seq

      @:CHECK_BOUNDS(SIZE(mask), SIZE(x_vecs, 2))

      ! Optimize grid construction by returning a mask that indices
      ! which vertices in x_vecs are usable (the remainder do not
      ! belong to complete interpolation cells, and therefore can be
      ! discarded)

      ! Start by creating a dummy vgrid

      rank = SIZE(x_vecs, 1)

      allocate(labels(rank))
      labels = ''

      vgrid = vgrid_t_(x_vecs, labels)

      ! Update rank (to handle removal of degenerate axes)

      call vgrid%get_rank(rank)

      ! Loop through cells of the grid

      allocate(c_vec(rank))

      mask = .FALSE.

      do c_lin = 1, vgrid%n_c_lin

         ! Check if the cell is complete

         if (vgrid%c_lin_mask(c_lin)) then

            c_vec = vgrid%c_vec_lin_(c_lin)

            ! Iterate over corner vertices of the cell

            do i = 1, 2**rank

               v_lin = vgrid%v_lin_cnr_(i, c_vec)
               v_seq = vgrid%v_seq_lin(v_lin)

               ! Update the mask

               mask(v_seq) = .TRUE.

            end do

         endif

      end do

      ! Finish

      return

   end subroutine optimize

end module vgrid_m
