module statistics_mod

  use :: iso_c_binding
  use :: fgsl

  implicit none

  private
  public :: mean_1d
  public :: variance_1d
  public :: stddev_1d
  public :: quantiles_1d
  public :: mean_over_rows
  public :: variance_over_rows
  public :: stddev_over_rows
  public :: quantiles_over_rows

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

  interface mean_over_rows
    module procedure mean_over_rows_i4
    module procedure mean_over_rows_i8
  end interface mean_over_rows

  interface variance_over_rows
    module procedure variance_over_rows_i4
    module procedure variance_over_rows_i8
  end interface variance_over_rows

  interface stddev_over_rows
    module procedure stddev_over_rows_i4
    module procedure stddev_over_rows_i8
  end interface stddev_over_rows

  interface quantiles_over_rows
   module procedure quantiles_over_rows_i4
   module procedure quantiles_over_rows_i8
  end interface quantiles_over_rows

contains

  function mean_1d(array) result(mean)
    implicit none
    real(kind=dp), dimension(:), intent(in) :: array
    real(kind=dp) :: mean
    mean = fgsl_stats_mean(array, 1_i8, int(size(array), kind=c_int64_t))
  end function mean_1d

  function variance_1d(array) result(var)
    implicit none
    real(kind=dp), dimension(:), intent(in) :: array
    real(kind=dp) :: var
    var = fgsl_stats_variance(array, 1_i8, int(size(array), kind=c_int64_t))
  end function variance_1d

  function stddev_1d(array) result(sd)
    implicit none
    real(kind=dp), dimension(:), intent(in) :: array
    real(kind=dp) :: sd
    sd = fgsl_stats_sd(array, 1_i8, int(size(array), kind=c_int64_t))
  end function stddev_1d

  function quantiles_1d(array, f) result(quantiles)
    implicit none
    real(kind=dp), dimension(:), intent(in) :: array
    real(kind=dp), dimension(:), intent(in) :: f
    real(kind=dp), dimension(:) :: quantiles(size(f))

    integer :: k

    integer(kind=c_int64_t) :: N
    real(kind=dp), dimension(:) :: a(size(array))

    N = int(size(array), kind=c_int64_t)

    a(:) = array(:)
    
    call fgsl_sort(a, 1_i8, N)
    do k=1,size(f)
      quantiles(k) = fgsl_stats_quantile_from_sorted_data( a, &
                                                          1_i8, &
                                                          N, &
                                                          f(k))
    end do

    return

  end function quantiles_1d


  subroutine mean_over_rows_i4(array, mean)
    implicit none
    integer(kind=c_int), dimension(:,:), intent(in) :: array
    real(kind=c_double), dimension(:), allocatable, intent(out) :: mean

    real(kind=c_double), dimension(:), allocatable :: a

    integer :: rows
    integer :: columns
    integer :: j

    if (allocated(mean)) deallocate(mean)

    rows = size(array(:,1))
    columns = size(array(1,:))

    allocate(mean(columns))
    allocate(a(rows))

    do j=1,columns
      !mean(j) = real(sum(array(:,j)), kind=c_double)/real(rows, kind=c_double)
      a = real(array(:,j), kind=c_double)
      mean(j) = fgsl_stats_mean(a, 1_i8, int(rows, kind=c_int64_t))
    end do

    deallocate(a)
    return

  end subroutine mean_over_rows_i4

  subroutine mean_over_rows_i8(array, mean)
    implicit none
    integer(kind=c_int64_t), dimension(:,:), intent(in) :: array
    real(kind=c_double), dimension(:), allocatable, intent(out) :: mean

    real(kind=c_double), dimension(:), allocatable :: a

    integer :: rows
    integer :: columns
    integer :: j

    if (allocated(mean)) deallocate(mean)

    rows = size(array(:,1))
    columns = size(array(1,:))

    allocate(mean(columns))
    allocate(a(rows))

    do j=1,columns
      ! mean(j) = real(sum(array(:,j)), kind=c_double)/real(rows, kind=c_double)
      a = real(array(:,j), kind=c_double)
      mean(j) = fgsl_stats_mean(a, 1_i8, int(rows, kind=c_int64_t))
    end do

    deallocate(a)
    return

  end subroutine mean_over_rows_i8

  subroutine variance_over_rows_i4(array, var)
    implicit none
    integer(kind=c_int), dimension(:,:), intent(in) :: array
    real(kind=c_double), dimension(:), allocatable, intent(out) :: var

    real(kind=c_double), dimension(:), allocatable :: a

    integer :: rows
    integer :: columns
    integer :: j

    if (allocated(var)) deallocate(var)

    rows = size(array(:,1))
    columns = size(array(1,:))

    allocate(var(columns))
    allocate(a(rows))

    do j=1,columns
      !mean(j) = real(sum(array(:,j)), kind=c_double)/real(rows, kind=c_double)
      a = real(array(:,j), kind=c_double)
      var(j) = fgsl_stats_variance(a, 1_i8, int(rows, kind=c_int64_t))
    end do

    deallocate(a)
    return

  end subroutine variance_over_rows_i4

  subroutine variance_over_rows_i8(array, var)
    implicit none
    integer(kind=c_int64_t), dimension(:,:), intent(in) :: array
    real(kind=c_double), dimension(:), allocatable, intent(out) :: var

    real(kind=c_double), dimension(:), allocatable :: a

    integer :: rows
    integer :: columns
    integer :: j

    if (allocated(var)) deallocate(var)

    rows = size(array(:,1))
    columns = size(array(1,:))

    allocate(var(columns))
    allocate(a(rows))

    do j=1,columns
      !mean(j) = real(sum(array(:,j)), kind=c_double)/real(rows, kind=c_double)
      a = real(array(:,j), kind=c_double)
      var(j) = fgsl_stats_variance(a, 1_i8, int(rows, kind=c_int64_t))
    end do

    deallocate(a)
    return

  end subroutine variance_over_rows_i8

  subroutine stddev_over_rows_i4(array, sd)
    implicit none
    integer(kind=c_int), dimension(:,:), intent(in) :: array
    real(kind=c_double), dimension(:), allocatable, intent(out) :: sd

    real(kind=c_double), dimension(:), allocatable :: a

    integer :: rows
    integer :: columns
    integer :: j

    if (allocated(sd)) deallocate(sd)

    rows = size(array(:,1))
    columns = size(array(1,:))

    allocate(sd(columns))
    allocate(a(rows))

    do j=1,columns
      !mean(j) = real(sum(array(:,j)), kind=c_double)/real(rows, kind=c_double)
      a = real(array(:,j), kind=c_double)
      sd(j) = fgsl_stats_sd(a, 1_i8, int(rows, kind=c_int64_t))
    end do

    deallocate(a)
    return

  end subroutine stddev_over_rows_i4

  subroutine stddev_over_rows_i8(array, sd)
    implicit none
    integer(kind=c_int64_t), dimension(:,:), intent(in) :: array
    real(kind=c_double), dimension(:), allocatable, intent(out) :: sd

    real(kind=c_double), dimension(:), allocatable :: a

    integer :: rows
    integer :: columns
    integer :: j

    if (allocated(sd)) deallocate(sd)

    rows = size(array(:,1))
    columns = size(array(1,:))

    allocate(sd(columns))
    allocate(a(rows))

    do j=1,columns
      !mean(j) = real(sum(array(:,j)), kind=c_double)/real(rows, kind=c_double)
      a = real(array(:,j), kind=c_double)
      sd(j) = fgsl_stats_sd(a, 1_i8, int(rows, kind=c_int64_t))
    end do

    deallocate(a)
    return

  end subroutine stddev_over_rows_i8

  subroutine quantiles_over_rows_i4(array, f, quantiles)
    implicit none
    integer(kind=c_int), dimension(:,:), intent(in) :: array
    real(kind=c_double), dimension(:), intent(in) :: f
    integer(kind=c_int), dimension(:,:), allocatable, intent(out) :: quantiles

    real(kind=c_double), dimension(:), allocatable :: a

    integer :: rows
    integer :: columns
    integer :: nf
    integer :: j, k

    if (allocated(quantiles)) deallocate(quantiles)

    nf = size(f)
    rows = size(array(:,1))
    columns = size(array(1,:))

    allocate(quantiles(columns,nf))
    allocate(a(rows))

    do j=1,columns
      !mean(j) = real(sum(array(:,j)), kind=c_double)/real(rows, kind=c_double)
      a = real(array(:,j), kind=c_double)
      call fgsl_sort(a, 1_i8, int(rows, kind=c_int64_t))
      do k=1,nf
        quantiles(j,k) = nint(fgsl_stats_quantile_from_sorted_data( &
                            a, 1_i8, int(rows, kind=c_int64_t), f(k)), &
                              kind=c_int)
      end do
    end do

    deallocate(a)
    return

  end subroutine quantiles_over_rows_i4

  subroutine quantiles_over_rows_i8(array, f, quantiles)
    implicit none
    integer(kind=c_int64_t), dimension(:,:), intent(in) :: array
    real(kind=c_double), dimension(:), intent(in) :: f
    integer(kind=c_int64_t), dimension(:,:), allocatable, intent(out) :: quantiles

    real(kind=c_double), dimension(:), allocatable :: a

    integer :: rows
    integer :: columns
    integer :: nf
    integer :: j, k

    if (allocated(quantiles)) deallocate(quantiles)

    nf = size(f)
    rows = size(array(:,1))
    columns = size(array(1,:))

    allocate(quantiles(columns,nf))
    allocate(a(rows))

    do j=1,columns
      !mean(j) = real(sum(array(:,j)), kind=c_double)/real(rows, kind=c_double)
      a = real(array(:,j), kind=c_double)
      call fgsl_sort(a, 1_i8, int(rows, kind=c_int64_t))
      do k=1,nf
        quantiles(j,k) = nint(fgsl_stats_quantile_from_sorted_data( &
                            a, 1_i8, int(rows, kind=c_int64_t), f(k)), &
                              kind=c_int64_t)
      end do
    end do

    deallocate(a)
    return

  end subroutine quantiles_over_rows_i8

end module statistics_mod