#:include "common.fypp"
#:set BETA_REAL_KINDS_TYPES = [item for item in REAL_KINDS_TYPES if item[0] in ('sp', 'dp', 'xdp')]
#:set BETA_CMPLX_KINDS_TYPES = [item for item in CMPLX_KINDS_TYPES if item[0] in ('sp', 'dp', 'xdp')]
#:set RC_KINDS_TYPES = BETA_REAL_KINDS_TYPES + BETA_CMPLX_KINDS_TYPES
module stdlib_stats_distribution_beta
    use ieee_arithmetic, only: ieee_value, ieee_quiet_nan, ieee_is_nan
    use stdlib_kinds, only : sp, dp, xdp
    use stdlib_error, only : error_stop
    use stdlib_optval, only : optval
    use stdlib_stats_distribution_uniform, only : uni=>rvs_uniform
    use stdlib_stats_distribution_gamma, only : rgamma=>rvs_gamma
    use stdlib_specialfunctions_gamma, only : beta, incomplete_beta, log_beta

    implicit none
    private

    public :: rvs_beta
    public :: pdf_beta
    public :: cdf_beta


    interface rvs_beta
    !! Version experimental
    !!
    !! Beta Distribution Random Variates
    !! ([Specification](../page/specs/stdlib_stats_distribution_beta.html#
    !! rvs_beta-beta-distribution-random-variates))
    !!
        #:for k1, t1 in RC_KINDS_TYPES
        module procedure beta_dist_rvs_${t1[0]}$${k1}$       ! 2 arguments
        #:endfor

        #:for k1, t1 in RC_KINDS_TYPES
        module procedure beta_dist_rvs_array_${t1[0]}$${k1}$ ! 3 arguments
        #:endfor
    end interface rvs_beta


    interface pdf_beta
    !! Version experimental
    !!
    !! Beta Distribution Probability Density Function
    !! ([Specification](../page/specs/stdlib_stats_distribution_beta.html#
    !! pdf_beta-beta-distribution-probability-density-function))
    !!
        #:for k1, t1 in RC_KINDS_TYPES
        module procedure beta_dist_pdf_${t1[0]}$${k1}$
        #:endfor

        #:for k1, t1 in RC_KINDS_TYPES
        module procedure beta_dist_pdf_impure_${t1[0]}$${k1}$
        #:endfor
    end interface pdf_beta


    interface cdf_beta
    !! Version experimental
    !!
    !! Beta Distribution Cumulative Distribution Function
    !! ([Specification](../page/specs/stdlib_stats_distribution_beta.html#
    !! cdf_beta-beta-distribution-cumulative-distribution-function))
    !!
        #:for k1, t1 in RC_KINDS_TYPES
        module procedure beta_dist_cdf_${t1[0]}$${k1}$
        #:endfor

        #:for k1, t1 in RC_KINDS_TYPES
        module procedure beta_dist_cdf_impure_${t1[0]}$${k1}$
        #:endfor
    end interface cdf_beta




contains

    #:for k1, t1 in BETA_REAL_KINDS_TYPES
    impure elemental function beta_dist_rvs_${t1[0]}$${k1}$(a, b, loc)      &
        result(res)
    ! Beta random variate using gamma variates
    ! If a < 1 or b < 1, uses uniform method, otherwise uses gamma method
    !
        ${t1}$, intent(in) :: a, b
        ${t1}$, intent(in), optional :: loc
        ${t1}$ :: res, x, y, xx(2)
        ${t1}$ :: loc_
        ${t1}$, parameter :: zero = 0.0_${k1}$, one = 1.0_${k1}$

        loc_ = optval(loc, 0.0_${k1}$)

        if(a <= zero .or. b <= zero) then
            res = ieee_value(1.0_${k1}$, ieee_quiet_nan)
            return
        end if

        if(a < one .or. b < one) then
            ! Use uniform method for small parameters
            do
                xx = uni(zero, one, 2)
                x = xx(1) ** (one / a)
                y = xx(2) ** (one / b)
                y = x + y
                if(y <= one .and. y /= zero) exit
            end do
        else
            ! Use gamma method for larger parameters
            do
                x = rgamma(a, one)
                y = rgamma(b, one)
                y = x + y
                if(y /= zero) exit
            end do
        end if
        
        res = x / y + loc_
    end function beta_dist_rvs_${t1[0]}$${k1}$

    #:endfor


    #:for k1, t1 in BETA_CMPLX_KINDS_TYPES
    impure elemental function beta_dist_rvs_${t1[0]}$${k1}$(a, b, loc)      &
        result(res)
    ! Complex parameter beta distributed. The real part and imaginary part are
    ! independent of each other.
    !
        ${t1}$, intent(in) :: a, b
        ${t1}$, intent(in), optional :: loc
        ${t1}$ :: res
        ${t1}$ :: loc_

        loc_ = optval(loc, cmplx(0.0_${k1}$, 0.0_${k1}$, kind=${k1}$))

        res = cmplx(beta_dist_rvs_r${k1}$(a%re, b%re) + loc_%re,                 &
                    beta_dist_rvs_r${k1}$(a%im, b%im) + loc_%im, kind=${k1}$)
    end function beta_dist_rvs_${t1[0]}$${k1}$

    #:endfor


    #:for k1, t1 in BETA_REAL_KINDS_TYPES
    function beta_dist_rvs_array_${t1[0]}$${k1}$(a, b, array_size, loc)     &
        result(res)
    !
        ${t1}$, intent(in) :: a, b
        integer, intent(in) :: array_size
        ${t1}$, intent(in), optional :: loc
        ${t1}$ :: res(array_size)
        integer :: i
        ${t1}$ :: loc_

        loc_ = optval(loc, 0.0_${k1}$)

        do i = 1, array_size
            res(i) = beta_dist_rvs_${t1[0]}$${k1}$(a, b, loc_)
        end do
    end function beta_dist_rvs_array_${t1[0]}$${k1}$

    #:endfor


    #:for k1, t1 in BETA_CMPLX_KINDS_TYPES
    function beta_dist_rvs_array_${t1[0]}$${k1}$(a, b, array_size, loc)     &
        result(res)
    ! Complex parameter beta distributed. The real part and imaginary part are
    ! independent of each other.
    !
        ${t1}$, intent(in) :: a, b
        integer, intent(in) :: array_size
        ${t1}$, intent(in), optional :: loc
        ${t1}$ :: res(array_size)
        integer :: i
        ${t1}$ :: loc_

        loc_ = optval(loc, cmplx(0.0_${k1}$, 0.0_${k1}$, kind=${k1}$))

        do i = 1, array_size
            res(i) = beta_dist_rvs_${t1[0]}$${k1}$(a, b, loc_)
        end do
    end function beta_dist_rvs_array_${t1[0]}$${k1}$

    #:endfor


    #:for k1, t1 in BETA_REAL_KINDS_TYPES
    elemental function beta_dist_pdf_${t1[0]}$${k1}$(x, a, b, loc)     &
        result(res)
    ! Beta distribution probability density function
    !
        ${t1}$, intent(in) :: x, a, b
        ${t1}$, intent(in), optional :: loc
        real(${k1}$) :: res
        ${t1}$ :: xs
        ${t1}$ :: loc_
        ${t1}$, parameter :: zero = 0.0_${k1}$, one = 1.0_${k1}$

        if(a <= zero .or. b <= zero) then
            res = ieee_value(1.0_${k1}$, ieee_quiet_nan)
            return
        end if

        loc_ = optval(loc, 0.0_${k1}$)
        xs = x - loc_

        if(xs <= zero .or. xs >= one) then
            res = zero
            return
        end if

        ! Use log formulation for numerical stability
        res = exp((a - one) * log(xs) + (b - one) * log(one - xs) - log_beta(a, b))
    end function beta_dist_pdf_${t1[0]}$${k1}$

    #:endfor


    #:for k1, t1 in BETA_REAL_KINDS_TYPES
    impure elemental function beta_dist_pdf_impure_${t1[0]}$${k1}$(x, a, b, err, loc) &
        result(res)
    ! Beta distribution probability density function (impure wrapper)
    !
        ${t1}$, intent(in) :: x, a, b
        integer, intent(out) :: err
        ${t1}$, intent(in), optional :: loc
        real(${k1}$) :: res

        res = beta_dist_pdf_${t1[0]}$${k1}$(x, a, b, loc)
        err = merge(1, 0, ieee_is_nan(res))
    end function beta_dist_pdf_impure_${t1[0]}$${k1}$

    #:endfor


    #:for k1, t1 in BETA_CMPLX_KINDS_TYPES
    elemental function beta_dist_pdf_${t1[0]}$${k1}$(x, a, b, loc)     &
        result(res)
    ! Complex parameter beta distributed. The real part and imaginary part are
    ! independent of each other.
    !
        ${t1}$, intent(in) :: x, a, b
        ${t1}$, intent(in), optional :: loc
        real(${k1}$) :: res
        ${t1}$ :: loc_

        loc_ = optval(loc, cmplx(0.0_${k1}$, 0.0_${k1}$, kind=${k1}$))

        res = beta_dist_pdf_r${k1}$(x%re, a%re, b%re, loc_%re)
        res = res * beta_dist_pdf_r${k1}$(x%im, a%im, b%im, loc_%im)
    end function beta_dist_pdf_${t1[0]}$${k1}$

    #:endfor


    #:for k1, t1 in BETA_CMPLX_KINDS_TYPES
    impure elemental function beta_dist_pdf_impure_${t1[0]}$${k1}$(x, a, b, err, loc) &
        result(res)
    ! Complex parameter beta distributed. The real part and imaginary part are
    ! independent of each other.
    !
        ${t1}$, intent(in) :: x, a, b
        integer, intent(out) :: err
        ${t1}$, intent(in), optional :: loc
        real(${k1}$) :: res
        ${t1}$ :: loc_

        loc_ = optval(loc, cmplx(0.0_${k1}$, 0.0_${k1}$, kind=${k1}$))

        res = beta_dist_pdf_r${k1}$(x%re, a%re, b%re, loc_%re)
        res = res * beta_dist_pdf_r${k1}$(x%im, a%im, b%im, loc_%im)
        err = merge(1, 0, ieee_is_nan(res))
    end function beta_dist_pdf_impure_${t1[0]}$${k1}$

    #:endfor


    #:for k1, t1 in BETA_REAL_KINDS_TYPES
    elemental function beta_dist_cdf_${t1[0]}$${k1}$(x, a, b, loc)   &
        result(res)
    ! Beta distribution cumulative distribution function
    !
        ${t1}$, intent(in) :: x, a, b
        ${t1}$, intent(in), optional :: loc
        real(${k1}$) :: res
        ${t1}$ :: xs
        ${t1}$ :: loc_
        ${t1}$, parameter :: zero = 0.0_${k1}$, one = 1.0_${k1}$

        loc_ = optval(loc, 0.0_${k1}$)
        xs = x - loc_

        if(a <= zero .or. b <= zero) then
            res = ieee_value(1.0_${k1}$, ieee_quiet_nan)
            return
        end if

        if(xs <= zero) then
            res = zero
            return
        end if

        if(xs >= one) then
            res = one
            return
        end if

        res = real(incomplete_beta(xs, a, b), kind=${k1}$)
    end function beta_dist_cdf_${t1[0]}$${k1}$

    #:endfor


    #:for k1, t1 in BETA_REAL_KINDS_TYPES
    impure elemental function beta_dist_cdf_impure_${t1[0]}$${k1}$(x, a, b, err, loc) &
        result(res)
    ! Beta distribution cumulative distribution function (impure wrapper)
    !
        ${t1}$, intent(in) :: x, a, b
        integer, intent(out) :: err
        ${t1}$, intent(in), optional :: loc
        real(${k1}$) :: res

        res = beta_dist_cdf_${t1[0]}$${k1}$(x, a, b, loc)
        err = merge(1, 0, ieee_is_nan(res))
    end function beta_dist_cdf_impure_${t1[0]}$${k1}$

    #:endfor


    #:for k1, t1 in BETA_CMPLX_KINDS_TYPES
    elemental function beta_dist_cdf_${t1[0]}$${k1}$(x, a, b, loc)   &
        result(res)
    ! Complex parameter beta distributed. The real part and imaginary part are
    ! independent of each other.
    !
        ${t1}$, intent(in) :: x, a, b
        ${t1}$, intent(in), optional :: loc
        real(${k1}$) :: res
        ${t1}$ :: loc_

        loc_ = optval(loc, cmplx(0.0_${k1}$, 0.0_${k1}$, kind=${k1}$))

        res = beta_dist_cdf_r${k1}$(x%re, a%re, b%re, loc_%re)
        res = res * beta_dist_cdf_r${k1}$(x%im, a%im, b%im, loc_%im)
    end function beta_dist_cdf_${t1[0]}$${k1}$

    #:endfor


    #:for k1, t1 in BETA_CMPLX_KINDS_TYPES
    impure elemental function beta_dist_cdf_impure_${t1[0]}$${k1}$(x, a, b, err, loc) &
        result(res)
    ! Complex parameter beta distributed. The real part and imaginary part are
    ! independent of each other.
    !
        ${t1}$, intent(in) :: x, a, b
        integer, intent(out) :: err
        ${t1}$, intent(in), optional :: loc
        real(${k1}$) :: res
        ${t1}$ :: loc_

        loc_ = optval(loc, cmplx(0.0_${k1}$, 0.0_${k1}$, kind=${k1}$))

        res = beta_dist_cdf_${t1[0]}$${k1}$(x, a, b, loc_)
        err = merge(1, 0, ieee_is_nan(res))
    end function beta_dist_cdf_impure_${t1[0]}$${k1}$

    #:endfor


end module stdlib_stats_distribution_beta