module indexer

! Indexer class. Useful to flatten a multi-dimensional indexed expansion
! onto a flat data structure. 

implicit none
character(*), parameter, private ::  version = &
     & '$Id: indexer.f90,v 1.7 2003/12/15 16:52:21 pletzer Exp $'

type indexer_obj
   integer counter ! flat array index
   integer ntot    ! flat array total size
   integer ndims   ! no of dimensions
   ! the following contain the sizes and starting indices 
   ! in each dimension
   integer, pointer :: sizes(:), start(:)
end type indexer_obj

contains

  subroutine indexer_init(this, sizes, start, ier)

    ! Constructor. 

    implicit none
    type(indexer_obj) :: this
    integer, intent(inout) :: sizes(:) ! An array containing the sizes 
    integer, intent(in) :: start(:)
    integer, intent(out) :: ier

    integer jer, j
    ier = 0
    this % counter = 0
    this % ndims = size(sizes)
    this % ntot = 1
    sizes = max(1, sizes)
    do j = 1, this % ndims
       this % ntot = this % ntot * sizes(j)
    enddo
    allocate(this % sizes(this % ndims), stat=jer)
    allocate(this % start(this % ndims), stat=jer)
    if(jer /= 0) ier = 1
    this % sizes = sizes
    this % start = start
  end subroutine indexer_init

  subroutine indexer_free(this, ier)
    
    ! Destructor

    implicit none
    type(indexer_obj) :: this
    integer, intent(out) :: ier
    
    integer jer
    ier = 0
    deallocate(this % sizes, stat=jer)
    deallocate(this % start, stat=jer)
    if(jer /= 0) ier = 2
  end subroutine indexer_free

  subroutine indexer_next(this, ier)

    ! Increment counter

    implicit none
    type(indexer_obj) :: this
    integer, intent(out) :: ier

    ier = 0
    this % counter = this % counter + 1
    if(this % counter > this % ntot) ier = 3

  end subroutine indexer_next

  subroutine indexer_get(this, indices, ier)

    ! Return sub-indices

    implicit none
    type(indexer_obj) :: this
    integer, intent(out) :: indices(:)
    integer, intent(out) :: ier

    integer j1, j2, periodicity(this % ndims)
    ier = 0
    periodicity = 1
    do j1 = 2, this % ndims
       periodicity(j1) = periodicity(j1-1) * this % sizes(j1-1)
    enddo
    indices = (this % counter - 1)/periodicity
    indices = mod(indices, this % sizes) + this % start
  end subroutine indexer_get

  subroutine indexer_save(this, filename, ier)

    ! Save state in netCDF file 'filename'

    use ezcdf
    implicit none
    type(indexer_obj) :: this
    character(*), intent(in) :: filename
    integer, intent(out) :: ier

    integer ncid, jer
    ier = 0
    call cdf_open(ncid, filename, 'w', jer)
    call cdf_define(ncid, 'version', (/len_trim(version),0,0/), 'CHAR', jer)
    call cdf_define(ncid, 'counter', (/0,0,0/), 'INT', jer)
    call cdf_define(ncid, 'ndims', (/0,0,0/), 'INT', jer)
    call cdf_define(ncid, 'sizes', (/this % ndims, 0, 0/), 'INT', jer)
    call cdf_define(ncid, 'start', (/this % ndims, 0, 0/), 'INT', jer)
    call cdf_write(ncid, 'version', version, jer)
    call cdf_write(ncid, 'counter', this % counter, jer)
    call cdf_write(ncid, 'ndims', this % ndims, jer)
    call cdf_write(ncid, 'sizes', this % sizes, jer)
    call cdf_write(ncid, 'start', this % start, jer)
    call cdf_close(ncid, jer) 
    if(jer /= 0) ier = 4
  end subroutine indexer_save

  subroutine indexer_load(this, filename, ier)

    ! Load state from netCDF file 'filename'

    use ezcdf
    implicit none
    type(indexer_obj) :: this
    character(*), intent(in) :: filename
    integer, intent(out) :: ier

    integer ncid, j, jer
    ier = 0
    call cdf_open(ncid, filename, 'r', jer)
    call cdf_read(ncid, 'counter', this % counter, jer)
    call cdf_read(ncid, 'ndims', this % ndims, jer)
    allocate(this % sizes(this % ndims))
    allocate(this % start(this % ndims))
    call cdf_read(ncid, 'sizes', this % sizes, jer)
    call cdf_read(ncid, 'start', this % start, jer)
    call cdf_close(ncid, jer)
    if(jer /= 0) ier = 5
    this % ntot = 1
    do j = 1, this % ndims
       this % ntot = this % ntot * this % sizes(j)
    enddo
  end subroutine indexer_load

  subroutine indexer_error(this, nunit, ier)
    
    ! Print error message

    implicit none
    type(indexer_obj) :: this
    integer, intent(in) :: nunit ! output channel number
    integer, intent(in) :: ier

    character(9), parameter :: header = 'Indexer::'

    select case(ier)
       case(1)
          write(nunit, *) header//'ERROR while allocating'
       case(2)
          write(nunit, *) header//'ERROR while deallocating'
       case(3)
          write(nunit, '(a,i6)') &
               & header//'WARNING counter exceeded declared size ', this % ntot
       case(4)
          write(nunit, *) header//'ERROR while saving'
       case(5)
          write(nunit, *) header//'ERROR while loading'
       case(6)
          write(nunit, *) header//'ERROR occurred in test unit'
       case default
    end select    
  end subroutine indexer_error

  subroutine indexer_test(this, ier)

    ! Test unit

    implicit none
    type(indexer_obj) :: this
    integer, intent(out) :: ier

    integer, allocatable :: sizes(:), start(:), indices(:)
    integer jer, ndims, Js(3), Ls(3), j
    ier = 0

    print *,'testing indexer ', trim(version)
    ! First index represents (x, y, z) field component
    ! Second 3 indices are the Fourier modes which run from -J ... +J
    ! Third 3 indices are spatial indices, running from 0...L

    ndims = 1 + 3 + 3
    allocate(sizes(ndims), start(ndims), indices(ndims))
    sizes(1) = 3 
    start(1) = 1
    Js = (/1,2,3/) ! Fourier
    sizes(2:4) = 2*Js + 1
    start(2:4) = -Js
    Ls = (/2,3,1/) ! spatial 
    sizes(5:7) = Ls+1
    start(5:7) = 0
    call indexer_init(this, sizes, start, jer)
    call indexer_error(this, 6, jer)
    do j = 1, this % ntot
       call indexer_next(this, jer)
       call indexer_error(this, 6, jer)
       call indexer_get(this, indices, jer)
       print *,'Flat array index=',this % counter, &
            & ' Max size=', this  % ntot, &
            & ' Sub-indices=', indices
       call indexer_error(this, 6, jer)       
    enddo
    call indexer_save(this, 'indexer_test.nc', jer)
    call indexer_error(this, 6, jer)
    call indexer_free(this, jer)
    call indexer_error(this, 6, jer)
    call indexer_load(this, 'indexer_test.nc', jer)
    call indexer_error(this, 6, jer)
    call indexer_save(this, 'indexer_test2.nc', jer)
    call indexer_error(this, 6, jer)
    if(jer /= 0) ier = 6
    
    deallocate(sizes, start, indices)
  end subroutine indexer_test  

end module indexer
