!
module histogram_mod
use bmad
implicit none

contains

! find the histogram
! ----------------------------------------------------------------
subroutine hist_real(array, sort_index, binsize, h, z, rev, binmin, binmax)
  !   usage:
  !       hist_real(array, sort_index, binsize, h, rev [, binmin, binmax])
  !   inputs:
  !       array: A 1-d array
  !       sort_index: A sort index into the array
  !       binsize: The bin size for the histogram
  !   optional inputs:
  !       binmin, binmax: The min,max values to consider. Default is
  !           the min and max of the input array
  !   outputs:
  !       h: The histogram
  !       z: The bined value
  !   optional outputs:
  !       rev: The reverse indices.
    
  real(rp), intent(in), dimension(:) :: array
  integer,  intent(in), dimension(:) :: sort_index
  real(rp), intent(in)               :: binsize
  
  integer, intent(inout), dimension(:), allocatable            :: h
  integer, intent(inout), dimension(:), allocatable, optional  :: rev
  real(rp), intent(inout), dimension(:), allocatable            :: z
  
  real(rp), optional :: binmin
  real(rp), optional :: binmax
  real(rp) mbinmin, mbinmax
    
  integer nbin
  logical dorev
    
  integer binnum, binnum_old, array_index, tbin, offset
  integer i

  real(rp) bininv
    
  if (size(array) /= size(sort_index)) then
     print *, "array and sort index must be same size!"
     stop
  endif

  dorev = .false. 
  if (present(rev)) then
     dorev=.true.
  endif
  
  if (present(binmin)) then
     mbinmin = binmin
  else
     mbinmin = array(sort_index(1))
  endif
  if (present(binmax)) then
     mbinmax = binmax
  else
     mbinmax = array(sort_index(size(array)))
  endif
  
  bininv = 1.0/binsize
  nbin = int( (mbinmax-mbinmin)*bininv) + 1
  
  ! allocate the outputs
  if (allocated(h)) then 
     deallocate(h)
  end if
  if (allocated(z)) then
    deallocate(z)
  end if
  allocate(h(nbin)); h=0
  allocate(z(nbin)); z=0.0_rp
  h=0
  z=0.0_rp
  do i=1,nbin
     z(i)=mbinmin+(i-1)*binsize
  end do
  z(nbin)=mbinmax
  
  if (dorev) then
     allocate(rev(size(array) + nbin + 1)); rev=0
     rev=(size(rev)+1)
  endif

  binnum_old = 0
  do i=1,size(array)
     array_index = sort_index(i)
     
     ! offset into rev
     offset = i + nbin + 1
     if (dorev) then
        rev(offset) = array_index
     endif

     binnum = int( (array(array_index)-mbinmin)*bininv ) + 1
     
     if ( (binnum >= 1) .and. (binnum <= nbin) ) then
        
        ! should we update the reverse indices?
        if (dorev .and. (binnum > binnum_old)) then
           tbin = binnum_old + 1
           do while(tbin <= binnum)
              rev(tbin) = offset
              tbin = tbin + 1
           end do
        endif

        ! update the histogram
        h(binnum) = h(binnum) + 1
        binnum_old = binnum
        
     endif
  end do
  
end subroutine hist_real

! interpolation integer
! ----------------------------------------------------------------
subroutine interp1d(z,h,z0,h0)
  integer,  intent(in), dimension(:), allocatable            :: h
  real(rp), intent(in), dimension(:), allocatable            :: z
  real(rp), intent(in) :: z0
  real(rp), intent(out) ::  h0
  
  real(rp), allocatable :: loc_comp_array(:)
  integer i,n,m
  
  if (size(h) /= size(z)) then
     print *, "hist and z must be same size!"
     stop
  endif

  h0=0.0_rp
  
  if (z0 > maxval(z) .or. z0 < minval(z)) then
     print *, "input z0 is out of range of z"
     stop
  end if
  
  if(allocated(loc_comp_array)) then
     deallocate(loc_comp_array)
  end if
  allocate(loc_comp_array(size(z))); loc_comp_array=0.0_rp

  do i=1, size(z)
     loc_comp_array(i) = abs(z(i) - z0)
  end do
  
  n=minloc(loc_comp_array,dim=1)

  if (n .eq. 1 .or. n .eq. size(z)) then
     h0=h(n)
  else
     if (loc_comp_array(n+1) .lt. loc_comp_array(n-1)) then
        m = n + 1
     else
        m = n - 1
     endif
     h0 = (h(n)*loc_comp_array(m)+h(m)*loc_comp_array(n))/(loc_comp_array(n)+loc_comp_array(m))
  end if

end subroutine interp1d

! interpolation real number
! ----------------------------------------------------------------
subroutine interp1d_real(z,h,z0,h0)
  real(rp), intent(in), dimension(:), allocatable            :: h
  real(rp), intent(in), dimension(:), allocatable            :: z
  real(rp), intent(in) :: z0
  real(rp), intent(out) ::  h0
  
  real(rp), allocatable :: loc_comp_array(:)
  integer i,n,m
  
  if (size(h) /= size(z)) then
     print *, "hist and z must be same size!"
     stop
  endif

  h0=0.0_rp
  
  if (z0 > maxval(z) .or. z0 < minval(z)) then
     print *, "input z0 is out of range of z: ", z0
!     print *, "maxz: ", maxval(z), "  minz: ", minval(z)
     return
  end if
  
  if(allocated(loc_comp_array)) then
     deallocate(loc_comp_array)
  end if
  allocate(loc_comp_array(size(z))); loc_comp_array=0.0_rp

  do i=1, size(z)
     loc_comp_array(i) = abs(z(i) - z0)
  end do
  
  n=minloc(loc_comp_array,dim=1)

  if (n .eq. 1 .or. n .eq. size(z)) then
     h0=h(n)
  else
     if (loc_comp_array(n+1) .lt. loc_comp_array(n-1)) then
        m = n + 1
     else
        m = n - 1
     endif
     h0 = (h(n)*loc_comp_array(m)+h(m)*loc_comp_array(n))/(loc_comp_array(n)+loc_comp_array(m))
  end if

end subroutine interp1d_real

! ----------------------------------------------------------------
! interpolation 1D real linear array
subroutine interp1d_real_linear(z,h,z0,h0)

  real(rp), intent(in), dimension(:), allocatable            :: h
  real(rp), intent(in), dimension(:), allocatable            :: z
  real(rp), intent(in) :: z0
  real(rp), intent(out) ::  h0

  integer m

  if (size(h) /= size(z)) then
     print *, "hist and z must be same size!"
     stop
  endif

  h0=0.0_rp
  
  call find_index_linear(z, z0, m)
  if (m == -1 ) return
  h0= (h(m+1) - h(m) ) / (z(m+1) - z(m)) *(z0 - z(m)) + h(m)

end subroutine interp1d_real_linear

! ----------------------------------------------------------------
! interpolation 2D real, bilinear interpolation

subroutine interp2d_real(x,y,z,x0,y0,z0)

  real(rp), intent(in), dimension(:), allocatable            :: x,y
  real(rp), intent(in), dimension(:,:), allocatable            :: z
  real(rp), intent(in) :: x0,y0
  real(rp), intent(out) ::  z0

  real(rp) z1, z2, z3, z4, t, u
  integer  m, n

  if (size(x) /= size(z(:,1)) .or. size(y) /= size(z(1,:)) ) then
     print *, "x and y zie must match z size!"
     stop
  endif

! General arrays, a little slower
! call find_index (x, x0, m) 
! call find_index (y, y0, n)  

  call find_index_linear (x, x0, m)
  call find_index_linear (y, y0, n)
  if (m == -1 .or. n == -1) then
     z0 = 0.0_rp
     return
  end if

!  print *,'m: ', m, ' n: ', n

  z1 = z(m,n)
  z2 = z(m+1,n)
  z3 = z(m+1,n+1)
  z4 = z(m,n+1)

  t= (x0-x(m)) / (x(m+1)-x(m))
  u= (y0-y(n)) / (y(n+1)-y(n))

  z0= (1-t)*(1-u)*z1 + t*(1-u)*z2 + t*u*z3 + (1-t)*u*z4
!  print *, 'x0: ', x0, ' y0: ', y0, ' z0: ', z0

end subroutine interp2d_real

! ----------------------------------------------------------------
! find the lower index of z0 in the array z

subroutine find_index (z, z0, idx)

  real(rp), intent(in), dimension(:), allocatable            :: z
  real(rp), intent(in) :: z0
  integer, intent(out) ::  idx
  
  real(rp), allocatable :: loc_comp_array(:)
  integer i, n

  
  if (z0 > maxval(z) .or. z0 < minval(z)) then
     print *, "input z0 is out of range of z: ", z0
     idx = 0
     return
  end if
  
  if(allocated(loc_comp_array)) then
     deallocate(loc_comp_array)
  end if
  allocate(loc_comp_array(size(z))); loc_comp_array=0.0_rp

  do i=1, size(z)
     loc_comp_array(i) = abs(z(i) - z0)
  end do
  
  n = minloc(loc_comp_array,dim=1)

  if (n .eq. size(z)) then
     idx = size(z)-1
  else
     if (loc_comp_array(n+1) .lt. loc_comp_array(n-1)) then
        idx = n
     else
        idx = n -1
     endif
  end if
 

end subroutine find_index


! ----------------------------------------------------------------
! find the index of z0 in a linear spaced array z

subroutine find_index_linear (z, z0, idx)

  real(rp), intent(in), dimension(:), allocatable            :: z
  real(rp), intent(in) :: z0
  integer, intent(out) ::  idx

  real(rp) d_z
  integer n

  if (z0 > maxval(z) .or. z0 < minval(z)) then
!     print *, "input z0 is out of range of z: ", z0
     idx = -1
     return
  end if

  d_z = ( z(size(z)) - z(1) ) / (size(z) - 1)

  idx = floor( ( z0 - z(1) ) / d_z ) + 1 
!  print *, 'd_z: ', d_z, ' idx: ', idx
  
end subroutine find_index_linear

! ----------------------------------------------------------------
! Extract the Ex file to seperate arrays

subroutine extract_2darrays_to_arrays(Ex_array, m, n, x, y, z)
  
  real(rp), intent(in),  dimension(:,:), allocatable  :: Ex_array
  integer, intent(in) :: m, n

  real(rp), intent(out), dimension(:), allocatable    :: x,y
  real(rp), intent(out), dimension(:,:), allocatable  :: z

  integer i, j, k

  if (size(Ex_array(1,:)) /=  3 ) then
     print *, "Not correct format of Ex_array! Stop!"
     stop 
  end if
  
  allocate(x(1:m))
  allocate(y(1:n))
  allocate(z(1:m,1:n))

  do j=1,m
     do k=1,n
        i = (j-1)*m + k
        z(j,k) = Ex_array(i,3)
        if (j .eq. 1) y(k) = Ex_array(i,2)
        if (k .eq. 1) x(j) = Ex_array(i,1)
     end do
  end do


end subroutine extract_2darrays_to_arrays

! 
end module histogram_mod