module dists

    use :: iso_c_binding
    use :: ieee_arithmetic
    use prob, only : pdists
    use :: fgsl
    use :: utilities, only : FatalErrorMessage
  
    implicit none
  
    private
  
    integer, parameter :: dp = c_double
    integer, parameter :: i4 = c_int
    integer, parameter :: i8 = c_int64_t

    type, abstract :: continuous_dist
        real(kind=dp) :: mean
        real(kind=dp) :: median
        real(kind=dp) :: mode
        real(kind=dp) :: variance
    contains
        procedure(pdffunc), deferred :: pdf
        procedure(cdffunc), deferred :: cdf
        procedure(quantilefunc), deferred :: quantile
        procedure(samplefunc), deferred :: sample
        procedure, pass(self) :: inverse_transform_sample_scalar
        generic, public :: sample_from_CDF => inverse_transform_sample_scalar
    end type continuous_dist

    abstract interface
        function pdffunc( self, x ) result(p)
            import continuous_dist
            import dp
            class(continuous_dist), intent(in) :: self
            real(kind=dp), intent(in) :: x
            real(kind=dp) :: p
        end function pdffunc

        function cdffunc( self, x ) result(F)
            import continuous_dist
            import dp
            class(continuous_dist), intent(in) :: self
            real(kind=dp), intent(in) :: x
            real(kind=dp) :: F
        end function cdffunc

        function quantilefunc( self, p ) result(x)
            import continuous_dist
            import dp
            class(continuous_dist), intent(in) :: self
            real(kind=dp), intent(in) :: p
            real(kind=dp) :: x
        end function quantilefunc

        function samplefunc( self ) result(rvs)
            import continuous_dist
            import dp
            class(continuous_dist), intent(in) :: self
            real(kind=dp) :: rvs
        end function samplefunc
    end interface

    type :: discrete_dist
        real(kind=dp) :: mean
        integer(kind=i4) :: median
        integer(kind=i4) :: mode
        real(kind=dp) :: variance
    end type discrete_dist

    type, public :: empirical_dist
        real(kind=dp), dimension(:), allocatable :: data_x
        real(kind=dp), dimension(:), allocatable :: data_p
        type(fgsl_interp), private :: interp_cdf, interp_quantile
    contains
        private
        procedure, public, pass(self) :: cdf => piecewise_linear_cdf
        procedure, public, pass(self) :: pdf => piecewise_linear_pdf
        procedure, public, pass(self) :: quantile => piecewise_linear_quantile
        procedure, public, pass(self) :: sample => piecewise_linear_sample
        final :: destroy_empirical
    end type empirical_dist

    interface empirical_dist
        module procedure init_empirical_dist, init_empirical_dist_weighted_combine
    end interface empirical_dist

    type, public, extends(continuous_dist) :: normal
        real(kind=dp) :: mu
        real(kind=dp) :: sigma
    contains
        private
        procedure, public, pass(self) :: pdf => normal_pdf
        procedure, public, pass(self) :: cdf => normal_cdf
        procedure, public, pass(self) :: quantile => normal_quantile
        procedure, public, pass(self) :: sample => normal_sample
    end type normal

    interface normal
        module procedure init_normal
    end interface normal

    type, public, extends(continuous_dist) :: lognormal
        real(kind=dp) :: mu
        real(kind=dp) :: sigma
    contains
        private
        procedure, public, pass(self) :: pdf => lognormal_pdf
        procedure, public, pass(self) :: cdf => lognormal_cdf
        procedure, public, pass(self) :: quantile => lognormal_quantile
        procedure, public, pass(self) :: sample => lognormal_sample
    end type lognormal

    interface lognormal
        module procedure init_lognormal
    end interface lognormal

    type, public, extends(continuous_dist) :: tdist
        real(kind=dp) :: df
        real(kind=dp) :: location
        real(kind=dp) :: scale
    contains
        private
        procedure, public, pass(self) :: pdf => tdist_pdf
        procedure, public, pass(self) :: cdf => tdist_cdf
        procedure, public, pass(self) :: quantile => tdist_quantile
        procedure, public, pass(self) :: sample => tdist_sample
    end type tdist

    interface tdist
        module procedure init_tdist
    end interface tdist

    type, public, extends(continuous_dist) :: TOST
        ! The piecewise skew-logistic distribution for TOST 
        ! from 
        ! The timing of COVID-19 transmission
        ! Ferretti et al.
        ! https://www.medrxiv.org/content/10.1101/2020.09.04.20188516v2
        real(kind=dp) :: mu ! = -4.0 days
        real(kind=dp) :: sigma ! = 1.85 days
        real(kind=dp) :: alpha ! = 5.85
        real(kind=dp) :: tau ! =5.42 (=mean incubation period)
        real(kind=dp) :: ti ! incubation time
        real(kind=dp) :: k ! normalization factor
    contains
        private
        procedure, public, pass(self) :: pdf => TOST_pdf
        procedure, public, pass(self) :: cdf => TOST_cdf
        procedure, public, pass(self) :: quantile => TOST_quantile
        procedure, public, pass(self) :: sample => TOST_sample
    end type TOST

    interface TOST
        module procedure init_TOST
    end interface TOST

    type, public, extends(discrete_dist) :: bernoulli
        real(kind=dp) :: p
    contains
        private
        procedure, public, pass(self) :: sample => bernoulli_sample
    end type bernoulli

    interface bernoulli
        module procedure init_bernoulli
    end interface bernoulli

contains

    function inverse_transform_sample_scalar(self) result(x)
        implicit none
        class(continuous_dist), intent(in) :: self
        real(kind=dp) :: x

        real(kind=dp) :: u

        call pdists%unif_rvs(u)

        x = self%quantile(u)
        return
    end function inverse_transform_sample_scalar

    function init_normal(mu, sigma) result(self)
        implicit none
        type(normal) :: self
        real(kind=dp), intent(in) :: mu
        real(kind=dp), intent(in) :: sigma
        self%mu = mu
        self%sigma = sigma

        self%mean = mu
        self%median = mu
        self%mode = mu
        self%variance = sigma*sigma
        return
    end function init_normal

    function normal_sample(self) result(rvs)
        implicit none
        class(normal), intent(in) :: self
        real(kind=dp) :: rvs
        rvs = pdists%normal_rvs(self%mu, self%sigma)
        return
    end function normal_sample

    function normal_pdf(self, x) result(p)
        implicit none
        class(normal), intent(in) :: self
        real(kind=dp), intent(in) :: x
        real(kind=dp) :: p
        p = fgsl_ran_gaussian_pdf(x-self%mu, self%sigma)
        return
    end function normal_pdf

    function normal_cdf(self, x) result(F)
        implicit none
        class(normal), intent(in) :: self
        real(kind=dp), intent(in) :: x
        real(kind=dp) :: F
        F = fgsl_cdf_gaussian_P(x-self%mu, self%sigma)
        return
    end function normal_cdf

    function normal_quantile(self, p) result(x)
        implicit none
        class(normal), intent(in) :: self
        real(kind=dp), intent(in) :: p
        real(kind=dp) :: x
        x = fgsl_cdf_gaussian_Pinv(p, self%sigma) + self%mu
        return
    end function normal_quantile
    
    function init_lognormal(mu, sigma) result(self)
        implicit none
        type(lognormal) :: self
        real(kind=dp), intent(in) :: mu
        real(kind=dp), intent(in) :: sigma

        real(kind=dp) :: var

        self%mu = mu
        self%sigma = sigma

        var = sigma*sigma

        self%mean = exp(mu + 0.5_dp*var)
        self%median = exp(mu)
        self%mode = exp(mu - var)
        self%variance = (exp(var)-1.0_dp)*exp(2.0_dp*mu + var)
        return
    end function init_lognormal

    function lognormal_sample(self) result(rvs)
        implicit none
        class(lognormal), intent(in) :: self
        real(kind=dp) :: rvs
        rvs = pdists%lognormal_rvs(self%mu, self%sigma)
        return
    end function lognormal_sample

    function lognormal_pdf(self, x) result(p)
        implicit none
        class(lognormal), intent(in) :: self
        real(kind=dp), intent(in) :: x
        real(kind=dp) :: p
        p = fgsl_ran_lognormal_pdf(x, self%mu, self%sigma)
        return
    end function lognormal_pdf

    function lognormal_cdf(self, x) result(F)
        implicit none
        class(lognormal), intent(in) :: self
        real(kind=dp), intent(in) :: x
        real(kind=dp) :: F
        F = fgsl_cdf_lognormal_P(x, self%mu, self%sigma)
        return
    end function lognormal_cdf

    function lognormal_quantile(self, p) result(x)
        implicit none
        class(lognormal), intent(in) :: self
        real(kind=dp), intent(in) :: p
        real(kind=dp) :: x
        x = fgsl_cdf_lognormal_Pinv(p, self%mu, self%sigma)
        return
    end function lognormal_quantile

    function init_tdist(df, location, scale) result(self)
        implicit none
        type(tdist) :: self
        real(kind=dp), intent(in) :: df
        real(kind=dp), intent(in), optional :: location
        real(kind=dp), intent(in), optional :: scale
        self%df = df

        if (present(location)) then
            self%location = location
        else
            self%location = 0.0_dp
        end if

        if (present(scale)) then
            self%scale = scale
        else
            self%scale = 1.0_dp
        end if

        if (self%df>1.0_dp) then
            self%mean = self%location
        else
            self%mean = ieee_value(1.0_dp, ieee_quiet_nan)
        end if
        self%median = self%location
        self%mode = self%location
        if (self%df > 2.0_dp) then
            self%variance = (self%scale**2)*self%df/(self%df-2.0_dp)
        elseif (self%df > 1.0_dp) then
            self%variance = ieee_value(1.0_dp, ieee_positive_inf)
        else
            self%variance = ieee_value(1.0_dp, ieee_quiet_nan)
        end if
        return
    end function init_tdist

    function tdist_sample(self) result(rvs)
        implicit none
        class(tdist), intent(in) :: self
        real(kind=dp) :: rvs
        rvs = self%location + self%scale*pdists%t_rvs(self%df)
        return
    end function tdist_sample

    function tdist_pdf(self, x) result(p)
        implicit none
        class(tdist), intent(in) :: self
        real(kind=dp), intent(in) :: x
        real(kind=dp) :: p

        real(kind=dp) :: y

        y = (x-self%location)/self%scale

        p = fgsl_ran_tdist_pdf(y, self%df)/self%scale
        return
    end function tdist_pdf

    function tdist_cdf(self, x) result(F)
        implicit none
        class(tdist), intent(in) :: self
        real(kind=dp), intent(in) :: x
        real(kind=dp) :: F

        real(kind=dp) :: y

        y = (x-self%location)/self%scale

        F = fgsl_cdf_tdist_P(y, self%df)
        return
    end function tdist_cdf

    function tdist_quantile(self, p) result(x)
        implicit none
        class(tdist), intent(in) :: self
        real(kind=dp), intent(in) :: p
        real(kind=dp) :: x

        real(kind=dp) :: y

        y = fgsl_cdf_tdist_Pinv(p, self%df)

        x = self%location + self%scale*y
        return
    end function tdist_quantile

    function init_bernoulli(p) result(self)
        implicit none
        type(bernoulli) :: self
        real(kind=dp), intent(in) :: p
        self%p = p

        self%mean = p
        self%median = merge(0,1, p<=0.5_dp)
        self%mode = merge(0,1, p<=0.5_dp)
        self%variance = p*(1.0_dp-p)
        return
    end function init_bernoulli

    function bernoulli_sample(self) result(rvs)
        implicit none
        class(bernoulli), intent(in) :: self
        logical :: rvs
        rvs = pdists%bernoulli_rvs(self%p)
        return
    end function bernoulli_sample

    function init_TOST(ti, mu, sigma, alpha, tau) result(self)
        implicit none
        type(TOST) :: self
        real(kind=dp), intent(in) :: ti ! incubation time
        real(kind=dp), intent(in), optional :: mu ! = -4.0 days
        real(kind=dp), intent(in), optional :: sigma ! = 1.85 days
        real(kind=dp), intent(in), optional :: alpha ! = 5.85
        real(kind=dp), intent(in), optional :: tau ! =5.42 (=mean incubation period)
        
        real(kind=dp) :: expm
        real(kind=dp) :: expt

        self%ti = ti

        if (present(mu)) then
            self%mu = mu
        else
            self%mu = -4.0_dp
        end if
        if (present(sigma)) then
            self%sigma = sigma
        else
            self%sigma = 1.85_dp
        end if
        if (present(alpha)) then
            self%alpha = alpha
        else
            self%alpha = 5.85_dp
        end if
        if (present(tau)) then
            self%tau = tau
        else
            self%tau = 5.42_dp
        end if
        
        expm = exp(self%mu/self%sigma)
        expt = exp((self%tau+self%mu)/self%sigma)

        self%k = self%sigma/self%alpha * ( (self%ti/self%tau - 1.0_dp)*((1.0_dp + expm)**(-self%alpha)) + 1.0_dp &
                                     - self%ti/self%tau*(1.0_dp + expt)**(-self%alpha) )

        return    
    end function init_TOST

    function TOST_pdf(self, x) result(p)
        implicit none
        class(TOST), intent(in) :: self
        real(kind=dp), intent(in) :: x
        real(kind=dp) :: p

        real(kind=dp) :: expf

        if (x < -self%ti) then
            p = 0.0_dp
        elseif (x < 0.0_dp) then
            expf = exp(-(x*self%tau/self%ti-self%mu)/self%sigma)
            p = expf/((1.0_dp + expf)**(self%alpha+1.0_dp))/self%k
        else
            expf = exp(-(x-self%mu)/self%sigma)
            p = expf/((1.0_dp + expf)**(self%alpha+1.0_dp))/self%k
        end if

        return
    end function TOST_pdf

    function TOST_cdf(self, x) result(F)
        implicit none
        class(TOST), intent(in) :: self
        real(kind=dp), intent(in) :: x
        real(kind=dp) :: F

        real(kind=dp) :: expm
        real(kind=dp) :: expt
        real(kind=dp) :: expf

        expm = exp(self%mu/self%sigma)
        expt = exp((self%tau+self%mu)/self%sigma)

        if (x < -self%ti) then
            F = 0.0_dp
        elseif (x < 0.0_dp) then
            expf = exp(-(x*self%tau/self%ti-self%mu)/self%sigma)
            F = self%sigma/self%alpha*self%ti/self%tau/self%k*( (1.0_dp + expf)**(-self%alpha) &
                                                               - (1.0_dp + expt)**(-self%alpha) )
        else
            expf = exp(-(x-self%mu)/self%sigma)
            F = self%sigma/self%alpha*self%ti/self%tau/self%k*( (1.0_dp+expm)**(-self%alpha) &
                                                                - (1.0_dp + expt)**(-self%alpha) ) &
                + self%sigma/self%alpha/self%k * ( (1.0_dp + expf)**(-self%alpha) - (1.0_dp + expm)**(-self%alpha) )
        end if

        return
    end function TOST_cdf

    function TOST_quantile(self, p) result(x)
        implicit none
        class(TOST), intent(in) :: self
        real(kind=dp), intent(in) :: p
        real(kind=dp) :: x

        real(kind=dp) :: expm
        real(kind=dp) :: expt
        real(kind=dp) :: p0

        expm = exp(self%mu/self%sigma)
        expt = exp((self%tau + self%mu)/self%sigma)

        p0 = self%sigma/self%alpha/self%k*self%ti/self%tau*( (1.0_dp+expm)**(-self%alpha) - (1.0_dp+expt)**(-self%alpha) )

        if (p < p0) then
            x = self%ti/self%tau*( self%mu & 
                                - self%sigma*log( (self%k*self%alpha/self%sigma*self%tau/self%ti*p &
                                                   + (1.0_dp + expt)**(-self%alpha) )**(-1.0_dp/self%alpha) &
                                                 - 1.0_dp ) )
        else
            x = self%mu - self%sigma*log( (self%k*self%alpha/self%sigma*p &
                                           + (1.0_dp - self%ti/self%tau)*(1+expm)**(-self%alpha) &
                                           + self%ti/self%tau*(1.0_dp + expt)**(-self%alpha))**(-1.0_dp/self%alpha) &
                                         - 1.0_dp)
        end if

        return
    end function TOST_quantile

    function TOST_sample(self) result(rvs)
        class(TOST), intent(in) :: self
        real(kind=dp) :: rvs
        rvs = self%inverse_transform_sample_scalar()
        return
    end function TOST_sample

    function init_empirical_dist(data_x, data_p) result(self)
        real(kind=dp), dimension(:), intent(in) :: data_x
        real(kind=dp), dimension(:), intent(in) :: data_p
        type(empirical_dist) :: self
        integer(kind=c_int) :: status

        integer(kind=c_int64_t) :: n
         
        if (size(data_x) .ne. size(data_p)) call FatalErrorMessage('In initializer for empirical distribution, data arrays must be same length')

        allocate(self%data_x, source=data_x)
        allocate(self%data_p, source=data_p)

        n = size(data_x)

        self%interp_cdf = fgsl_interp_alloc(FGSL_INTERP_LINEAR, n)
        status = fgsl_interp_init(self%interp_cdf, self%data_x, self%data_p)

        self%interp_quantile = fgsl_interp_alloc(FGSL_INTERP_LINEAR, n)
        status = fgsl_interp_init(self%interp_quantile, self%data_p, self%data_x)

        return

    end function init_empirical_dist

    function init_empirical_dist_weighted_combine(ecdf1, ecdf2, weight) result(self)
        class(empirical_dist), intent(in) :: ecdf1
        class(empirical_dist), intent(in) :: ecdf2
        real(kind=c_double), intent(in) :: weight
        type(empirical_dist) :: self

        integer(kind=c_int) :: status

        integer(kind=c_int64_t) :: n = 100
        integer(kind=c_int64_t) :: j

        allocate(self%data_x(n), self%data_p(n))

        do j=1,n
            self%data_p(j) = (j-1.0_dp)/(n-1.0_dp)
            self%data_x(j) = weight*ecdf1%quantile(self%data_p(j)) + (1.0_dp - weight)*ecdf2%quantile(self%data_p(j))
        end do

        self%interp_cdf = fgsl_interp_alloc(FGSL_INTERP_LINEAR, n)
        status = fgsl_interp_init(self%interp_cdf, self%data_x, self%data_p)

        self%interp_quantile = fgsl_interp_alloc(FGSL_INTERP_LINEAR, n)
        status = fgsl_interp_init(self%interp_quantile, self%data_p, self%data_x)

        return
    end function init_empirical_dist_weighted_combine

    function piecewise_linear_cdf(self, x) result(p)
        class(empirical_dist), intent(in) :: self
        real(kind=dp), intent(in) :: x
        real(kind=dp) :: p

        type(fgsl_interp_accel) :: accl_cdf

        accl_cdf = fgsl_interp_accel_alloc()
        p = fgsl_interp_eval(self%interp_cdf, self%data_x, self%data_p, x, accl_cdf)
        call fgsl_interp_accel_free(accl_cdf)

        return
    end function piecewise_linear_cdf

    function piecewise_linear_pdf(self, x) result(p)
        class(empirical_dist), intent(in) :: self
        real(kind=dp), intent(in) :: x
        real(kind=dp) :: p

        type(fgsl_interp_accel) :: accl_cdf

        accl_cdf = fgsl_interp_accel_alloc()
        p = fgsl_interp_eval_deriv(self%interp_cdf, self%data_x, self%data_p, x, accl_cdf)
        call fgsl_interp_accel_free(accl_cdf)
        return
    end function piecewise_linear_pdf

    function piecewise_linear_quantile(self, p) result(x)
        class(empirical_dist), intent(in) :: self
        real(kind=dp), intent(in) :: p
        real(kind=dp) :: x

        type(fgsl_interp_accel) :: accl_quantile

        integer(kind=c_int64_t) :: e

        accl_quantile = fgsl_interp_accel_alloc()
        e = fgsl_interp_eval_e(self%interp_quantile, self%data_p, self%data_x, p, accl_quantile, x)
        
        if (e .ne. 0_i8) then
            print *, 'Error in piecewise_linear_quantile'
            print *, p, x, e, minval(self%data_p), maxval(self%data_p)    
        end if
        call fgsl_interp_accel_free(accl_quantile)
        return
    end function piecewise_linear_quantile

    function piecewise_linear_sample(self) result(rvs)
        class(empirical_dist), intent(in) :: self
        real(kind=dp) :: rvs

        real(kind=dp) :: u

        call pdists%unif_rvs(u)

        rvs = self%quantile(u)
        return
    end function piecewise_linear_sample

    subroutine destroy_empirical(self)
        type(empirical_dist) :: self
        
        if (allocated(self%data_x)) deallocate(self%data_x)
        if (allocated(self%data_p)) deallocate(self%data_p)

    end subroutine destroy_empirical

end module dists