!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!
!
! mfu_sampler.f90
!
! You are welcome to use this code for your own research but please 
! properly cite and attibute:
!
! Author: 
! Cameron Bracken, cameron.bracken@gmail.com
!
! Citation: 
! C. Bracken, K. Holman, B. Rajagopalan, H. Moradkhani, A Bayesian 
! hierarchical approach to multivariate nonstationary hydrologic frequency 
! analysis, Water Resources Research, 2017, Under Review. 
!
! Multivariate from univariate (mfu) slice sampler. Allows for slice sampling 
! of multivariate density functions, one paramter at a time. 
! 
! The slice sampler is origianlly by Radford M. Neal, March 2008, translated 
! to Fortran by Cameron Bracken, October, 2015. 
!
!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!

module mfu_sampler
  use model
  implicit none

  integer, allocatable:: n_expands(:), n_shrinks(:)
  double precision, allocatable:: widths(:)
  integer:: n_proposals, n_iter_per_adapt
  double precision:: target_ratio_mfu

  contains 

! FUNCTIONS FOR PERFORMING UNIVARIATE SLICE SAMPLING.
!
! Radford M. Neal, 17 March 2008.
!
! Implements, with slight modifications and extensions, the algorithm described
! in Figures 3 and 5 of the following paper:
!
!   Neal, R. M (2003) "Slice sampling" (with discussion), Annals of Statistics,
!      vol. 31, no. 3, pp. 705-767.
!
! See the documentation for the function uni.slice below for how to use it.
! The function uni.slice.test was used to test the uni.slice function.

! public function: multivariate sampler, utilizing univariate slice samplers through a Gibbs wrapper
subroutine mfu_sample(x,iter,warmup)

    double precision:: x(:), x1, fx1, fx0
    integer:: k, n, nlpi, iter, warmup

    n = size(x)

    if(mod(iter, n_iter_per_adapt) .eq. 0 .and. iter <= warmup) then
      !write(*,'(a21,100f7.3)')'Tuning slice widths:', widths
      !write(*,*)'Tuning slice widths.'
      call tune_widths()
    end if 
    
    do k = 1,n

      !write(*,'(i5)', advance='no') k
      !nlpi = nlp

      if(k .eq. 1) fx0 = lp_fun(x)

      call mfu_slice_sampler(x, k, widths(k), lower(k), upper(k), x1, fx0)

      ! x should get modified in place, but just to be safe
      x(k) = x1
      fx0 = lp_fun(x)

      !write(*,*)nlp-nlpi

    end do

end subroutine

subroutine mfu_slice_sampler(x, k, w, lb, ub, x1, fx0) 

  ! UNIVARIATE SLICE SAMPLING WITH STEPPING OUT AND SHRINKAGE.
  !
  ! Performs a slice sampling update from an initial point to a new point that 
  ! leaves invariant the distribution with the specified log density function.
  !
  ! Arguments:
  !
  !   x0    Initial point
  !   g     Function returning the log of the probability density (plus constant)
  !   w     Size of the steps for creating interval (default 1)
  !   m     Limit on steps (default infinite)
  !   lower Lower bound on support of the distribution (default -Inf)
  !   upper Upper bound on support of the distribution (default +Inf)
  !   fx0   Value of g(x0), if known (default is not known)
  !
  ! The log density function may return -Inf for points outside the support 
  ! of the distribution.  If a lower and/or upper bound is specified for the
  ! support, the log density function will not be called outside such limits.
  !
  ! The value of this function is the new point sampled.
  !
  ! WARNING:  If you provide a value for g(x0), it must of course be correct!
  ! In addition to giving wrong answers, wrong values for fx0 may result in
  ! the uni.slice function going into an infinite loop.

  double precision:: x0, x(:), w, lb, ub, fx0, x1, fx1
  integer:: k, n
  double precision:: logy, u, L, R

  x0 = x(k)
  !x1 = x0

  ! Determine the slice level, in log terms.
  logy = fx0 - rexp(1d0)
  
  ! Find the initial interval to sample from.
  u = runif(0d0,w)
  L = x0 - u
  R = x0 + (w-u)  ! should guarantee that x0 is in [L,R], even with roundoff
  
  ! Expand the interval until its ends are outside the slice, or until
  ! the limit on steps is reached.
  do 
    if (L <= lb) exit
    x(k) = L
    if (lp_fun(x) <= logy) exit
    L = L - w
    n_expands(k) = n_expands(k) + 1
  end do  
    
  do
    if (R >= ub) exit
    x(k) = R
    if (lp_fun(x) <= logy) exit
    R = R + w
    n_expands(k) = n_expands(k) + 1
  end do

  !write(*,*)k, fx0, logy, L, R, x0

  
  ! else if (m>1)  ! limit on steps, bigger than one
  ! { 
  !   J = floor(runif(1,0,m))
  !   K = (m-1) - J
    
  !   while (J>0)
  !   { if (L<=lower) break
  !     !uni.slice.evals <= uni.slice.evals + 1
  !     if (g(L)<=logy) break
  !     L = L - w
  !     J = J - 1
  !   }
    
  !   while (K>0)
  !   { if (R>=upper) break
  !     !uni.slice.evals <= uni.slice.evals + 1
  !     if (g(R)<=logy) break
  !     R = R + w
  !     K = K - 1
  !   }
  ! }
  
  ! Shrink interval to lower and upper bounds.
  if (L < lb) L = lb
  if (R > ub) R = ub

  n = 0
  do
    n = n + 1
    x1 = runif(L,R)
    
    x(k) = x1 
    fx1 = lp_fun(x)
    
    !write(*,*) k, fx0, fx1, logy, fx1 >= logy, L, R, x0, x1

    if (fx1 >= logy) exit
    
    if (x1 > x0) then 
      R = x1 
    else 
      L = x1
    end if
    n_shrinks(k) = n_shrinks(k) + 1
  end do
  
end subroutine

subroutine tune_widths()

    double precision:: ratio, multiplier, denom
    integer:: i

    !write(*,*)'Tuning widths'
    !write(*,*)n_expands
    !write(*,*)n_shrinks

    do i=1,n_pars
        denom = dble(n_expands(i) + n_shrinks(i))

        if ( denom .gt. 0 ) then
        
            ratio = n_expands(i) / denom
            ! write(*,*)ratio

            if ( ratio .eq. 0d0 ) then
              ratio = 1d0 / denom
              !write(*,*)'ratio', ratio
            end if

            multiplier = (ratio / target_ratio_mfu)
            !write(*,*)'multiplier', multiplier, target_ratio_mfu


            ! Modify Interval Width
            widths(i) = widths(i) * multiplier

        end if
    end do

    !write(*,*)'Done Tuning widths'
    !write(*,*) w
    !  Reset Interval Width Counters
    n_expands   = 0
    n_shrinks   = 0
    n_proposals = 0

    ! Double Number of Iterations to Next Adaptation
    n_iter_per_adapt = n_iter_per_adapt * 2

end subroutine

subroutine mfu_init()
  
  integer::i,j
  allocate(n_expands(n_pars), n_shrinks(n_pars), widths(n_pars))

  n_expands = 0
  n_shrinks = 0 
  widths = 1d0

  n_proposals = 0
  n_iter_per_adapt = 1
  target_ratio_mfu = 0.5d0


end subroutine


function rexp(lambda)
  double precision:: lambda, r, rexp

  call random_number(r)
  rexp = - log(r) * lambda
end function

function runif(low, high)
  double precision:: runif, high, low, r

  call random_number(r)
  !write(*,*)high, low, r
  runif = r * (high - low) + low
end function

end module