module table_2D_mod

 use flat_interpolator_mod

 implicit none
 
 integer, private, parameter :: NDIM = 2 

 type :: table_2D
   
   private   
   character(LEN=80)        :: name
   logical                  :: initialized = .false.
   integer                  :: nderiv
   integer                  :: nmax1
   integer                  :: nmax2
   integer, dimension(NDIM) :: nx
   logical, dimension(NDIM) :: log_spc
   real(8), dimension(NDIM) :: x0
   real(8), dimension(NDIM) :: xmax
   real(8), dimension(NDIM) :: dx
   logical, dimension(NDIM) :: zero = .false. ! Allows for zero point on logarithmic grids
                                              ! If true, the first position in the table 
                                              ! along a dimension corresponds to zero and 
                                              ! the second point has the value x0 (instead 
                                              ! of the first)
   real(8), dimension(:,:), allocatable :: tab
   real(8), dimension(:,:), allocatable :: xx
   type(flat_interpolator)  :: interpol  
 
 contains

   procedure :: init          => init_table_2D
   procedure :: get_xmin      => get_xmin_table_2D
   procedure :: get_xmax      => get_xmax_table_2D
   procedure :: load_slab_der => load_slab_der_table_2D 
   procedure :: load_slab_pos => load_slab_pos_table_2D 
   procedure :: interp        => interp_table_2D
 
 end type table_2D 
 
 private :: init_table_2D
 private :: get_xmin_table_2D
 private :: get_xmax_table_2D
 private :: load_slab_der_table_2D
 private :: load_slab_pos_table_2D
 private :: interp_table_2D
 private :: table_position

contains

  function get_xmin_table_2D(this,id) result(xmin)
    real(8) :: xmin
    class(table_2D), intent(in) :: this
    integer, intent(in) :: id
    if (id<=NDIM.and.id>0) then
      xmin = this%x0(id)
    else
      xmin = this%x0(NDIM)
    endif 
    return
  end function get_xmin_table_2D

  function get_xmax_table_2D(this,id) result(xmax)
    real(8) :: xmax
    class(table_2D), intent(in) :: this
    integer, intent(in) :: id
    if (id<=NDIM.and.id>0) then
      xmax = this%xmax(id)
    else
      xmax = this%xmax(NDIM)
    endif 
    return
  end function get_xmax_table_2D
  
  subroutine init_table_2D(this,nx,nderiv,x0,dx,log_spc,tab,zero)
    
    class(table_2D), intent(out) :: this
    integer, dimension(NDIM),      intent(in)  :: nx   
    integer, dimension(NDIM),      intent(in)  :: nderiv  
    real(8), dimension(NDIM),      intent(in)  :: x0,dx
    logical, dimension(NDIM),      intent(in), optional :: log_spc
    logical, dimension(NDIM),      intent(in), optional :: zero
    real(8), dimension(:,:,0:,0:), intent(in), optional :: tab

    integer :: ierr,i,j,i1,i2,i3,i4
    
    ! Set the array properties 
    this%nx      = nx
    this%x0      = x0
    this%dx      = dx
    this%nderiv  = MINVAL(nderiv(1:NDIM))
    if (present(zero).and.present(log_spc)) then
      do i=1,NDIM
        this%zero(i)    = zero(i) .and. log_spc(i)
        this%log_spc(i) = log_spc(i)
      enddo
    elseif (present(log_spc)) then
      this%log_spc = log_spc
      this%zero    = .false.
    else
      this%zero    = .false.
      this%log_spc = .false.
      this%xmax    = x0 + dx*(nx-1.d0)
    endif   
    
    ! Build tables of point values
    ! This is done only to prevent numerous 
    ! recalculations of 10.d0**x for 
    ! logarithmic tables (i.e. a hash is used 
    ! to find the table positions, not a table search).
    if (allocated(this%xx)) deallocate(this%xx)
    print *,1
    allocate(this%xx(MAXVAL(this%nx),NDIM))
    print *,2
    do i=1,NDIM
      if (this%log_spc(i).and.this%zero(i)) then
        this%xx(1,i) = 0.d0
        do j=2,this%nx(i)
          this%xx(j,i) = this%x0(i)*10.d0**(REAL(j-2)*this%dx(i))
        enddo  
      elseif (this%log_spc(i)) then
        do j=1,this%nx(i)
          this%xx(j,i) = this%x0(i)*10.d0**(REAL(j-1)*this%dx(i))
        enddo  
      else
        do j=1,this%nx(i)
          this%xx(j,i) = this%x0(i) + REAL(j-1)*this%dx(i)
        enddo  
      endif   
    enddo  

    ! Set the flat array sizes  
    this%nmax1 = (this%nderiv+1)**NDIM-1
    this%nmax2 = 1
    do i=1,NDIM
      this%nmax2 = this%nx(i)*this%nmax2
    enddo
    this%nmax2 = this%nmax2-1
    
    ! Allocate the table 
    if (allocated(this%tab)) deallocate(this%tab)
    allocate(this%tab(0:this%nmax1,0:this%nmax2))
   
    this%initialized = .true.
    
    ! Load table if present 
    if (present(tab)) then
      do i1=0,this%nderiv 
      do i2=0,this%nderiv 
        ierr = this%load_slab_der(tab(:,:,i1,i2),(/i1,i2/))
        if (ierr .ne. 0) print *, 'Table load error ',ierr
      enddo   
      enddo   
    endif


    return
      
  end subroutine 

  function load_slab_der_table_2D(this,dat,ideriv) result(ierr)
     
     integer :: ierr
 
     class(table_2D), intent(inout) :: this 
     real(8), dimension(:,:), intent(in) :: dat
     integer, dimension(:), intent(in)   :: ideriv
     
     integer :: i,id,ii,i1,i2,i3,i4

     ierr = 0

     ! Can't load an uninitialized table
     if (.not.this%initialized) then
       ierr = -1
       return
     endif     
     
     ! Check that the array is large enough
     if (SIZE(dat)<this%nmax2+1) then
       ierr = i
       return
     endif   

     ! Check derivative position specification
     if (SIZE(ideriv)<NDIM) then
       ierr = -2
       return
     endif   
     do i=1,NDIM
       if (ideriv(i)<0 .or. ideriv(i)>this%nderiv) then
         ierr = -i-2
         return
       endif   
     enddo 
     
     ! Put it in the table
     id = 0
     do i=NDIM,1,-1
       id = (this%nderiv+1)*id + ideriv(i)
     enddo

     do i1=0,this%nx(1)-1
     do i2=0,this%nx(2)-1
       ii = this%nx(1)*i2 + i1
       this%tab(id,ii) = dat(i1+1,i2+1)
     enddo
     enddo

     return 
  
  end function load_slab_der_table_2D

  function load_slab_pos_table_2D(this,dat,ipos) result(ierr)
     
     integer :: ierr
 
     class(table_2D), intent(inout) :: this 
     real(8), dimension(:,:), intent(in) :: dat
     integer, dimension(:), intent(in)   :: ipos
     
     integer :: i,id,ii,i1,i2,i3,i4

     ierr = 0

     ! Can't load an uninitialized table
     if (.not.this%initialized) then
       ierr = -1
       return
     endif     
     
     ! Check that the array is large enough
     if (SIZE(dat)<this%nmax1+1) then
       ierr = i
       return
     endif   

     ! Check derivative position specification
     if (SIZE(ipos)<NDIM) then
       ierr = -2
       return
     endif   
     do i=1,NDIM
       if (ipos(i)<1 .or. ipos(i)>this%nx(i)) then
         ierr = -i-2
         return
       endif   
     enddo 
     
     ! Put it in the table
     ii = 0
     do i=NDIM,1,-1
       ii = this%nx(i)*id + ipos(i) - 1
     enddo

     do i1=0,this%nderiv
     do i2=0,this%nderiv
       id = (this%nderiv+1)*i2 + i1
       this%tab(id,ii) = dat(i1+1,i2+1)
     enddo
     enddo

     return 
  
  end function load_slab_pos_table_2D
  
  function interp_table_2D(this,x,der,order) result(f)
    
    real(8) :: f
     
    class(table_2D)   :: this
    real(8)           :: x(NDIM),dd(NDIM),delta(NDIM)
    integer           :: idx(NDIM)
    integer, optional :: der(NDIM)
    integer, optional :: order

    integer :: ierr,ipos(0:2**NDIM-1)
    integer :: ii,i1,i2,i3,i4

    if (present(order)) then
      print *,'Not allowing you to choose the order for now'
    endif
    
    if (this%nderiv==0) then
      ierr = this%interpol%set_order('linear')  
    elseif (this%nderiv==1) then
      ierr = this%interpol%set_order('cubic') 
    else
      ierr = this%interpol%set_order('quintic')
    endif
   
    call table_position(this,x,dd,idx,delta)

!    write(*,*) "luketbl: dd = ", dd
!    write(*,*) "luketbl: delta = ", delta
    do i1=0,1
    do i2=0,1
      ii = 2*i2 + i1
      ipos(ii) =                       this%nx(1)*(idx(2)-1 + i2) &
                +                                 (idx(1)-1 + i1)
    enddo
    enddo

    if (present(der)) then
      f = this%interpol%f2D(dd,delta,this%tab(:,ipos),der,this%nderiv+1)
    else
      f = this%interpol%f2D(dd,delta,this%tab(:,ipos),(/0,0/),this%nderiv+1)
    endif
    
    return

  end function interp_table_2D

  subroutine table_position(this,x,dd,idx,delta)
    
    class(table_2D),  intent(in)  :: this
    real(8), dimension(NDIM), intent(in)  :: x
    real(8), dimension(NDIM), intent(out) :: dd
    real(8), dimension(NDIM), intent(out), optional :: delta
    integer, dimension(NDIM), intent(out) :: idx
    
    real(8), dimension(NDIM) :: xt
    integer :: i 
    real(8) :: xl,xh,x0
    integer, save :: nprint = 0

    ! Transform logarithmic dimensions
    do i=1,NDIM
      if (this%log_spc(i)) then
        xt(i) = log10(x(i))
      else
        xt(i) = x(i)
      endif   
    enddo 
    
    ! Hash into table
    do i=1,NDIM

      if (this%log_spc(i)) then
        x0 = log10(this%x0(i))
      else
        x0 = this%x0(i)
      endif   

      if (xt(i)<x0) then
        idx(i) = 1
        if (this%log_spc(i) .and. this%zero(i)) then
          continue
        elseif (this%log_spc(i)) then
          if (nprint<20) write(6,'(a,i2,a,2es12.3)')'Value ',i,' off bottom of table.',10.d0**xt(i),10.d0**x0
        else
          if (nprint<20) write(6,'(a,i2,a,2es12.3)')'Value ',i,' off bottom of table.',xt(i),x0
        endif 
        nprint = nprint + 1
      elseif (xt(i)==(x0 + (this%nx(i)-1)*this%dx(i))) then
        idx(i) = this%nx(i)-1
      elseif (xt(i)>(x0 + (this%nx(i)-1)*this%dx(i))) then
        idx(i) = this%nx(i)-1
        if (nprint<20) write(6,'(a,i2,a,2f5.1,i3)')'Value ',i,' off top of table.',xt(i),this%dx(i),this%nx(i)
        nprint = nprint + 1
      elseif (xt(i).ne.xt(i)) then
        idx(i) = 1
        if (nprint<20) write(6,'(a,i2,a,es12.3)')'Value ',i,' is NaN.',xt(i)
        nprint = nprint + 1
      else  
        idx(i) = INT(FLOOR((xt(i)-x0)/this%dx(i))) + 1
      endif

    enddo   
        
    ! Calculate fractional distances
    do i=1,NDIM
      if (idx(i)==1 .and. this%log_spc(i) .and. this%zero(i)) then
        dd(i) = x(i)/this%x0(i)  
        if (present(delta)) delta(i) = this%x0(i)
      elseif (this%log_spc(i) .and. this%zero(i)) then
        idx(i) = idx(i) + 1 ! Set to the correct grid position
        xl     = this%xx(idx(i),i)
        xh     = this%xx(idx(i)+1,i)
        dd(i)  = (x(i)-xl)/(xh-xl)
        if (present(delta)) delta(i) = xh-xl
      elseif (this%log_spc(i)) then
        xl = this%xx(idx(i),i)
        xh = this%xx(idx(i)+1,i)
        dd(i) = (x(i)-xl)/(xh-xl)
        if (present(delta)) delta(i) = xh-xl
      else  
        dd(i) = (x(i)-(idx(i)-1)*this%dx(i) - this%x0(i))/this%dx(i)
        if (present(delta)) delta(i) = this%dx(i)
      endif
      if (dd(i).ne.dd(i)) dd(i) = 0.d0
      dd(i) = min(dd(i),1.0)
      dd(i) = max(dd(i),0.0)
    enddo   
    
    return    
     
  end subroutine table_position 

end module table_2D_mod