program test_prob
  use :: iso_c_binding
  use :: sort
  use :: prob, only : pdists, prob_dists
  use :: dists
  use :: statistics_mod
  implicit none

  integer, parameter :: dp = c_double
  integer, parameter :: i8 = c_int64_t

  ! type(prob_dists) :: pdists
  integer(kind=c_int64_t), dimension(:) :: rarray(5)
  integer(kind=c_int64_t), dimension(:), allocatable :: rchoose
  logical, dimension(:), allocatable :: rbernoulli

  integer(kind=c_int), dimension(:) :: items(5), sort_indx(5)
  real(kind=c_double), dimension(:) :: sort_items(5)
  real(kind=c_double), dimension(:) :: weights(5)

  integer(kind=c_int), dimension(:) :: wchoose(3)

  integer(kind=c_int), dimension(:) :: items_permute(5)

  integer :: j

  type(lognormal) :: ln
  type(normal) :: n
  type(TOST) :: t
  type(tdist) :: td

  real(kind=c_double) :: mu, sigma
  real(kind=c_double), dimension(:), allocatable :: lns, ns, ts

  type(lognormal) :: logN_c, logN_a

  real(kind=dp), dimension(:), allocatable :: data1_x, data1_p
  real(kind=dp), dimension(:), allocatable :: data2_x, data2_p
  type(empirical_dist) :: contacts1
  type(empirical_dist) :: contacts2
  type(empirical_dist) :: contacts
  real(kind=dp) :: weight
  real(kind=dp) :: contacts_rvs

  pdists = prob_dists()

  print *, 'poisson = ', pdists%poisson_rvs(1.0_dp)

  rarray = pdists%sample(5_i8)

  rchoose = pdists%choose(rarray,1)

  print *, 'rarray = ', rarray
  print *, 'rchoose =', rchoose

  print *, 'bernoulli_rvs = ', pdists%bernoulli_rvs(real(0.25,kind=c_double))

  rbernoulli = pdists%bernoulli_rvs(real(0.25,kind=c_double),int(5,kind=c_int64_t))

  print *, 'rbernoulli = ', rbernoulli

  do j=1,10
    print *, pdists%loglogistic_rvs(shape=1.0_dp,scale=1.0_dp),','
  end do

  items = (/1, 2, 3, 4, 5/)
  weights = (/0.1_dp, 0.05_dp, 0.4_dp, 0.45_dp, 0.0_dp/)

  do j=1,10
    wchoose = pdists%weighted_choose(3, items, weights)
    print *, j, '   ', wchoose
  end do

  print *, 'items'
  print *, items
  print *, 'permute'
  do j=1,10
    items_permute = pdists%shuffle(items)
    print *, items_permute
  end do

  sort_items = (/3_i8, 4_i8, 5_i8, 1_i8, 2_i8/)

  print *, 'pre-sort: ', sort_items
  call Qsort(sort_items, sort_indx)
  print *, 'post-sort: ', sort_items
  print *, 'index: ', sort_indx

  sort_items = (/3.14_dp, 4.41_dp, 5.712_dp, 1.452_dp, 2.7863_dp/)

  print *, 'pre-sort: ', sort_items
  call Qsort(sort_items, sort_indx, order=-1)
  print *, 'post-sort: ', sort_items
  print *, 'index: ', sort_indx


  mu = 5.7_dp
  sigma = 3.5_dp

  print *, 'lognormal, normal'
  ln = lognormal(1.63_dp, 0.5_dp)
  n = normal(mu, sigma)

  allocate(lns(1000), ns(1000))
  do j=1,1000
    lns(j) = ln%sample()
    ns(j) = n%sample()
  end do

  print *, 'stats'
  print *, 'normal'
  print *, 'sample  :   expected'
  print *, mean_1d(ns), n%mean
  print *, variance_1d(ns), n%variance
  print *, stddev_1d(ns), sqrt(n%variance)

  print *, 'lognormal'
  print *, 'sample  :   expected'
  print *, mean_1d(lns), ln%mean
  print *, variance_1d(lns), ln%variance
  print *, stddev_1d(lns), sqrt(ln%variance)
  
  t = TOST(5.42_dp)
  allocate(ts(1000000))
  do j=1,1000000
    ts(j) = t%sample()
  end do
  print *, 'tost'
  print *, 'sample  :   expected'
  print *, mean_1d(ts), 0.1731545471_dp
  print *, variance_1d(ts), 6.267666129_dp
  print *, stddev_1d(ts), sqrt(6.267666129_dp)
  deallocate(ts)

  print *, 't-distribution'
  td = tdist(3.0_dp)
  print *, 'nu = 3 : variance = ', td%variance
  td = tdist(1.5_dp)
  print *, 'nu = 1.5 : variance = ', td%variance
  td = tdist(0.5_dp)
  print *, 'nu = 0.5 : variance = ', td%variance

  td = tdist(3.3454_dp, location=-0.0747_dp, scale=1.8567_dp)

  print *, td%sample()
  allocate(ts(1000000))
  do j=1,1000000
    ts(j) = td%sample()
  end do
  print *, 'tost'
  print *, 'sample  :   expected'
  print *, mean_1d(ts), td%mean
  print *, variance_1d(ts), td%variance
  print *, stddev_1d(ts), sqrt(td%variance)
  deallocate(ts)

  print *, td%quantile(0.05_dp), td%quantile(0.25_dp), td%quantile(0.5_dp), td%quantile(0.75_dp), td%quantile(0.95_dp)


  print *, 'unwell time distributions'

  logN_c = lognormal(1.609_dp, sqrt(0.923_dp))
  logN_a = lognormal(2.398_dp, sqrt(0.840_dp))

  print *, 'Child :', logN_c%median, logN_c%mean, logN_c%quantile(1.0_dp - 0.031_dp), (1.0_dp - logN_c%cdf(28.0_dp))*100_dp
  
  print *, 'Adult : ', logN_a%median, logN_a%mean, logN_a%quantile(1.0_dp - 0.133_dp), (1.0_dp - logN_a%cdf(28.0_dp))*100_dp

  print *, 'Empirical distribution'

  allocate(data1_x(35), data1_p(35)) 
  data1_x = (/0.0_dp, 1.0_dp, 4.0_dp, 5.0_dp, 6.0_dp, 7.0_dp, 8.0_dp, 9.0_dp, 10.0_dp, 12.0_dp, 13.0_dp, 14.0_dp, 15.0_dp, 16.0_dp, 20.0_dp, 22.0_dp, 24.0_dp, 30.0_dp, 33.0_dp, 34.0_dp, 35.0_dp, 40.0_dp, 41.0_dp, 43.0_dp, 45.0_dp, 48.0_dp, 50.0_dp, 60.0_dp, 62.0_dp, 65.0_dp, 68.0_dp, 80.0_dp, 100.0_dp, 150.0_dp, 151.0_dp/)
  data1_p = (/0.0_dp, 0.003065639_dp, 0.009238066_dp, 0.063602629_dp, 0.118351948_dp, 0.173101558_dp, 0.227851167_dp, 0.282601402_dp, 0.33740609_dp, 0.44752922_dp, 0.502747731_dp, 0.557988658_dp, 0.613207396_dp, 0.668428206_dp, 0.889313243_dp, 0.891248229_dp, 0.893184187_dp, 0.903604983_dp, 0.908542623_dp, 0.91060152_dp, 0.91266046_dp, 0.922964469_dp, 0.925025411_dp, 0.928929996_dp, 0.932723992_dp, 0.938083412_dp, 0.941656333_dp, 0.954571595_dp, 0.957079356_dp, 0.959963346_dp, 0.966724723_dp, 0.993769873_dp, 0.995515816_dp, 0.999866681_dp, 1.0_dp/)

  allocate(data2_x(23), data2_p(23))
  data2_x = (/0.0_dp, 0.001_dp, 1.0_dp, 2.0_dp, 3.0_dp, 4.0_dp, 5.0_dp, 6.0_dp, 7.0_dp, 8.0_dp, 10.0_dp, 12.0_dp, 14.0_dp, 15.0_dp, 16.0_dp, 18.0_dp, 20.0_dp, 23.0_dp, 28.0_dp, 35.0_dp, 45.0_dp, 75.0_dp, 76.0_dp/)
  data2_p = (/0.0_dp, 0.002558737_dp, 0.007851272_dp, 0.013530773_dp, 0.221824454_dp, 0.430697977_dp, 0.5366136_dp, 0.642337377_dp, 0.74806127_dp, 0.854549502_dp, 0.861047298_dp, 0.880665976_dp, 0.900283987_dp, 0.907873309_dp, 0.915903065_dp, 0.931966576_dp, 0.98986558_dp, 0.990809823_dp, 0.992053894_dp, 0.993795611_dp, 0.995318253_dp, 0.999885807_dp, 1.0_dp/)

  contacts1 = empirical_dist(data1_x, data1_p)
  contacts2 = empirical_dist(data2_x, data2_p)

  contacts = empirical_dist(contacts1, contacts2, 0.5_dp)

  do j=1,10
    contacts_rvs = contacts%sample()
    print *, contacts_rvs
  end do


end program test_prob