!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!
!
! model.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. 
!
! This program defines the parameters and log-likelihood function that is called by 
! the sampler. It has three main components: 
!
!    (1) A module called `model` that defines the variables in your 
!        model likelihood function, 
!    (2) A subroutine called `read_data` that reads the data from the data 
!        files and allocates any variables specific to your model
!    (3) A function called lp_fun, that takes a vector of parameter 
!        values and computes the log of the posterior density (prior 
!        times likelihood)
!
!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!

module model
    use functions

    implicit none

    ! model parameters
    double precision, allocatable, dimension(:):: beta_loc_snow(:), &!beta_scale_snow(:), &
                                                  beta_loc_flow(:)!beta_scale_flow(:), &
                                                  !beta_loc_elev(:)   !beta_scale_elev(:)
    double precision:: scale_snow, scale_flow, scale_elev
    double precision:: shape_snow, shape_flow, shape_elev
    double precision:: cor_snow_flow, cor_snow_elev, cor_flow_elev, a, b, c
    !double precision:: mu_loc(3), sd_loc(3)

    ! data
    integer:: np_snow, np_flow, nt, n_pars, nlp = 0

    ! precip, flow, res elevation, covariates, initial values
    double precision, allocatable:: y(:), z(:), r(:), covars_snow(:,:), covars_flow(:,:), x0(:)

    ! parameter bounds
    double precision, allocatable:: lower(:), upper(:), w(:)

    ! constants
    double precision, parameter:: zero=0d0, one=1d0, three=3d0, five=5d0, ten=10d0
    double precision, parameter:: onehundred=100d0, threehundred=300d0, nugget=0.01d0
    double precision, parameter:: p1 = 0.1d0, p3 = 0.3d0, p5 = 0.5d0, p7 = 0.7d0
    double precision, parameter:: infinity = 1.7976931348623157d300


    contains

subroutine read_data()
    
    integer:: i,j

    open(12,file='model_data/sizes')
    open(13,file='model_data/init')
    open(14,file='model_data/y')
    open(15,file='model_data/z')
    open(16,file='model_data/r')
    open(17,file='model_data/covars_snow')
    open(18,file='model_data/covars_flow')
    open(19,file='model_data/lower')
    open(20,file='model_data/upper')
    

    read(12,*) np_snow, np_flow, nt, n_pars
    !write(*,*) np_snow, np_flow, nt, n_pars

    ! parameters
    allocate(beta_loc_snow(np_snow), &!beta_scale_snow(np), &
             beta_loc_flow(np_flow))!, &!beta_scale_flow(np), &
             !beta_loc_elev(np))!, beta_scale_elev(np))

    ! data
    allocate(x0(n_pars), y(nt), z(nt), r(nt), covars_snow(nt,np_snow), covars_flow(nt,np_flow))

    ! need to allocate 
    allocate(lower(n_pars), upper(n_pars), w(n_pars))
    w = 1d0

    read(13,*) x0
    read(14,*) y
    read(15,*) z
    read(16,*) r
    read(17,*) covars_snow
    read(18,*) covars_flow
    read(19,*) lower
    read(20,*) upper

    ! write(*,'(1000f5.2)') x0
    ! write(*,*)

    ! do i=1,nt
    !     write(*,'(100f8.2)') (y(i,j),j=1,ns)
    ! end do 
    ! write(*,*)
    
    ! write(*,*) z
    ! write(*,*)

    ! write(*,'(1000f10.2)') r
    ! write(*,*)

    ! do i=1,nt
    !     write(*,'(100f6.2)') (covars(i,j),j=1,np)
    ! end do 
    ! write(*,*)

    ! write(*,'(1000f9.2)') lower
    ! write(*,*)

    ! write(*,'(1000f9.2)') upper

    close(12)
    close(13)
    close(14)
    close(15)
    close(16)
    close(17)
    close(18)
    close(19)
    close(20)

end subroutine

function lp_fun(x)

  double precision:: x(:), lp_fun, lp
  integer:: t, s, k, n, p, i, j, ip

  !write(*,*)'lp fun'
  n = size(x)
  !write(*,*)'npars=',n

  !write(*,*)x

  p = 1
  ! beta_loc_snow
  do ip = 1,np_snow
    beta_loc_snow(ip) = x(p)
    !write(*,'(a20,2i5,f20.5)')'beta_loc_snow', ip, p, x(p)
    p = p + 1
  end do

  ! ! beta_scale_snow
  ! do ip = 1,np
  !   beta_scale_snow(ip) = x(p)
  !   !write(*,'(a20,2i5,f20.5)')'beta_scale_snow', ip, p, x(p)
  !   p = p + 1
  ! end do

  ! beta_loc_flow
  do ip = 1,np_flow
    beta_loc_flow(ip) = x(p)
    !write(*,'(a20,2i5,f20.5)')'beta_loc_flow', ip, p, x(p)
    p = p + 1
  end do

  ! ! beta_scale_flow
  ! do ip = 1,np
  !   beta_scale_flow(ip) = x(p)
  !   !write(*,'(a20,2i5,f20.5)')'beta_scale_flow', ip, p, x(p)
  !   p = p + 1
  ! end do

  ! beta_loc_elev
  ! do ip = 1,np
  !   beta_loc_elev(ip) = x(p)
  !   !write(*,'(a20,2i5,f20.5)')'beta_loc_elev', ip, p, x(p)
  !   p = p + 1
  ! end do

  ! ! beta_scale_elev
  ! do ip = 1,np
  !   beta_scale_elev(ip) = x(p)
  !   !write(*,'(a20,2i5,f20.5)')'beta_scale_elev', ip, p, x(p)
  !   p = p + 1
  ! end do

  !write(*,*) 'here, p = ',p

  scale_snow = x(p) 
  !write(*,'(a20,i10,f20.5)')'scale_snow', p, scale_snow
  p = p + 1

  scale_flow = x(p) 
  !write(*,'(a20,i10,f20.5)')'scale_flow', p, scale_flow
  p = p + 1

  scale_elev = x(p) 
  !write(*,'(a20,i10,f20.5)')'scale_elev', p, scale_elev
  p = p + 1

  shape_snow = x(p) 
  !write(*,'(a20,i10,f20.5)')'shape_snow', p, x(p)
  p = p + 1

  shape_flow = x(p) 
  !write(*,'(a20,i10,f20.5)')'shape_flow', p+1, x(p+1)
  p = p + 1

  shape_elev = x(p) 
  !write(*,'(a20,i10,f20.5)')'shape_elev', p+2, x(p+2)
  p = p + 1

  cor_snow_flow = x(p) 
  !write(*,'(a20,i10,f20.5)')'cor_snow_flow', p+3, x(p+3)
  p = p + 1

  cor_snow_elev = x(p) 
  !write(*,'(a20,i10,f20.5)')'cor_snow_elev', p+4, x(p+4)
  p = p + 1

  cor_flow_elev = x(p) 
  p = p + 1

  a = x(p) 
  p = p + 1

  b = x(p) 
  p = p + 1

  c = x(p) 
  !write(*,'(a20,i10,f20.5)')'cor_flow_elev', p+5, x(p+5)

  !write(*,*) 'here, p = ',p, ' n = ', n

  if(p .ne. n) then
    write(*,*) 'p not equal to number of pars!!!!'
    write(*,*)'here, p = ',p, ' n = ', n
  end if


  !write(*,'(a)', advance='no')'starting lp calc'
  call lp_compute(lp)
  !write(44,*) x
  !write(*,*) lp
  lp_fun = lp

  nlp = nlp + 1

  return 
end function


subroutine lp_compute(lp)

  ! global parameter variables come from the module

  ! local variables
  double precision:: loc_snow, loc_flow, loc_elev
  double precision:: lpdf(3), cdf(3), cov_loc(3,3), chol_loc(3,3), loc(3)

  double precision:: lp, ll
  integer:: i, j, k, t, s, ip, info

  ! double precision:: cor_snow_flow, cor_snow_elev, cor_flow_elev

  ! write(*,*)'lp compute'

  lp = 0
  ll = 0

  ! priors
  lp = lp + &
    norm(shape_flow, zero, p3) + &
    norm(shape_snow, zero, p3) + &
    norm(shape_elev, zero, p3) + &
    norm(a,  9.33d0, 0.002d0) + & 
    norm(b, 0.003d0, 0.002d0) + & 
    norm(c,    20d0, 10d0) 
    !std_norm(sd_loc) !+ &
    !std_norm(mu_loc) + & 
    ! norm(cor_snow_flow, 0.8463796d0, 0.05d0) + & 
    ! norm(cor_snow_elev, 0.3956583d0, 0.05d0) + & 
    ! norm(cor_flow_elev, 0.5478896d0, 0.05d0) 
    ! norm(beta_loc_snow, zero, five) + & 
    ! norm(beta_scale_snow, zero, five) + & 
    ! norm(beta_loc_flow, zero, five) + & 
    ! norm(beta_scale_flow, zero, five) + & 
    ! norm(beta_loc_elev, zero, five) + & 
    ! norm(beta_scale_elev, zero, five)
  
  !write(*,*)'After priors',lp
  if(lp <= -infinity) return

  ! write(*,*)'After priors'


  ! joint distribution for the data
  ! for now, cov_loc is the dependence matrix
  ! standard deviations on the diagonal, squared to make variances
  ! correlations on the off diagonal, converted to covariances
  cov_loc(1,1) = 1
  cov_loc(2,2) = 1
  cov_loc(3,3) = 1

  !cor_snow_flow = 0.8463796d0
  !cor_snow_elev = 0.3956583d0
  !cor_flow_elev = 0.5478896d0

  cov_loc(1,2) = cor_snow_flow
  cov_loc(1,3) = cor_snow_elev
  cov_loc(2,3) = cor_flow_elev

  cov_loc(2,1) = cor_snow_flow
  cov_loc(3,1) = cor_snow_elev
  cov_loc(3,2) = cor_flow_elev


  ! write(*,*)'Before cholesky'
  ! do i=1,3
  !   write(*,'(100f7.4)') (cov_loc(i,j),j=1,3)
  ! end do 
  ! write(*,*)

  !write(*,*)'Copula'
  ! L is now the cholesky factor of the dependogram matrix
  ! covariance matrix is now cholesky factorization
  ! L = cholesky(L)
  i = 3
  chol_loc = cov_loc 
  call DPOTRF( 'L', i, chol_loc, i, info )

  ! compute gev liklihoods
  do t=1,nt

    ! compute these on the log scale to keep them positive
    loc_snow = sum(covars_snow(t,:) * beta_loc_snow)
    !scale_snow = exp(sum(covars(t,:) * beta_scale_snow))
    !write(*,*)loc_snow

    loc_flow = sum(covars_flow(t,:) * beta_loc_flow)
    !scale_flow = exp(sum(covars(t,:) * beta_scale_flow))

    loc_elev = a/(1+b*c**(-loc_flow))!sum(covars(t,:) * beta_loc_elev)
    !scale_elev = exp(sum(covars(t,:) * beta_scale_elev))

    ! write(*,*)loc_snow, scale_snow, loc_flow, scale_flow, loc_elev, scale_elev
    ! write(*,*)shape_snow, shape_flow, shape_elev

    ! check constraints
    ! do s = 1,ns
    !   ! gev precip constraint 1 + xi*(x-mu)/sigma > 0
    !   if(y(t,s) .gt. -99)then
    !     if(one + shape_snow*(y(t,s)-loc_snow)/scale_snow .lt. zero)then
    !       !write(*,*) 'GEV constraint: snow'
    !       lp = -infinity 
    !     end if
    !   end if
    ! end do

    ! ! loc_flow and scale_flow > 0 
    ! if(loc_snow .lt. zero .or. scale_snow .lt. zero)then
    !   !write(*,*) 'snow loc or scale lt 0: snow', t, loc_snow, scale_snow
    !   !write(*,*) beta_loc_snow
    !   lp = -infinity 
    ! end if

    ! ! write(*,*) 'After Constraints snow'

    ! ! loc_flow and scale_flow > 0 
    ! if(loc_flow .lt. zero .or. scale_flow .lt. zero)then
    !   !write(*,*) 'flow loc or scale lt 0: flow'
    !   lp = -infinity 
    ! end if

    ! ! write(*,*) 'After Constraints flow 1'

    ! !  gev contraint 1 + xi*(x-mu)/sigma > 0
    ! if(one + shape_flow*(z(t)-loc_flow)/scale_flow .lt. zero)then
    !   !write(*,*) 'GEV constraint: flow'
    !   lp = -infinity 
    ! end if

    ! ! write(*,*) 'After Constraints flow 2'

    ! ! loc_elev and scale_elev > 0 
    ! if(loc_elev .lt. zero .or. scale_elev .lt. zero)then
    !   !write(*,*) 'elev loc or scale lt 0: flow'
    !   lp = -infinity 
    ! end if

    ! !  gev contraint 1 + xi*(x-mu)/sigma > 0
    ! if(one + shape_elev*(z(t)-loc_elev)/scale_elev .lt. zero)then
    !   !write(*,*) 'GEV constraint: elev'
    !   lp = -infinity 
    ! end if

    ! if(lp >= infinity) lp = -infinity
    ! if(lp <= -infinity) then 
    !   !write(*,*) 'Breaking on lp before likelihoods'
    !   return
    ! end if

    ! write(*,*) 'After Constraints'
    ! gev snow likelihood
    lpdf(1) = gev(y(t), loc_snow, scale_snow, shape_snow)
    !write(*,*)'After snow',lp
    !write(*,*)y(t,s), loc_snow, scale_snow, shape_snow

    ! gev flow likelihood
    lpdf(2) = gev(z(t), loc_flow, scale_flow, shape_flow)
    !write(*,*)'After flow',gev(z(t), loc_flow, scale_flow, shape_flow)
    !write(*,*)z(t), loc_flow, scale_flow, shape_flow

    ! gev elevation likelihood
    lpdf(3) = gev(r(t), loc_elev, scale_elev, shape_elev)
    !write(*,*)'After elev',gev(r(t), loc_elev, scale_elev, shape_elev)
    !write(*,*)r(t), loc_elev, scale_elev, shape_elev


    ! ! need to make a copy because chol_mvnorm modifies the cholesky matrix in place
    !chol_loc = cov_loc

    ! pass in cholesky of covariance matrix, log likelihood output
    !i = 3
    !call chol_mvnorm(loc, mu_loc, chol_loc, i, ll, info)
    !lp = lp + ll

    cdf(1) = gev_cdf(y(t), loc_snow, scale_snow, shape_snow)
    cdf(2) = gev_cdf(z(t), loc_flow, scale_flow, shape_flow)
    cdf(3) = gev_cdf(r(t), loc_elev, scale_elev, shape_elev)

    !write(*,*) mu_loc,sd_loc
    !write(*,*) loc
    !write(*,*) lpdf, cdf
    ll = gauss_eliptical_copula(lpdf, cdf, chol_loc)
    !write(*,*)ll
    lp = lp + ll + sum(lpdf)

    !write(*,*)'After mvn',lp

    if(lp >= infinity) lp = -infinity
    if(lp <= -infinity) then 
      !write(*,*) 'Breaking on lp after likelihoods'
      return
    end if
    !stop
  end do
  !stop
  
  !write(*,*)'End' 
  
  return 
end subroutine 

end module