/* 

  Probabilistic Programming Distributions in Picat.
  
  PPL: Probabilistic Programming Light
  
  This model contains probability distributions. Most of them
  supports the following scheme:

  - Random generation: <dist>_dist(Parameters)
  - Random generation (n samples): <dist>_dist_n(Parameters,N)   
  - PDF/PMF          : <dist>_dist_pdf(Parameters,X)   
  - CDF              : <dist>_dist_cdf(Parameters,X)   
  - Quantile         : <dist>_dist_quantile(Parameters,Q)
  - Mean             : <dist>_dist_mean(Parameters)   
  - Variance         : <dist>_dist_variance(Parameters)   

  The following distributions has been implemented:
  (Cf the list in gamble_distributions.rkt)
  - benford_dist: dist, pdf, cdf, quantile, mean
  - bernoulli_dist: dist, pdf, cdf, quantile
    aliased as bern
  - beta_binomial: dist, pdf, cdf, quantile
  - beta_negative_binomial: dist, pdf, cdf, quantile, mean, variance
  - beta_prime/2: dist, pdf, cdf, quantile, mean, variance
  - beta_prime/3: dist, pdf, cdf, quantile, mean, variance
  - beta_prime/4: dist, pdf, cdf, quantile, mean, variance
  - binomial_dist: dist, pdf, cdf, quantile, mean, variance
  - binomial_process: dist, pdf, cdf, quantile, mean, variance
  - birthday_dist: dist, pdf, quantile,median 
    (also as rbirthday,pbirthday,qbirthday)
  - categorical_dist, dist, pdf, cdf, (quantile), mean, variance
  - cauchy_dist: dist, pdf, cdf, quantile, mean, variance
  - chi_dist: dist, pdf, cdf, quantile, mean, variance
  - chi_squared_dist: dist, pdf, cdf, quantile (est), mean
  - coupon_collector_dist: dist, pdf, cdf, quantile, mean, variance
  - crp (Chinese Restaurant Process): dist, pdf, cdf, quantile, mean
  - dice: dist, pdf, cdf, quantile, mean, variance
  - dirichlet_dist: dist, pdf, cdf, quantile, mean, variance
  - discrete_laplace: dist, pdf, cdf, quantile, mean, variance
  - discrete_markov_process: dist, pdf, cdf, quantile (and stationary)
  - discrete_uniform: dist, pdf, cdf, quantile, mean, variance
  - erlang_dist: dist, pdf, cdf, mean
  - exponential_dist: dist, pdf, cdf, quantile, mean
  - extreme_value_dist: dist, pdf, cdf, quantile, mean, variance
  - f_dist: dist, pdf, cdf, quantile, mean, variance
  - flip/0: dist, pdf, cdf, quantile, mean, variance
  - flip/1: dist, pdf, cdf, quantile, mean, variance
  - frechet_dist/2: dist, pdf, cdf, quantile, mean, variance
  - frechet_dist/3: dist, pdf, cdf, quantile, mean, variance
  - gamma_dist: dist, pdf, cdf, quantile (est), mean
  - geometric_dist: dist, pdf, cdf, quantile, mean, variance
  - generalized_extreme_value_dist: dist, pdf, cdf, quantile, mean
  - gumbel_dist: dist, pdf, cdf, quantile, mean, variance
  - hypergeometric1: dist, pdf, cdf, quantile, mean
  - hypergeometric: dist, pdf, cdf, quantile
  - k_record: random, pdf, cdf, quantile (the quantile is restricted)
  - kumaraswamy_dist: dist, pdf, cdf, quantile, mean
  - laplace: dist, pdf, cdf, quantile
  - lognormal: dist, pdf, cdf, quantile, mean, variance
  - log_gamma: dist, pdf, cdf, quantile, mean, variance
  - logistic/0: dist, pdf, cdf, quantile, mean, variance
  - logistic/2: dist, pdf, cdf, quantile, mean, variance
  - logseries: dist, pdf, cdf, quantile, mean, variance
  - matching: dist, pdf, cdf, quantile, mean, variance
  - max_stable: dist, pdf, cdf, quantile, mean, variance
  - min_stable: dist, pdf, cdf, quantile, mean, variance
  - multinomial_dist: dist, pdf, mean, variance
  - multivariate_hypergeometric: dist, pdf, (cdf), mean
  - negative_binomial: dist, pdf, cdf, quantile
  - negative_hypergeometric: dist, pdf, cdf, quantile, mean
  - normal_dist: dist, pdf, cdf, quantile, mean, variance
  - num_records_dist: dist, pdf, cdf, quantile, mean (and harmonic number estimator)
  - order statistics: estimator_of_m_u, estimator_of_m_u_all, estimator_of_m_v
  - order_statistics_continuous: pdf, cdf, quantile, median
  - order_statistics_with_replacement_discrete: pdf, cdf
  - order_statistics_without_replacement: dist, pdf, cdf, quantile, mean, variance
  - pareto1: dist, pdf, cdf, quantile, mean, variance
  - pareto2: dist, pdf, cdf, quantile, mean, variance
  - pareto3: dist, pdf, cdf, quantile, mean, variance
  - pareto4: dist, pdf, cdf, quantile, mean, variance
  - pascal_dist: dist, pdf, cdf, quantile, mean, variance
  - poisson_dist: dist, pdf, cdf, quantile, mean
  - poisson_process: dist, pdf, cdf, quantile, mean, variance
  - prob-n-heads-after-k-in-max-m-tosses: dist, pdf, cdf, quantile, mean, variance
  - probability-of-run-size: (not a proper distribution)
  - rademacher_dist: dist, pdf, cdf, quantile
  - random_integer: dist, pdf, cdf, quantile, mean, variance
  - random_integer1: dist, pdf, cdf, quantile, mean, variance
  - random_walk_process_dist: dist, pdf, cdf, quantile, mean, variance
  - shifted_geometric_dist: dist, pdf, cdf, quantile, mean, variance
  - student_t/1: dist, pdf, cdf, quantile, mean, variance  
  - student_t/3: dist, pdf, cdf, quantile, mean, variance  
  - sum_prob: dist, pdf, cdf, quantile, mean, variance
  - triangular_dist/[0,2]: dist, pdf, cdf, quantile, mean
    triangular2_dist: dist, pdf, cdf, quantile, mean
  - uniform: dist, pdf, cdf, quantile, mean, variance
  - weibull/2: dist, pdf, cdf, quantile, mean, variance
  - weibull/3: dist, pdf, cdf, quantile, mean, variance
  - wiener_process: dist, pdf, cdf, quantile, mean, variance
  - zipf_dist/1: dist, pdf, cdf, quantile, mean, variance 
    zipf_mma_dist/1: dist, pdf, cdf, quantile, mean, variance 
    zipf_dist/2: dist, pdf, cdf, quantile, mean, variance
    zipf_mma_dist/2: dist, pdf, cdf, quantile, mean, variance

  TODO:
  - Skellan
  - Raleigh
  - Yule-Simon (WaringYuleDistribution in Mathematica)
  - Zipf–Mandelbrot


  These are only random generation.
  TODO: make PDF, CDF etc
  - chi_squared_inverse_dist: dist
  - geometric_zero_truncated_dist: dist
  - inverse_exponential: dist
  - inverse_exponential_shifted: dist
  - polya: dist
  - polya_eggenberg: dist


$ perl -nle 'print $1 if m!^([a-z][^\s]+?) =!' ppl_distributions.pi | sort
add_scaled_row(W,Row,Scale,J)
adjust_down(C,K,N,P)
adjust_up(C,K,N,P)
benford_dist(B)
benford_dist_mean(B)
benford_dist_n(B,N)
bern_n(P,N)
bernoulli_dist_mean(P)
bernoulli_dist_n(P,N)
bernoulli_dist(P)
bernoulli_dist_rand(P)
bernoulli_dist_sum(P,N)
bernoulli_dist_variance(P)
bern(P)
bern_sum(P,N)
betacf(A,B,X)
beta_func(A,B)
binary_search_quantile(CDF,TargetP,Tolerance,Lower,Upper)
binomial_dist1(N,P)
binomial_dist_cdf(N,P,K)
binomial_distf_cdf(N,P,K)
binomial_distf_pdf(N,P,K)
binomial_distf_quantile(N,P,Q)
binomial_dist_mean(N,P)
binomial_dist_n(N,P,Num)
binomial_dist(N,P)
binomial_dist_pdf(N,P,K)
binomial_dist_quantile(N,P,Q)
binomial_dist_smart_n(N,P,Num)
binomial_dist_smart(N,P)
binomial_dist_variance(N,P)
binomialf_float(N,K)
binomialf(N,K)
binomialg(N,K)
binomial_process_dist_cdf(P,T,X)
binomial_process_dist_mean(P,T)
binomial_process_dist_pdf(P,T,X)
binomial_process_dist(P,T)
binomial_process_dist_quantile(P,T,Q)
binomial_process_dist_variance(P,T)
birthday_dist()
birthday_dist(Classes,Coincident)
birthday_dist_median(Classes,Coincident)
birthday_dist_n(Classes,Coincident,N)
birthday_dist_n(N)
birthday_dist_pdf(Classes,Coincident,X)
birthday_dist_pdf(X)
birthday_dist_quantile(Classes,Coincident,X)
birthday_dist_quantile(X)
categorical_cdf(Probs,Values,X)
categorical_mean(Probs,Values)
categorical_pdf(Probs,Values,X)
categorical_quantile_index(Probs,X)
categorical_quantile_value(Probs,Values,X)
categorical_variance(Probs,Values)
cauchy_dist(A,B)
cauchy_dist_cdf(A,B,X)
cauchy_dist_mean(A,B)
cauchy_dist_n(A,B,N)
cauchy_dist_pdf(A,B,X)
cauchy_dist_quantile(A,B,X)
cauchy_dist_variance(A,B)
cdf_all(Dist,Params)
cdf_all(DistParams)
cdf_all(Dist,Params,FromQ,ToQ)
cdf_all(DistParams,FromQ,ToQ)
cdf(Dist,Params,X)
cdf(DistParams,X)
chi_dist_mean(Nu)
chi_dist(Nu)
chi_dist_variance(Nu)
chi_square_dist_mean(Nu)
chi_square_dist(Nu)
chi_square_dist_variance(Nu)
chi_square_inverse_dist(K)
chi_square_inverse_dist_n(K,N)
count_partitions_sum(A,B,N,S)
coupon_collector_dist_cdf(M,K,N)
coupon_collector_dist_mean_euler(N)
coupon_collector_dist_mean(M,K)
coupon_collector_dist(M,K)
coupon_collector_dist_n(M,K,N)
coupon_collector_dist_pdf(M,K,N)
coupon_collector_dist_quantile(M,K,Q)
coupon_collector_dist_variance(M,K)
coupon_collectors_problem_theoretical(N)
crp_dist_cdf(Theta,N,K)
crp_dist_mean(Theta,N)
crp_dist_n(Theta,N,Num)
crp_dist_pdf(Theta,N,K)
crp_dist_quantile(Theta,N,Q)
crp_dist(Theta,N)
dice()
dice(N)
dice_n(N)
dice_n(N,Num)
digamma_asymptotic(X)
digammaf(X)
dirichlet_dist(Alpha)
dirichlet_dist_cdf(_Alpha,_Xs)
dirichlet_dist_mean(Alpha)
dirichlet_dist_n(Alpha,N)
dirichlet_dist_pdf(Alpha,Xs)
dirichlet_dist_quantile(_Alpha,_P)
dirichlet_dist_variance(Alpha)
discrete_markov_process_dist_cdf(TM,Init,T,State)
discrete_markov_process_dist_n(TM,Init,T,N)
discrete_markov_process_dist_pdf_loop(TM,Init,T,State,N,K)
discrete_markov_process_dist_pdf(TM,Init,T,State)
discrete_markov_process_dist_quantile(TM,Init,Step,Q)
discrete_markov_process_dist(TM,Init,T)
dist_sigma(Dist,Sigma)
ensure_dirichlet_params(Alpha)
erf_approx(X)
erlang_dist_cdf(K,Lambda,X)
erlang_dist(K,Lambda)
erlang_dist_mean(K,Lambda)
erlang_dist_n(K,Lambda,N)
erlang_dist_pdf(K,Lambda,X)
erlang_dist_quantile(K,Lambda,X)
erlang_dist_quantile(K,Lambda,X)
erlang_dist_variance(K,Lambda)
euler_gamma
expected_tosses_needed_for_n_heads(N)
expm1(X)
exponential_dist(Lambda)
exponential_dist_mean(Lambda)
exponential_dist_n(Lambda,N)
exponential_dist_variance(Lambda)
extreme_value_dist()
extreme_value_dist(Alpha,Beta)
extreme_value_dist_cdf(X)
extreme_value_dist_mean()
extreme_value_dist_n(Alpha,Beta,N)
extreme_value_dist_n(N)
extreme_value_dist_pdf(X)
extreme_value_dist_quantile(X)
extreme_value_dist_variance()
factorialf(N)
factorial_mem(N)
flatten1(List)
flip()
flip_n(N)
flip(P)
gamma_func(X)
generalized_extreme_value_dist_cdf(Mu,Sigma,Xi,X)
generalized_extreme_value_dist_mean(Mu,Sigma,Xi)
generalized_extreme_value_dist(Mu,Sigma,Xi)
generalized_extreme_value_dist_n(Mu,Sigma,Xi,N)
generalized_extreme_value_dist_pdf(Mu,Sigma,Xi,X)
generalized_extreme_value_dist_quantile(Mu,Sigma,Xi,Q)
generalized_extreme_value_dist_reduced(Mu,Sigma,Xi)
generalized_extreme_value_dist_reduced_n(Mu,Sigma,Xi,N)
geometric1(P,X)
geometric_dist1(P)
geometric_dist_mean(P)
geometric_dist(P)
geometric_dist_variance(P)
geometric_zero_truncated1(P,N)
geometric_zero_truncated_dist_n(P,N)
geometric_zero_truncated_dist(P)
gumbel_dist()
gumbel_dist(Alpha,Beta)
gumbel_dist_cdf(X)
gumbel_dist_mean()
gumbel_dist_n(Alpha,Beta,N)
gumbel_dist_n(N)
gumbel_dist_pdf(X)
gumbel_dist_quantile(X)
gumbel_dist_variance()
harmonic_number(N)
hazard_function_cont(Dist,Params)
hazard_function_discrete(Dist,Params)
hypergeometric1_dist_bool(Kk,N,K,Nn)
hypergeometric1_dist_helper(Kk,N,K,Nn,Count)
hypergeometric1_dist(Kk,N,K,Nn)
hypergeometric1_dist_n(Kk,N,K,Nn,Num)
incbeta_reg(A,B,X)
inverse_exponential_dist(Lambda)
inverse_exponential_dist_n(Lambda,N)
invnorm_acklam_enforced(P)
invnorm_acklam(P)
invnorm_acklam_symmetric(P)
ith_smallest(Xs,I)
ith_smallest(Xs,I,Sorted)
k_record_dist_cdf(K,N)
k_record_distf_cdf(K,N)
k_record_distf_pdf(K,N)
k_record_distf_quantile(K,Q)
k_record_dist(K)
k_record_dist_mean(K,N)
k_record_dist_pdf(K,N)
k_record_dist_quantile(K,Q)
kumaraswamy_dist(Alpha,Beta)
kumaraswamy_dist_cdf(Alpha,Beta,X)
kumaraswamy_dist_mean(Alpha,Beta)
kumaraswamy_dist_n(Alpha,Beta,N)
kumaraswamy_dist_pdf(Alpha,Beta,X)
kumaraswamy_dist_quantile(Alpha,Beta,X)
kumaraswamy_dist_variance(Alpha,Beta)
laplace_dist_n(Mu,B,N)
lgamma(X)
log1p(X)
log_beta(A,B)
log_gamma_dist(Alpha,Beta,Mu)
log_gamma_dist_cdf(Alpha,Beta,Mu,X)
log_gamma_dist_mean(Alpha,Beta,Mu)
log_gamma_dist_n(Alpha,Beta,Mu,N)
log_gamma_dist_pdf(Alpha,Beta,Mu,X)
log_gamma_dist_quantile(Alpha,Beta,Mu,X)
log_gamma_dist_variance(Alpha,Beta,Mu)
logistic_dist()
logistic_dist_cdf(X)
logistic_dist_mean()
logistic_dist_n(N)
logistic_dist_pdf(X)
logistic_dist_quantilef(X)
logistic_dist_variance()
lognormal_dist_n(Mu,Sigma,N)
logseries_dist_mean(Theta)
logseries_dist(Theta)
logseries_dist_variance(Theta)
log_zero()
matching_dist_cdf(N,R)
matching_dist_mean(N)
matching_dist(N)
matching_dist_n(N,Num)
matching_dist_pdf(N,R)
matching_dist_quantile(N,Q)
matching_dist_variance(N)
meanf(Dist,Params)
meanf(DistParams)
multivariate_hypergeometric_dist1(RemainingTrials,RemainingBalls,NumBalls,Aux)
multivariate_hypergeometric_dist_cdf(N,NumBalls,Ps)
multivariate_hypergeometric_dist_mean(N,NumBalls)
multivariate_hypergeometric_dist_n(N,NumBalls,Num)
multivariate_hypergeometric_dist(N,NumBalls)
multivariate_hypergeometric_dist_pdf(N,NumBalls,Ps)
negative_binomial_dist_cdf(N,P,K)
negative_binomial_dist_mean(N,P)
negative_binomial_dist(M,P)
negative_binomial_dist_n(M,P,N)
negative_binomial_dist_pdf(N,P,K)
negative_binomial_dist_quantile(N,P,Q)
negative_hypergeometric_dist_cdf(W,WTot,BTot,X)
negative_hypergeometric_dist_mean(W,WTot,BTot)
negative_hypergeometric_dist_n(W,WTot,BTot,N)
negative_hypergeometric_dist_pdf(W,WTot,BTot,X)
negative_hypergeometric_dist_quantile(W,WTot,BTot,X)
negative_hypergeometric_dist(W,WTot,BTot)
normal01()
normal_dist_mean(Mu,Sigma)
normal_dist(Mean,Stdev)
normal_dist_n(Mean,Stdev,N)
normal_dist_variance(Mu,Sigma)
normalize_dist(V)
normal_std_cdf(Z)
num_records_dist_cdf(N,K)
num_records_dist_mean_h_n(N)
num_records_dist_mean(N)
num_records_dist(N)
num_records_dist_pdf(N,K)
num_records_dist_quantile(N,Q)
order_statistics_m_estimator_u_all(Xs)
order_statistics_m_estimator_u(Xs,I)
order_statistics_m_estimator_v(Xs)
order_statistics_with_replacement_discrete_leq_cdf(PDF,CDF,N,K,X)
order_statistics_with_replacement_discrete_leq_quantile(PDF,CDF,N,K,X)
order_statistics_with_replacement_discrete_lt_cdf(PDF,CDF,N,K,X)
pascal_dist_cdf(R,P,K)
pascal_dist_mean(N,P)
pascal_dist_n(N,P,Num)
pascal_dist(N,P)
pascal_dist_pdf(N,P,X)
pascal_dist_variance(N,P)
pbirthday(Classes,Coincident,N)
pdf_all(Dist,Params)
pdf_all(DistParams)
pdf_all(Dist,Params,FromQ,ToQ)
pdf_all(DistParams,FromQ,ToQ)
pdf(Dist,Params,X)
pdf(DistParams,X)
phi_marsaglia(Z)
poisson_dist(Lambda)
poisson_dist_mean(Lambda)
poisson_dist_validate(Lambda)
poisson_dist_variance(Lambda)
poisson_gt0_dist(Lambda)
poisson_gt0_dist_n(Lambda,N)
poisson_process_dist_cdf(Mu,T,X)
poisson_process_dist_mean(Mu,T)
poisson_process_dist(Mu,T)
poisson_process_dist_n(Mu,T,N)
poisson_process_dist_pdf(Mu,T,X)
poisson_process_dist_quantile(Mu,T,Q)
poisson_process_dist_variance(Mu,T)
polya_eggenberg_dist_n(N,W,B,C,Num)
polya_eggenberg_dist(N,W,B,C)
probability(Formula,Dist)
probability_of_run_size(N,P,R)
prob_n_heads_after_k_in_max_m_tosses_dist_cdf(P,M,N,K)
prob_n_heads_after_k_in_max_m_tosses_dist_pdf(P,M,N,K)
prob_n_heads_after_k_in_max_m_tosses_dist(P,M,N)
prob_n_heads_after_k_in_max_m_tosses_dist_quantile(P,M,N,Q)
prob_n_heads_after_k_in_max_m_tosses_list(P,M,N)
qbirthday(Classes,Coincident,Prob)
quantile_all(Dist,Params)
quantile_all(DistParams)
quantile_all(Dist,Params,Qs)
quantile(Dist,Params,X)
quantile(DistParams,X)
quantile_qs()
rademacher_dist()
rademacher_dist_cdf(K)
rademacher_dist_n(N)
rademacher_dist_pdf(K)
rademacher_dist_quantile(K,Q)
random_integer1_cdf(N,X)
random_integer1_dist_cdf(N,X)
random_integer1_dist_mean(N)
random_integer1_dist(N)
random_integer1_dist_pdf(N,X)
random_integer1_dist_quantile(N,X)
random_integer1_dist_variance(N,X)
random_integer1_mean(N)
random_integer1(N)
random_integer1_pdf(N,X)
random_integer1_quantile(N,X)
random_integer1_variance(N,X)
random_integer_cdf(N,X)
random_integer_dist_cdf(N,X)
random_integer_dist_mean(N,X)
random_integer_dist(N)
random_integer_dist_pdf(N,X)
random_integer_dist_quantile(N,X)
random_integer_dist_variance(N,X)
random_integer_mean(N,X)
random_integer(N)
random_integer_pdf(N,X)
random_integer_quantile(N,X)
random_integer_variance(N,X)
random_walk_process_dist_cdf(P,T,X)
random_walk_process_dist_mean(P,T)
random_walk_process_dist_n(P,T,N)
random_walk_process_dist_pdf(P,T,X)
random_walk_process_dist(P,T)
random_walk_process_dist_quantile(P,T,Q)
random_walk_process_dist_variance(P,T)
rbirthday(Classes,Coincident)
replace_nth([H|Rest],N,Val)
replace_nth([_|Rest],1,Val)
row_vec_times_matrix_acc(V,P,I,W0)
row_vec_times_matrix(V,P)
shifted_exponential_dist(Lambda,T)
shifted_exponential_dist_n(Lambda,T,N)
shifted_geometric_dist_mean(P)
shifted_geometric_dist_n(P,N)
shifted_geometric_dist(P)
shifted_geometric_dist_variance(P)
stationary_dist1(P)
stationary_dist(P)
stationary_dist(P,Tolerance,MaxIter)
stationary_power_iteration(P,Pi,Tolerance,MaxIter)
stationary_power_iteration(P,Pi,Tolerance,MaxIter,Iter)
stdevf(Dist,Params)
stdevf(DistParams)
stdnormal_pdf(Z)
std_normal_rand()
stirling1(0,0)
stirling1(0,K)
stirling1_approx(N,K)
stirling1(N,0)
stirling1(N,K)
student_t_dist_mean(Nu)
student_t_dist_n(Nu,N)
student_t_dist(Nu)
student_t_dist_variance(Nu)
sum_prob_dist(A,B,N)
sum_prob_dist_cdf(A,B,N,S)
sum_prob_dist_mean(A,B,N)
sum_prob_dist_n(A,B,N,Num)
sum_prob_dist_pdf(A,B,N,S)
sum_prob_dist_quantile(A,B,N,Q)
sum_prob_dist_variance(A,B,N)
surprise(Dist)
surprise(Dist,Level)
survival_function(Dist,Params)
test_probability
triangular2_dist_cdf(Min,Max,C,X)
triangular2_dist_mean(Min,Max,C)
triangular2_dist(Min,Max,C)
triangular2_dist_n(Min,Max,C,N)
triangular2_dist_pdf(Min,Max,C,X)
triangular2_dist_quantile(Min,Max,C,X)
triangular2_dist_variance(Min,Max,C)
triangular_dist()
triangular_dist_cdf(Min,Max,X)
triangular_dist_cdf(X)
triangular_dist_mean()
triangular_dist_mean(Min,Max)
triangular_dist(Min,Max)
triangular_dist_n(Min,Max,N)
triangular_dist_n(N)
triangular_dist_pdf(Min,Max,X)
triangular_dist_pdf(X)
triangular_dist_quantile(Min,Max,X)
triangular_dist_quantile(X)
triangular_dist_variance()
triangular_dist_variance(Min,Max)
u01()
uniform_dist(Low,Up)
uniform_dist_n(Low,Up,N)
uniform_draw(List)
uniform_draw_mean(L)
uniform_draw_n(List,N)
uniform_draw_variance(L)
uniform(Low,Up)
uniform_n(Low,Up,N)
variancef(Dist,Params)
variancef(DistParams)
wiener_process_dist_cdf(Mu,Sigma,T,X)
wiener_process_dist_cdf(T,X)
wiener_process_dist_mean(Mu,_Sigma,T)
wiener_process_dist_mean(T)
wiener_process_dist(Mu,Sigma,T)
wiener_process_dist_n(Mu,Sigma,T,N)
wiener_process_dist_n(T,N)
wiener_process_dist_pdf(Mu,Sigma,T,X)
wiener_process_dist_pdf(T,X)
wiener_process_dist_quantile(Mu,Sigma,T,Q)
wiener_process_dist_quantile(T,X)
wiener_process_dist(T)
wiener_process_dist_variance(_Mu,Sigma,T)
wiener_process_dist_variance(T)
zeta(S)
zipf_dist_mean(S)
zipf_dist_n(N,S,Num)
zipf_dist_n(S,N)
zipf_dist(S)
zipf_dist_variance(S)
zipf_mma_dist_mean(Rho)
zipf_mma_dist_n(S,N)
zipf_mma_dist(Rho)
zipf_mma_dist_variance(Rho)


  Full disclosure: ChatGPT was used - in small or large part - for some of 
  these functions, especially the one requires advanced numerical methods.

  That being said, all PPL models and most of the functions in ppl_utils.pi
  was written without LLM.

  TODO: White a small helper text for each distribution, e.g.
  Picat> help_dist(binomial)  % or help_dist(binomial_dist)
  binomial(N,P): Given N trials each with probability P, what is the probability of success?

  This Picat model was created by Hakan Kjellerstrand, hakank@gmail.com
  See also my Picat page: http://www.hakank.org/picat/

*/

module ppl_distributions.
import ppl_common_utils.

%
% General utilities
%
% infinity() = 1.0e100.
% negative_ifinity() = -1.0e100.

all_nonneg([]).
all_nonneg([X|Xs]) :- X >= 0.0, all_nonneg(Xs).

% A safe sentinel for "log(0)"
log_zero() = -1.0e300.

euler_gamma = 0.57721566490153286061.

table
factorial_mem(N) = factorial(N).


% Exact version: Can be very slow for large N and K
binomialf(N,K) = factorial(N) // factorial(K) // factorial(N-K).

% 
binomialg(N,K) = factorialf(N) // factorialf(K) // factorialf(N-K).



% Float version of binomial
binomialf_float(N,K) = R =>
    if K < 0 ; K > N then
        R = 0.0
    else
        R = exp(lgamma(N+1.0) - lgamma(K+1.0) - lgamma(N-K+1.0))
    end.



/*
------------------------------------------------------------
stirling1(N, K)
------------------------------------------------------------
Unsigned Stirling numbers of the first kind (|s(N,K)|).

Meaning:
  Number of permutations of N elements with exactly K cycles.

Recurrence:
  s(0,0) = 1
  s(n,0) = 0 for n>0
  s(0,k) = 0 for k>0
  s(n,k) = s(n-1,k-1) + (n-1)*s(n-1,k)

This version:
  * Uses tabling for fast reuse.
  * Computes exact integers (no floats).
  * Works efficiently up to large N (~300+).

Example:
  stirling1(5,2) = 50
  stirling1(10,3) = 9330
------------------------------------------------------------
*/

table
stirling1(0,0) = 1.
stirling1(N,0) = 0 => N > 0.
stirling1(0,K) = 0 => K > 0.
stirling1(N,K) = S =>
    S = stirling1(N-1,K-1) + (N-1) * stirling1(N-1,K).

stirling1_approx(N,K) = round(exp(lgamma(N+1) - lgamma(K+1) - (N-K))).



% log gamma function using the Lanczos approximation
lgamma(X) = R =>
    if X < 0.5 then
        % Reflection formula for better accuracy on small X
        R = log(pi) - log(sin(pi*X)) - lgamma(1.0 - X)
    else
        % Lanczos coefficients for g=7, n=9
        G = 7.0,
        Coeffs = [
            0.99999999999980993,
            676.5203681218851,
           -1259.1392167224028,
            771.32342877765313,
           -176.61502916214059,
             12.507343278686905,
             -0.13857109526572012,
              9.9843695780195716e-6,
              1.5056327351493116e-7
        ],
        Y = X - 1.0,
        A = Coeffs[1],
        N = length(Coeffs),
        foreach(I in 2..N)
            A := A + Coeffs[I] / (Y + I - 1)
        end,
        T = Y + G + 0.5,
        R = 0.5*log(2*pi) + (Y+0.5)*log(T) - T + log(A)
    end.


%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%
% gamma_func(X)
%
% Lanczos approximation of the Gamma function.
% Accurate to about 15 digits for X > 0.
%
% For negative non-integer X, uses reflection formula:
%   Γ(x) = π / (sin(πx) * Γ(1 - x))
%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%
gamma_func(X) = G =>
    if X <= 0.0, (X == floor(X)) then
        throw($error('gamma_func: undefined for nonpositive integer'))
    else
        if X < 0.5 then
            % Reflection formula for better stability
            G = math.pi / (sin(math.pi * X) * gamma_func(1.0 - X))
        else
            % Lanczos coefficients (g=7, n=9)
            Coeffs = [
                0.99999999999980993,
                676.5203681218851,
                -1259.1392167224028,
                771.32342877765313,
                -176.61502916214059,
                12.507343278686905,
                -0.13857109526572012,
                9.9843695780195716e-6,
                1.5056327351493116e-7
            ],
            Y = X - 1.0,
            A = Coeffs[1],
            foreach(I in 2..length(Coeffs))
                A := A + Coeffs[I] / (Y + I - 1.0)
            end,
            T = Y + 7.5,   % g + 0.5
            G = sqrt(2.0 * math.pi) * safe_pow(T, Y + 0.5) * exp(-T) * A
        end
    end.

factorialf(N) = gamma_func(N+1).

% log Beta via lgamma
log_beta(A,B) = lgamma(A) + lgamma(B) - lgamma(A+B).

% beta function
beta_func(A,B) = exp(log_beta(A,B)).

%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%
% Regularized incomplete beta I_x(a,b)
%   I_x(a,b) = B_x(a,b) / B(a,b)
% Uses symmetry and Lentz's continued fraction (betacf).
%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%
incbeta_reg(A,B,X) = R =>
    % Parameter checks
    if A =< 0.0 ; B =< 0.0 ; X < 0.0 ; X > 1.0 then
        throw($error('incbeta_reg: A =< 0.0 ; B =< 0.0 ; X < 0.0 ; X > 1.0'))
    elseif X == 0.0 then
        R = 0.0
    elseif X == 1.0 then
        R = 1.0
    else
        % Use symmetry to improve convergence
        if X < (A+1.0)/(A+B+2.0) then
            R = exp( A*log(X) + B*log(1.0-X) - log_beta(A,B) ) * betacf(A,B,X) / A
        else
            R = 1.0 - exp( B*log(1.0-X) + A*log(X) - log_beta(B,A) ) * betacf(B,A,1.0-X) / B
        end
    end.

% Continued fraction for incomplete beta (Numerical Recipes style)
betacf(A,B,X) = H =>
    MaxIter = 200, Eps = 1.0e-12, FPMIN = 1.0e-300,
    M2 = 0.0,
    AA = 0.0,
    C = 1.0,
    D = 1.0 - (A+B)*X/(A+1.0),
    if abs(D) < FPMIN then D := FPMIN end,
    D := 1.0 / D,
    H = D,
    M = 1,
    OK = true,
    while (M =< MaxIter,OK == true)
        M2 := 2.0 * M,

        % First term (even step)
        AA := M*(B - M)*X / ((A + M2 - 1.0)*(A + M2)),
        D := 1.0 + AA*D, if abs(D) < FPMIN then D := FPMIN end, D := 1.0/D,
        C := 1.0 + AA/C, if abs(C) < FPMIN then C := FPMIN end,
        H := H * D * C,

        % Second term (odd step)
        AA := -(A + M)*(A + B + M)*X / ((A + M2)*(A + M2 + 1.0)),
        D := 1.0 + AA*D, if abs(D) < FPMIN then D := FPMIN end, D := 1.0/D,
        C := 1.0 + AA/C, if abs(C) < FPMIN then C := FPMIN end,
        Delta = D*C,
        H := H * Delta,

        if abs(Delta - 1.0) =< Eps then
           OK := false
        end,
        M := M + 1
    end.

%
% Regularized incomplete beta: I_x(A,B)
%   A,B > 0, 0 <= X <= 1
%
reg_incomplete_beta(A, B, X) = R =>
    if A =< 0.0 then
        throw($error('reg_incomplete_beta: A must be > 0'))
    elseif B =< 0.0 then
        throw($error('reg_incomplete_beta: B must be > 0'))
    elseif X < 0.0 ; X > 1.0 then
        throw($error('reg_incomplete_beta: X must be in [0,1]'))
    elseif X =:= 0.0 then
        R = 0.0
    elseif X =:= 1.0 then
        R = 1.0
    else
        % bt = exp( lgamma(A+B) - lgamma(A) - lgamma(B) + A*log(X) + B*log(1-X) )
        Bt = exp(lgamma(A + B) - lgamma(A) - lgamma(B) + A*log(X) + B*log(1.0 - X)),
        Threshold = (A + 1.0) / (A + B + 2.0),
        if X < Threshold then
            Cf = ibetacf(A, B, X),
            R = Bt * Cf / A
        else
            % symmetry: I_x(a,b) = 1 - I_{1-x}(b,a)
            Cf = ibetacf(B, A, 1.0 - X),
            R = 1.0 - Bt * Cf / B
        end
    end.

%
% Continued fraction for incomplete beta (modified Lentz method)
%
ibetacf(A, B, X) = H =>
    Tiny = 1.0e-300,
    Eps = 1.0e-15,
    MaxIter = 200,

    Qab = A + B,
    Qap = A + 1.0,
    Qam = A - 1.0,

    C = 1.0,
    D = 1.0 - Qab * X / Qap,
    if abs(D) < Tiny then D := Tiny end,
    D := 1.0 / D,
    H0 = D,

    M = 1,
    Hcur = H0,
    Done = false,
    while (M =< MaxIter, not Done) do
        M2 = 2 * M,

        % First step (aa1)
        Aa1 = M * (B - M) * X / ((Qam + M2) * (A + M2)),
        D := 1.0 + Aa1 * D,
        if abs(D) < Tiny then D := Tiny end,
        C := 1.0 + Aa1 / C,
        if abs(C) < Tiny then C := Tiny end,
        D := 1.0 / D,
        Delta = C * D,
        Hcur := Hcur * Delta,

        % Second step (aa2)
        Aa2 = -(A + M) * (Qab + M) * X / ((A + M2) * (Qap + M2)),
        D := 1.0 + Aa2 * D,
        if abs(D) < Tiny then D := Tiny end,
        C := 1.0 + Aa2 / C,
        if abs(C) < Tiny then C := Tiny end,
        D := 1.0 / D,
        Delta := C * D,
        Hcur := Hcur * Delta,

        if abs(Delta - 1.0) =< Eps then
            Done := true
        else
            M := M + 1
        end
    end,
    H = Hcur.

%
% Inverse regularized incomplete beta:
%   Finds x in (0,1) such that reg_incomplete_beta(A,B,x) = P
%
inverse_reg_incomplete_beta(A, B, P) = X =>
    if P =< 0.0 then
        X = 0.0
    elseif P >= 1.0 then
        X = 1.0
    else
        % initial guess from mean
        X0 = A / (A + B),
        X1 = X0,
        Eps = 1.0e-10,
        MaxIter = 100,
        I = 1,
        Done = false,
        while (I =< MaxIter, not Done) do
            F = reg_incomplete_beta(A, B, X1) - P,
            % derivative: dI_x/dx = x^{A-1} (1-x)^{B-1} / B(A,B)
            Df = exp((A-1.0)*log(X1) + (B-1.0)*log(1.0 - X1)
                     - (lgamma(A) + lgamma(B) - lgamma(A + B))),
            if Df =:= 0.0 then
                Done := true
            else
                Step = F / Df,
                X2 = X1 - Step,
                if X2 < 0.0 then X2 := 0.5 * X1 end,
                if X2 > 1.0 then X2 := 0.5 * (1.0 + X1) end,
                if abs(X2 - X1) < Eps then
                    Done := true
                else
                    X1 := X2,
                    I := I + 1
                end
            end
        end,
        X = X1
    end.

% Standard normal CDF Φ(Z) with ≈ machine-precision absolute accuracy.
phi_marsaglia(Z) = P =>
    % Handle extreme tails explicitly (prevents tiny negative / >1 due to rounding)
    if Z < -8.0 then
        P = 0.0
    elseif Z > 8.0 then
        P = 1.0
    else
        Q = Z*Z,
        % Series accumulation (stops automatically when adding no longer changes S)
        S = Z, T = 0.0, B = Z, I = 1.0,
        while (S != T)
            T := S,
            I := I + 2.0,
            B := B * Q / I,     % next odd-term factor
            S := T + B
        end,
        % 0.5 + s * exp(-z^2/2) / sqrt(2π)
        % where log(sqrt(2π)) ≈ 0.91893853320467274178
        P = 0.5 + S * exp(-0.5*Q - 0.9189385332046727)
    end.
    
%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%
% erf_approx(X) — approximation of the error function
% Abramowitz & Stegun 7.1.26, max error ~1.5e-7
%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%
erf_approx(X) = E =>
    T = 1.0 / (1.0 + 0.3275911 * abs(X)),
    % Coefficients
    A1 = 0.254829592,
    A2 = -0.284496736,
    A3 = 1.421413741,
    A4 = -1.453152027,
    A5 = 1.061405429,
    Poly = ((((A5*T + A4)*T + A3)*T + A2)*T + A1)*T,
    E = cond(X >= 0.0,
             1.0 - Poly * exp(-X*X),
            -1.0 + Poly * exp(-X*X)).


%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%
% invnorm_acklam(P)
% --------------------------------------------------------------
% Standard-normal quantile z such that Φ(z) = P
% Based on Peter John Acklam, 2003.
% Corrected tail signs:
%   Lower tail  ->  z = -Num/Den
%   Upper tail  ->  z = +Num/Den
%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%
invnorm_acklam(P) = Z =>
    % Coefficients for central region
    A0 = -3.969683028665376e+01, A1 = 2.209460984245205e+02,
    A2 = -2.759285104469687e+02, A3 = 1.383577518672690e+02,
    A4 = -3.066479806614716e+01, A5 = 2.506628277459239e+00,
    B0 = -5.447609879822406e+01, B1 = 1.615858368580409e+02,
    B2 = -1.556989798598866e+02, B3 = 6.680131188771972e+01,
    B4 = -1.328068155288572e+01,

    % Coefficients for tail regions
    C0 = -7.784894002430293e-03, C1 = -3.223964580411365e-01,
    C2 = -2.400758277161838e+00, C3 = -2.549732539343734e+00,
    C4 =  4.374664141464968e+00, C5 =  2.938163982698783e+00,
    D0 =  7.784695709041462e-03, D1 =  3.224671290700398e-01,
    D2 =  2.445134137142996e+00, D3 =  3.754408661907416e+00,

    % Breakpoints for regions
    PLow  = 0.02425,
    PHigh = 1.0 - PLow,  % = 0.97575

    if P < PLow then
        % ---------- LOWER TAIL ----------
        Q = sqrt(-2.0 * log(P)),
        Num = (((((C0*Q + C1)*Q + C2)*Q + C3)*Q + C4)*Q + C5),
        Den = ((((D0*Q + D1)*Q + D2)*Q + D3)*Q + 1.0),
        Z = - Num / Den          % negative z in lower tail

    elseif P =< PHigh then
        % ---------- CENTRAL REGION ----------
        Q = P - 0.5,
        R = Q * Q,
        Num = (((((A0*R + A1)*R + A2)*R + A3)*R + A4)*R + A5) * Q,
        Den = ((((B0*R + B1)*R + B2)*R + B3)*R + B4) * R + 1.0,
        Z = Num / Den

    else
        % ---------- UPPER TAIL ----------
        Q = sqrt(-2.0 * log(1.0 - P)),
        Num = (((((C0*Q + C1)*Q + C2)*Q + C3)*Q + C4)*Q + C5),
        Den = ((((D0*Q + D1)*Q + D2)*Q + D3)*Q + 1.0),
        Z =  Num / Den           % positive z in upper tail
    end.


%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%
% invnorm_acklam_symmetric(P)
% --------------------------------------------------------------
% Standard-normal quantile Φ^{-1}(P) using Acklam coefficients,
% implemented with symmetry to avoid tail sign mistakes.
%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%
invnorm_acklam_symmetric(P) = Z =>
    % Central region coefficients
    A0 = -3.969683028665376e+01, A1 = 2.209460984245205e+02,
    A2 = -2.759285104469687e+02, A3 = 1.383577518672690e+02,
    A4 = -3.066479806614716e+01, A5 = 2.506628277459239e+00,
    B0 = -5.447609879822406e+01, B1 = 1.615858368580409e+02,
    B2 = -1.556989798598866e+02, B3 = 6.680131188771972e+01,
    B4 = -1.328068155288572e+01,

    % Tail region coefficients
    C0 = -7.784894002430293e-03, C1 = -3.223964580411365e-01,
    C2 = -2.400758277161838e+00, C3 = -2.549732539343734e+00,
    C4 =  4.374664141464968e+00, C5 =  2.938163982698783e+00,
    D0 =  7.784695709041462e-03, D1 =  3.224671290700398e-01,
    D2 =  2.445134137142996e+00, D3 =  3.754408661907416e+00,

    PLow  = 0.02425,
    PMid  = 0.5,

    % Symmetry: reflect top half to bottom half and remember the sign
    Sign = cond(P < PMid, -1.0, 1.0),
    R = cond(P < PMid, P, 1.0 - P),

    if R < PLow then
        % ---------- TAIL (always positive base; apply Sign at end) ----------
        Q = sqrt(-2.0 * log(R)),
        Num = (((((C0*Q + C1)*Q + C2)*Q + C3)*Q + C4)*Q + C5),
        Den = ((((D0*Q + D1)*Q + D2)*Q + D3)*Q + 1.0),
        Zpos = Num / Den,              % positive magnitude
        Z = Sign * Zpos
    else
        % ---------- CENTRAL ----------
        Q = (cond(Sign < 0.0, P, 1.0 - P)) - 0.5,  % equals P-0.5 in original side
        R2 = Q * Q,
        Num = (((((A0*R2 + A1)*R2 + A2)*R2 + A3)*R2 + A4)*R2 + A5) * Q,
        Den = ((((B0*R2 + B1)*R2 + B2)*R2 + B3)*R2 + B4) * R2 + 1.0,
        Z = Num / Den
    end.



%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%
% invnorm_acklam_enforced(P)
% Acklam quantile with explicit sign enforcement in tails
%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%
invnorm_acklam_enforced(P) = Z =>
    % Central region coefficients
    A0 = -3.969683028665376e+01, A1 = 2.209460984245205e+02,
    A2 = -2.759285104469687e+02, A3 = 1.383577518672690e+02,
    A4 = -3.066479806614716e+01, A5 = 2.506628277459239e+00,
    B0 = -5.447609879822406e+01, B1 = 1.615858368580409e+02,
    B2 = -1.556989798598866e+02, B3 = 6.680131188771972e+01,
    B4 = -1.328068155288572e+01,

    % Tail region coefficients
    C0 = -7.784894002430293e-03, C1 = -3.223964580411365e-01,
    C2 = -2.400758277161838e+00, C3 = -2.549732539343734e+00,
    C4 =  4.374664141464968e+00, C5 =  2.938163982698783e+00,
    D0 =  7.784695709041462e-03, D1 =  3.224671290700398e-01,
    D2 =  2.445134137142996e+00, D3 =  3.754408661907416e+00,

    PLow  = 0.02425,
    PHigh = 1.0 - PLow,  % 0.97575

    if P < PLow then
        % ---------- LOWER TAIL ----------
        Q = sqrt(-2.0 * log(P)),
        Num = (((((C0*Q + C1)*Q + C2)*Q + C3)*Q + C4)*Q + C5),
        Den = ((((D0*Q + D1)*Q + D2)*Q + D3)*Q + 1.0),
        Z = -abs(Num / Den)    % force negative
    elseif P =< PHigh then
        % ---------- CENTRAL ----------
        Q = P - 0.5,
        R = Q*Q,
        Num = (((((A0*R + A1)*R + A2)*R + A3)*R + A4)*R + A5) * Q,
        Den = ((((B0*R + B1)*R + B2)*R + B3)*R + B4) * R + 1.0,
        Z = Num / Den          % sign comes from Q
    else
        % ---------- UPPER TAIL ----------
        Q = sqrt(-2.0 * log(1.0 - P)),
        Num = (((((C0*Q + C1)*Q + C2)*Q + C3)*Q + C4)*Q + C5),
        Den = ((((D0*Q + D1)*Q + D2)*Q + D3)*Q + 1.0),
        Z =  abs(Num / Den)    % force positive
    end.




%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%
% Incomplete gamma
% inc_gamma_lower_reg(A, X) = P(A,X)
%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%
inc_gamma_lower_reg(A, X) = P =>
    (A =< 0.0 ; X < 0.0) -> throw($error('inc_gamma_lower_reg: require A>0, X>=0'))
    ;
    if X == 0.0 then
        P = 0.0
    elseif X < A + 1.0 then
        % ---- Series for P(A,X) ----
        GLN = lgamma(A),
        AP = A,
        Sum = 1.0 / A,
        Del = Sum,
        Eps = 1.0e-14,
        while (abs(Del) > Eps)
            AP := AP + 1.0,
            Del := Del * (X / AP),
            Sum := Sum + Del
        end,
        Pref = exp(A*log(X) - X - GLN),
        P = Sum * Pref
    else
        % ---- Continued fraction for Q(A,X), then P = 1 - Q ----
        GLN = lgamma(A),
        Eps = 1.0e-14,
        FPMIN = 1.0e-300,
        B0 = X + 1.0 - A,
        C = 1.0 / FPMIN,
        D = 1.0 / max(B0, FPMIN),
        H = D,
        I = 1,
        Delta = 0.0,
        while (abs(Delta - 1.0) > Eps)
            AN = -I * (I - A),
            B = B0 + 2.0 * I,
            D := 1.0 / max(B + AN * D, FPMIN),
            C := max(B + AN / C, FPMIN),
            Delta := C * D,
            H := H * Delta,
            I := I + 1
        end,
        Q = exp(A*log(X) - X - GLN) * H,
        P = 1.0 - Q
    end.


/*
  General binary search for quantile estimation based on a given CDF
  Used for the quantiles which there's no available close form/method

  Example:
  Picat> X=binary_search_quantile($gamma_dist_cdf(180,8),0.999,1.0e-16,0,1.0e100)
  X = 1794.590717246031545

*/
binary_search_quantile(CDF,TargetP,Tolerance,Lower,Upper) = Ret =>
  Mid = (Lower + Upper) / 2,
  CDFMid = apply(CDF,Mid),
  Ret = _,
  if abs(CDFMid - TargetP) < Tolerance then
    % Stop if within tolerance of the target probability
    Ret = Mid
  elseif abs(Upper-Lower) < Tolerance then
    % Stop if lower and upper bounds are very close
    Ret = Mid
  elseif CDFMid < TargetP then
    % Continue searching by adjusting bounds
    Ret = binary_search_quantile(CDF,TargetP,Tolerance,Mid,Upper)
  else
   Ret = binary_search_quantile(CDF,TargetP,Tolerance,Lower,Mid)
  end.


/* safe_pow(Base,Exp)
   A log-safe power function that avoids underflow/overflow errors.
   Returns 0.0 for underflow, 1.0e308 for overflow,
   and throws explicit domain errors for invalid inputs.
*/
safe_pow(Base, Exp) = Res =>
    if Exp =:= 0.0 then
        Res = 1.0
    elseif Base =:= 0.0 then
        if Exp > 0.0 then
            Res = 0.0
        else
            throw($error('domain_error(number, pow(Base, Exp)), pow/2'))
        end
    elseif Base < 0.0, not integer(Exp) then
        throw($error('domain_error(number, pow(Base, Exp)), pow/2'))
    else
        LogVal = Exp * log(abs(Base)),
        % Double precision safe log range
        if LogVal > 709.7 then
            Res = 1.0e308
        elseif LogVal < -708.4 then
            Res = 0.0
        else
            Res0 = exp(LogVal),
            if Base < 0.0, integer(Exp), (Exp mod 2 =:= 1) then
                Res = -Res0
            else
                Res = Res0
            end
        end
    end.


/*
  zeta(S)
  --------
  High-precision, fast Riemann zeta for S > 1.
  Uses partial sum + Euler–Maclaurin remainder (2 terms).
  Typically ~1e-14 relative error for S≥1.5.
*/
% Fast with just 1.0e-10 precision
zeta(S) = Z =>
    if S =< 1.0 then
        throw($error('zeta: S must be > 1'))
    else
        % choose cutoff dynamically
        Nmax = cond(S < 2.0, 200_000, 50_000),

        Sum = sum([1.0/(I**S) : I in 1..Nmax]),

        % Euler–Maclaurin remainder terms
        B2 = 1.0/6.0,
        B4 = -1.0/30.0,
        Tail = (Nmax**(1.0 - S))/(S - 1.0)
              + 0.5*(Nmax**(-S))
              + (B2*S)*(Nmax**(-S-1.0))/2.0
              + (B4*S*(S+1.0)*(S+2.0)*(S+3.0))*(Nmax**(-S-3.0))/24.0,

        Z = Sum + Tail
    end.



%
% Regularized lower incomplete gamma:
%   gamma_regularized(A, X) = P(A,X) = gamma(A,X) / Gamma(A)
% Domain: A > 0, X >= 0
%
gamma_regularized(A, X) = P =>
    if A =< 0.0 then
        throw($error('gamma_regularized: A must be > 0'))
    elseif X < 0.0 then
        throw($error('gamma_regularized: X must be >= 0'))
    elseif X =:= 0.0 then
        P = 0.0
    elseif X < A + 1.0 then
        P = gamma_reg_series(A, X)
    else
        Q = gamma_reg_cf_q(A, X),
        P = 1.0 - Q
    end.

%
% Series expansion for P(A,X) when X < A+1
%
gamma_reg_series(A, X) = P =>
    Eps = 1.0e-15,
    MaxIter = 100000,
    LogG = lgamma(A),
    Pref = exp(A*log(X) - X - LogG),
    Sum0 = 1.0 / A,
    Term0 = Sum0,
    N0 = 1,
    Sum = gamma_reg_series_loop(A, X, Sum0, Term0, N0, Eps, MaxIter),
    P = Pref * Sum.

gamma_reg_series_loop(A, X, Sum, Term, N, Eps, MaxIter) = Res =>
    if N > MaxIter then
        Res = Sum
    else
        Term1 = Term * X / (A + N),
        Sum1  = Sum + Term1,
        if abs(Term1) =< Eps * abs(Sum1) then
            Res = Sum1
        else
            Res = gamma_reg_series_loop(A, X, Sum1, Term1, N+1, Eps, MaxIter)
        end
    end.

%
% Continued fraction (Lentz) for Q(A,X) when X >= A+1
%
gamma_reg_cf_q(A, X) = Q =>
    Tiny = 1.0e-300,
    Eps = 1.0e-15,
    MaxIter = 100000,
    LogG = lgamma(A),
    Pref = exp(A*log(X) - X - LogG),

    B0 = X + 1.0 - A,
    B0a = cond(abs(B0) < Tiny, Tiny, B0),
    C0 = 1.0 / Tiny,
    D0 = 1.0 / B0a,
    H0 = D0,
    I0 = 1,
    H = gamma_reg_cf_q_loop(A, X, B0a, C0, D0, H0, I0, Tiny, Eps, MaxIter),
    Q = Pref * H.

gamma_reg_cf_q_loop(A, X, B0, C, D, H, I, Tiny, Eps, MaxIter) = Res =>
    if I > MaxIter then
        Res = H
    else
        Ai = I * (A - I),
        Bi = B0 + 2.0 * I,

        Den0 = Bi + Ai * D,
        Den = cond(abs(Den0) < Tiny, Tiny, Den0),
        D1 = 1.0 / Den,

        Num0 = Ai / C,
        C1t = Bi + Num0,
        C1 = cond(abs(C1t) < Tiny, Tiny, C1t),

        Delta = C1 * D1,
        H1 = H * Delta,

        if abs(Delta - 1.0) =< Eps then
            Res = H1
        else
            Res = gamma_reg_cf_q_loop(A, X, B0, C1, D1, H1, I+1, Tiny, Eps, MaxIter)
        end
    end.



%
% Inverse regularized lower incomplete gamma function
% inverse_gamma_regularized(A, P)
%   Finds X such that gamma_regularized(A, X) = P
%
inverse_gamma_regularized(A, P) = X =>
    if A =< 0.0 then
        throw($error('inverse_gamma_regularized: A must be > 0'))
    elseif P =< 0.0 then
        X = 0.0
    elseif P >= 1.0 then
        X = 1.0e100
        % throw($error('inverse_gamma_regularized: P must be > 0'))
    else
        X0 = inverse_gamma_regularized_initial(A, P),
        X = inverse_gamma_regularized_iter(A, P, X0)
    end.

%
% Initial approximation for X
%
inverse_gamma_regularized_initial(A, P) = X0 =>
    if A > 1.0 then
        % Wilson–Hilferty approximation for chi-square relation
        T = sqrt(-2.0 * log(cond(P < 0.5, P, 1.0 - P))),
        S = sign(P - 0.5),
        Approx = A * (1.0 - 1.0/(9.0*A) - S * T/(3.0*sqrt(A)))**3,
        X0 = cond(Approx > 0.0, Approx, 1.0e-8)
    else
        % Small A: power approximation
        G = gamma_func(A + 1.0),
        X0 = cond(P < 0.5, (P * G)**(1.0/A), (-log(1.0 - P))**A)
    end.

%
% Newton–Raphson iteration
%
inverse_gamma_regularized_iter(A, P, X0) = X =>
    Eps = 1.0e-12,
    MaxIter = 100,
    Xcur = X0,
    N = 1,
    Done = false,
    while (N =< MaxIter, not Done) do
        Gp = gamma_regularized(A, Xcur) - P,
        if abs(Gp) < Eps then
            Done := true
        else
            % derivative: f'(X) = exp(-X) * X^(A-1) / Gamma(A)
            % Fp = exp(-Xcur) * Xcur**(A - 1.0) / gamma_func(A),
            Fp = exp(-Xcur) * safe_pow(Xcur,(A - 1.0)) / gamma_func(A),
            if Fp =:= 0.0 then
                Done := true
            else
                Xnew = Xcur - Gp / Fp,
                if Xnew =< 0.0 then
                    Xnew := Xcur / 2.0
                end,
                Xcur := Xnew,
                N := N + 1
            end
        end
    end,
    X = Xcur.


/*
  Harmonic_number
*/
harmonic_number(N) = harmonic_number(N,1).
harmonic_number(N, A) = H =>
    (N < 1 -> throw($error('harmonic_number: N must be >= 1')) ; true),
    H = sum([K**(-A) : K in 1..N]).





% --- robust U(0,1) from integers: uses rand/2, never returns 0 or 1 ---
u01() = U =>
    % Draw integer I uniformly from 0..2147483646 (2^31-1 - 1)
    I = random(0, 2147483646),
    % Map to (0,1): add 0.5 and divide by 2^31-1
    U = (I + 0.5) / 2147483647.0.

% Internal: standard normal via Box–Muller using u01()
std_normal_rand() = Z =>
    U1 = u01(),                % ∈ (0,1), safe for log
    U2 = u01(),
    R  = sqrt(-2.0 * log(U1)),
    A  = 2.0 * 3.141592653589793 * U2,
    Z  = R * cos(A).



/*
----------------------------------------------------------------------
digamma(X)
----------------------------------------------------------------------
Computes the digamma function psi(X) = d/dX ln(Gamma(X)).

This implementation:
  - uses argument reduction for X < 6
  - uses an asymptotic expansion for X >= 6
  - valid for X > 0
  - throws domain_error for invalid X

Asymptotic expansion:
  psi(x) ≈ ln(x) - 1/(2x)
           - 1/(12x^2)
           + 1/(120x^4)
           - 1/(252x^6)
           + 1/(240x^8)
           - 5/(660x^10)

This gives fast and accurate values for all positive X.

Accuracy:
  ~1e-12 relative error across most of the positive real line.
----------------------------------------------------------------------
*/

digammaf(X) = Val =>
    if X =< 0.0 then
        throw($error('domain_error(positive_real,X),digamma/1'))
    end,

    % Argument reduction: push X up until X >= 6
    Val = digamma_reduce(X, 0.0).


digamma_reduce(X, Acc) = Val =>
    if X < 6.0 then
        % psi(x) = psi(x+1) - 1/x
        Acc1 = Acc - 1.0 / X,
        X1 = X + 1.0,
        Val = digamma_reduce(X1, Acc1)
    else
        % Now apply asymptotic series for large X
        Val = Acc + digamma_asymptotic(X)
    end.


/*
----------------------------------------------------------------------
Asymptotic expansion valid for X >= 6
----------------------------------------------------------------------
*/
digamma_asymptotic(X) = A =>
    Inv = 1.0 / X,
    Inv2 = Inv * Inv,
    Inv4 = Inv2 * Inv2,
    Inv6 = Inv4 * Inv2,
    Inv8 = Inv4 * Inv4,
    Inv10 = Inv8 * Inv2,

    A = log(X)
        - 0.5 * Inv
        - (1.0 / 12.0) * Inv2
        + (1.0 / 120.0) * Inv4
        - (1.0 / 252.0) * Inv6
        + (1.0 / 240.0) * Inv8
        - (5.0 / 660.0) * Inv10.


% type(gaussian_dist,continous).
% type(exponential_dist,continous).
% type(exponential_dist2,continous).

% type(geometric_dist,discrete).
% type(bernoulli_dist,discrete).
% type(binomial_dist,discrete).
% type(negative_binomial_dist,discrete).
% type(poisson_dist,discrete).
% type(pascal_dist,discrete).




/*
  probability(formula,[distributions])

  probability($X >= 3,[X=binomial_dist(10,1/2)])

  Picat> X=$binomial_dist(10,1/2), F=[cond($(X.apply > 4),1,0) : I in 1..1000].sum/1000
  X = binomial_dist(10,1 / 2)
  F = 0.649

  TODO!

*/
/*
probability(Formula,Dist) = Res =>
  Formula =.. F,
  println(formula=Formula=F),
  Dist =.. D,
  println(dist=Dist=D),  
  Res = [R : I in 1..10, R=cond((Dist,Formula),1,0)].
  
test_probability =>
 T = probability($(X >= 3),$(X=binomial_dist(10,1/2))),
 println(T).
*/


/*
  Survival function: 1-CDF

  Picat> X=1-binomial_dist_cdf(10,1/2,0.9)
  X = 0.9990234375

  Picat> X=survival_function(binomial,[10,1/2,0.9])
  X = 0.9990234375


*/
survival_function(Dist,Params) = Survival =>
  bp.atom_concat(Dist,'_dist_cdf',CDF),
  C =.. [CDF,Params].flatten,
  Survival=1-C.apply.
  

/*
  Hazard function:
  - for continuous distributions: PDF/(1-CDF)
  - for discrete distributions: PDF/(1-CDF(k-1))
*/
hazard_function_cont(Dist,Params) = Hazard =>
  S = survival_function(Dist,Params),
  bp.atom_concat(Dist,'_dist_pdf',PDF),
  P =.. [PDF,Params].flatten,
  Hazard = P.apply/S.

% K (last parameter in Params) should be K-1 in the PDF.
% TODO: Check that the discrete distributions handle K=0 correctly!
hazard_function_discrete(Dist,Params) = Hazard =>
  S = survival_function(Dist,Params),
  Len = Params.len,
  K2 = Params.last - 1,
  Params2 = Params[1..Len-1],  
  bp.atom_concat(Dist,'_dist_pdf',PDF),
  P =.. [PDF,Params2,K2].flatten,
  Hazard = P.apply/S.


/*
  surprise(Dist, Level)

  Returns the values of the distribution Dist for which there would
  be a surprise to see those values.

  Level is the lower and upper quantiles that constitutes the
  suprises, default 0.001 (and 1.0.001)
  
  Example: Flip a coin 100 times. How many heads (or tails) should we be
  surprised to see?

  At (default) surprise level 0.001
  Picat> X=surprise($binomial(100,1/2))     
  X = [35,65]

  Let's be more skeptical with a level of 0.00001:
  Picat> X=surprise($binomial(100,1/2),0.00001)
  X = [29,71]

  For (self-reported) IQ-scores:
  Picat> X=surprise($normal(100,15),0.00001)   
  X = [36.026638091193696,163.973361908743811]
 
*/
surprise(Dist) = surprise(Dist,0.001).
surprise(Dist,Level) = [Q1,Q2] =>
  Q1 = quantile(Dist,Level),
  Q2 = quantile(Dist,1-Level).


%
% flatten1(List)
%
% Flatten the list one level
% flatten1([ [1,2,3,4], [[5,6,7],[8,9,10]] ,[11]])=A
% A = [1,2,3,4,[5,6,7],[8,9,10],11]
%
flatten1(List) = fold(++,[], List).

/*
   pdf(Dist,Params,X) 
   pdf(DistParams,X) 

   returns the values of "Dist(Params,X)".
   
   The distribution can be either of the form:
   - <dist>_dist
   - <dist>

   - Params is a list of parameter(s) to the distribution
     E.g.   Dist = binomial, Params = [10,1/2]
     For the DistParam variant, the parameters are inside
     the distribution term, e.g binomial_dist(10,1/2)
   
   For example:

   Picat> pdf($binomial_dist(10,1/2),4)=X                          
   X = 0.205078125

   Picat> pdf($binomial(10,1/2),4)=X                          
   X = 0.205078125

   Picat> pdf($binomial,[10,1/2],4)=X 
   X = 0.205078125

   Picat> pdf($binomial_dist,[10,1/2],4)=X
   X = 0.205078125
   

   cdf(Dist,Params) 
   returns $<dist_cdf(Params)

   quantile(Dist,Params) 
   returns $<dist_quantile(Params)

*/

pdf(DistParams,X) = pdf(Dist,Params,X) =>
    DistParams =.. [Dist|Params].
pdf(Dist,Params,X) = Res =>
    bp.atom_chars(Dist,DistL),
    if once(append(_,[d,i,s,t],DistL)) then
      bp.atom_concat(Dist,'_pdf',PDF)
    else 
      bp.atom_concat(Dist,'_dist_pdf',PDF)
    end,
    F =.. [PDF,Params,X].flatten,
    Res = apply(F).

cdf(DistParams,X) = cdf(Dist,Params,X)  =>
    DistParams =.. [Dist|Params].
cdf(Dist,Params,X) = Res =>
    bp.atom_chars(Dist,DistL),
    if once(append(_,[d,i,s,t],DistL)) then
      bp.atom_concat(Dist,'_cdf',CDF)
    else 
      bp.atom_concat(Dist,'_dist_cdf',CDF)
    end,
    F =.. [CDF,Params,X].flatten,
    Res = apply(F).

quantile(DistParams,X) = quantile(Dist,Params,X)  =>
    DistParams =.. [Dist|Params].
quantile(Dist,Params,X) = Res =>
    bp.atom_chars(Dist,DistL),
    if once(append(_,[d,i,s,t],DistL)) then
      bp.atom_concat(Dist,'_quantile',Quantile)
    else 
      bp.atom_concat(Dist,'_dist_quantile',Quantile)
    end,
    F =.. [Quantile,Params,X].flatten,
    Res = apply(F).

meanf(DistParams) = meanf(Dist,Params) =>
  DistParams =.. [Dist|Params].
meanf(Dist,Params) = Res =>
    bp.atom_chars(Dist,DistL),
    if once(append(_,[d,i,s,t],DistL)) then
      bp.atom_concat(Dist,'_mean',MeanF)
    else 
      bp.atom_concat(Dist,'_dist_mean',MeanF)
    end,
    F =.. [MeanF,Params].flatten,
    Res = apply(F).
  
variancef(DistParams) = variancef(Dist,Params) =>
  DistParams =.. [Dist|Params].
variancef(Dist,Params) = Res =>
    bp.atom_chars(Dist,DistL),
    if once(append(_,[d,i,s,t],DistL)) then
      bp.atom_concat(Dist,'_variance',VarianceF)
    else 
      bp.atom_concat(Dist,'_dist_variance',VarianceF)
    end,
    F =.. [VarianceF,Params].flatten,
    Res = apply(F).

stdevf(DistParams) = stdevf(Dist,Params) =>
  DistParams =.. [Dist|Params].
stdevf(Dist,Params) = Res => 
    bp.atom_chars(Dist,DistL),
    if once(append(_,[d,i,s,t],DistL)) then
      bp.atom_concat(Dist,'_variance',VarianceF)
    else 
      bp.atom_concat(Dist,'_dist_variance',VarianceF)
    end,
    F =.. [VarianceF,Params].flatten,
    Res = sqrt(apply(F)).


/*
  pdf_all(Dist,Params,FromQ,ToQ)
  pdf_all(DistParam,FromQ,ToQ)
  pdf_all(Dist,Param)
  pdf_all(DistParam)

  Returns a list of value=PDF for the
  values in the range FromQ (default 0.01) to ToQ (default 0.99).

  - Dist can be either of the form
    $normal_dist([100,15]})
    $normal([100,15])

  - DistParam can be in either form
    $normal_dist(100,15)
    $normal(100,15)

  I.e. '_dist' is optional.
  Note the '$' escapes of the term. (Without '$', Picat tries to evaluate it as function.)

  Example: 
  Picat> pdf_all(binomial,[10,1/2],0.01,0.99).print_list             
  pdf = binomial_dist_pdf(10,0.5)
  1 = 0.00976562
  2 = 0.0439453
  3 = 0.1171875
  4 = 0.205078
  5 = 0.246094
  6 = 0.205078
  7 = 0.1171875
  8 = 0.0439453
  9 = 0.00976562

  Show according to probability:
  Picat> pdf_all(binomial,[10,1/2],0.01,0.99).sort_down(2).print_list
  pdf = binomial_dist_pdf(10,0.5)
  5 = 0.246094
  6 = 0.205078
  4 = 0.205078
  7 = 0.1171875
  3 = 0.1171875
  8 = 0.0439453
  2 = 0.0439453
  9 = 0.00976562
  1 = 0.00976562

  Picat> cdf_all(normal,[100,15],0.1,0.9).print_list                 
  cdf = normal_dist_cdf(100,15)
  80.7767 = 0.1
  81.7767 = 0.112205
  82.7767 = 0.125439
  83.7767 = 0.139726
  84.7767 = 0.15508
  85.7767 = 0.171509
  86.7767 = 0.18901
  ... 
  111.777 = 0.783807
  112.777 = 0.802832
  113.777 = 0.820808
  114.777 = 0.837716
  115.777 = 0.85355
  116.777 = 0.868312
  117.777 = 0.882014
  118.777 = 0.894675


  Heuristic: 
  If Dist is a continuous distribution and the the difference
  between Q1 and Q2 <= 5, then the precision is reduced to
  give about 15 values.
  
  cdf_all(Dist,Params)
  cdf_all(DistParams)
  Same as pdf_all/1-2, but for the CDF.

  Picat> cdf_all(binomial,[10,1/2],0.01,0.99).print_list
  [q1 = 1,q2 = 9,diff = 8]
  cdf = binomial_dist_cdf(10,0.5)
  1 = 0.0107422
  2 = 0.0546875
  3 = 0.171875
  4 = 0.376953
  5 = 0.623047
  6 = 0.828125
  7 = 0.9453125
  8 = 0.989258
  9 = 0.999023

  TODO: 
  Distributions with lists as parameter are not supported. For example, 
  this throws an error:
  dirichlet does not work:
  Picat> X=pdf_all($dirichlet,[180,8],0.1,0.9)
  *** Undefined procedure: dirichlet_dist_quantile/3

  multinomial does not work:
  Picat> X=pdf_all($multinomial_dist(10,[1,2,3]))
  *** Undefined procedure: multinomial_dist_quantile/5
  Multinomial does not have a quantile!


  Picat> X = pdf_all($categorical,[[1/4,3/4],[1,2]],0.001,0.999999999)
  Ah, but here categorical_dist has no categorical_dist_quantile, it has 
  quantile_index and quantile_value!
  

*/
% pdf_all(DistParams) = pdf_all(Dist,Params) =>
%   println($pdf_all0(DistParams)),
%   DistParams =.. [Dist|Params],
%   println([dist=Dist,params=Params]).
  
% pdf_all(DistParams,FromQ,ToQ) = pdf_all(Dist,Params,FromQ,ToQ) =>
%   println($pdf_all1(DistParams,FromQ,ToQ)),
%   DistParams =.. [Dist|Params],println([dist=Dist,params=Params]).
  
pdf_all(DistParams) = pdf_all(Dist,Params,0.01,0.99) =>
  DistParams =.. [Dist|Params].
pdf_all(Dist,Params) = pdf_all(Dist,Params,0.01,0.99).
pdf_all(DistParams,FromQ,ToQ) = pdf_all(Dist,Params,FromQ,ToQ) =>
  DistParams =.. [Dist|Params].
pdf_all(Dist,Params,FromQ,ToQ) = Res =>
  Q1 = quantile(Dist,Params,FromQ),
  Q2 = quantile(Dist,Params,ToQ),
  Diff = Q2-Q1,
  Prec = 1,
  if float(Q1), Diff <= 5 then
    Prec := Diff/15
  end,  
  % println([q1=Q1,q2=Q2,diff=Diff,prec=Prec]),  
  L = [],
  foreach(V in Q1..Prec..Q2)
    P = pdf(Dist,Params,V),
    L := L ++ [V=P]
  end,
  Res = L.  

cdf_all(DistParams) = cdf_all(Dist,Params,0.01,0.99) =>
  DistParams =.. [Dist|Params].
cdf_all(Dist,Params) = cdf_all(Dist,Params,0.01,0.99).
cdf_all(DistParams,FromQ,ToQ) = cdf_all(Dist,Params,FromQ,ToQ) =>
  DistParams =.. [Dist|Params].
cdf_all(Dist,Params,FromQ,ToQ) = Res =>
  Q1 = quantile(Dist,Params,FromQ),
  Q2 = quantile(Dist,Params,ToQ),
  Diff = Q2-Q1,
  Diff = Q2-Q1,
  Prec = 1,
  if float(Q1), Diff <= 5 then
    Prec := Diff/15
  end,
  % println([q1=Q1,q2=Q2,diff=Diff,prec=Prec]),  
  L = [],
  foreach(V in Q1..Prec..Q2)
    P = cdf(Dist,Params,V),
    L := L ++ [V=P]
  end,
  Res = L.

quantile_qs() =
  [0.000001,0.00001,0.001,0.01,0.025,0.05,0.1,0.25,0.50,0.75,0.84,0.9,0.975,0.99,0.999,0.99999,0.999999].

quantile_all(DistParams) = quantile_all(Dist,Params,quantile_qs()) =>
  DistParams =.. [Dist|Params].
quantile_all(Dist,Params) = quantile_all(Dist,Params,quantile_qs()).
quantile_all(Dist,Params,Qs) = L =>
  L = [],
  foreach(Q in Qs)
    V = quantile(Dist,Params,Q),
    L := L ++ [Q=V]
  end.

/*

  dist_sigma(Dist, Sigma)

  Given a distribution Dist (e.g. $normal_dist(Mean,Stdev))
  and a multiple Sigma, compute:

    Val = Mean + Sigma * Stdev
    CDF = 1.0 - cdf(Dist, Val)

  Returns a pair [Val, CDF]:
    Val : the numeric cutoff value
    CDF : the upper-tail probability of exceeding Val

  Requirements:
    * Dist must support mean/1, variance/1, and cdf/2
    * Sigma must be a number (can be float or integer)

  
  Returns the Sigma*standard devation and its 1-CDF for the distribution Dist,
  i.e. the probability of getting at least a value as large as this.

  Picat> [S=dist_sigma($normal_dist(100,15),S) : S in 0..10].printf_list   
   0 [100.0,0.5] 
   1 [115.0,0.158655253931457] 
   2 [130.0,0.022750131948179] 
   3 [145.0,0.00134989803163] 
   4 [160.0,0.000031671241833] 
   5 [175.0,0.000000286651572] 
   6 [190.0,0.000000000986588] 
   7 [205.0,0.000000000001281] 
   8 [220.0,0.000000000000002] 
   9 [235.0,0.0] 
  10 [250.0,0.0] 

  Cf data_sigma/2 (in ppl_utils)

  Picat> D=data_sigma(normal_dist_n(100,15,1000000),4)
  D = [160.05431345860336,0.000027]

  Picat> D=dist_sigma($normal_dist(100,15),4)  
  D = [160.0,0.000031671241833]


*/
dist_sigma(Dist,Sigma) = [Val,CDF] =>
  Mean = meanf(Dist),
  Val = Mean+Sigma*stdevf(Dist),
  CDF = 1-cdf(Dist,Val).

/*

  Uniform distribution (continuous)

  Models a continuous value X uniformly distributed on the interval [A,B].

  PDF(x) = 1/(B-A) for A <= x <= B, else 0
  CDF(x) = 0 for x < A
         (x - A)/(B - A) for A <= x < B
         1 for x >= B
  Quantile(u) = A + u*(B - A), with u in [0,1]
  Mean = (A + B)/2
  Variance = (B - A)^2 / 12

  Parameters:
  A - lower bound (float)
  B - upper bound (float), must satisfy A < B

  Functions:
  uniform_dist_pdf(A, B, X)
      Density at X.

  uniform_dist_cdf(A, B, X)
      Cumulative probability P(X' <= X).

  uniform_dist_quantile(A, B, U)
      Value x such that CDF(x) >= U (left-continuous). Handles U=0 -> A, U=1 -> B.

  uniform_dist(A, B)
      Random sample using u01(); returns a value in (A,B) (open interval).

  uniform_dist_mean(A, B)
      Expected value.

  uniform_dist_variance(A, B)
      Variance.

  Usage examples:
  println(uniform_dist_pdf(0.0, 1.0, 0.3)).
  println(uniform_dist_cdf(2.0, 5.0, 3.0)).
  println(uniform_dist_quantile(2.0, 5.0, 0.75)).
  X = uniform_dist(-1.0, 1.0), printf("Sample: %w\n", X).
  println(uniform_dist_mean(2.0, 5.0)).
  println(uniform_dist_variance(2.0, 5.0)).

*/
uniform(Low,Up) = frand(Low,Up). % Using the built in frand/2
uniform_n(Low,Up,N) = [uniform(Low,Up) : _ in 1..N].

uniform_dist(Low,Up) = uniform(Low,Up).
uniform_dist_n(Low,Up,N) = [uniform_dist(Low,Up) : _ in 1..N].

% ----- Random generator -----
% Uses u01() to avoid endpoints; returns a sample in (A,B)
% uniform_dist(A, B) = X =>
%     uniform_dist_validate(A, B),
%     U = u01(),
%     X = A + U * (B - A).


/* ========= Uniform distribution on (A,B) =========
   Support: A < X < B  (density defined on [A,B], RNG uses open interval via u01())

   PDF(x) = 1/(B-A) for A <= x <= B, else 0
   CDF(x) = 0                  for x < A
            (x - A)/(B - A)    for A <= x < B
            1                  for x >= B
   Quantile(u) = A + u*(B - A), with endpoints: u=0 -> A, u=1 -> B
   Mean = (A + B)/2
   Variance = (B - A)^2 / 12
*/

% ----- validation -----
uniform_dist_validate(A, B) =>
    if A >= B then
        throw('uniform_dist: require A < B')
    end.

% ----- PDF -----
uniform_dist_pdf(A, B, X) = P =>
    uniform_dist_validate(A, B),
    if (X < A) then
        P = 0.0
    elseif (X > B) then
        P = 0.0
    else
        P = 1.0 / (B - A)
    end.

% ----- CDF -----
uniform_dist_cdf(A, B, X) = F =>
    uniform_dist_validate(A, B),
    if (X < A) then
        F = 0.0
    elseif (X >= B) then
        F = 1.0
    else
        F = (X - A) / (B - A)
    end.

% ----- Quantile (left-continuous) -----
uniform_dist_quantile(A, B, U) = X =>
    uniform_dist_validate(A, B),
    if (U < 0.0 ; U > 1.0) then
        throw('uniform_dist_quantile: U must be in [0,1]')
    elseif U =:= 0.0 then
        X = A
    elseif U =:= 1.0 then
        X = B
    else
        X = A + U * (B - A)
    end.


% ----- Mean and Variance -----
uniform_dist_mean(A, B) = M =>
    uniform_dist_validate(A, B),
    M = 0.5 * (A + B).

uniform_dist_variance(A, B) = V =>
    uniform_dist_validate(A, B),
    D = (B - A),
    V = (D * D) / 12.0.


/*
  Discrete Uniform distribution

  Represents an integer-valued random variable uniformly distributed over
  the inclusive range [A, B].

  PMF(x) = 1 / (B - A + 1) for A <= x <= B, else 0
  CDF(x) = 0 for x < A
         (floor(x) - A + 1)/(B - A + 1) for A <= x < B
         1 for x >= B
  Quantile(u) = A + floor(u * (B - A + 1)), for u in [0,1]
  Mean = (A + B)/2
  Variance = ((B - A + 1)^2 - 1)/12

  Parameters:
  A - lower bound (integer)
  B - upper bound (integer), must satisfy A <= B

  Functions:
  discrete_uniform_dist_pdf(A, B, X)
      Probability mass for value X.

  discrete_uniform_dist_cdf(A, B, X)
      Cumulative probability P(X' <= X).

  discrete_uniform_dist_quantile(A, B, U)
      Smallest integer x such that CDF(x) >= U (left-continuous).

  discrete_uniform_dist(A, B)
      Random integer uniformly sampled from [A,B].

  discrete_uniform_dist_mean(A, B)
      Expected value.

  discrete_uniform_dist_variance(A, B)
      Variance.

  Usage examples:
  println(discrete_uniform_dist_pdf(1, 6, 3)).
  println(discrete_uniform_dist_cdf(1, 6, 3)).
  println(discrete_uniform_dist_quantile(1, 6, 0.5)).
  X = discrete_uniform_dist(1, 6), printf("Random draw: %w\n", X).
  println(discrete_uniform_dist_mean(1, 6)).
  println(discrete_uniform_dist_variance(1, 6)).

  Conceptual notes:
  - The discrete uniform distribution is the simplest discrete model,
    assigning equal probability to each integer in [A,B].
  - It is commonly used to model fair dice, random integers, or
    categorical selections with equal likelihood.
  - As (B - A + 1) grows, it approximates the continuous uniform
    distribution on [A,B].

*/

% uniform_discrete(From,To) = random(From,To). % Picat built-in

/* ========= Discrete Uniform distribution =========
   Support: X ∈ {A, A+1, ..., B}
   with A and B integers and A <= B

   PMF(x) = 1/(B - A + 1)  for A <= x <= B, else 0
   CDF(x) = 0                     for x < A
            (floor(x) - A + 1)/(B - A + 1)  for A <= x < B
            1                     for x >= B
   Quantile(u) = A + floor(u * (B - A + 1))
   Mean = (A + B)/2
   Variance = ((B - A + 1)^2 - 1) / 12
*/

% % ----- Random generator -----
% Uses u01() that avoids 0/1
discrete_uniform_dist(A, B) = X =>
    discrete_uniform_dist_validate(A, B),
    U = u01(),
    X = A + floor(U * (B - A + 1)),
    if X > B then X := B end.  % guard for rounding

discrete_uniform_dist_n(A, B, N) = [discrete_uniform_dist(A, B) : _ in 1..N].

% ----- validation -----
discrete_uniform_dist_validate(A, B) =>
    if not integer(A) then
        throw('discrete_uniform_dist: A must be an integer')
    elseif not integer(B) then
        throw('discrete_uniform_dist: B must be an integer')
    elseif A > B then
        throw('discrete_uniform_dist: require A <= B')
    end.

% ----- PDF (PMF) -----
discrete_uniform_dist_pdf(A, B, X) = P =>
    discrete_uniform_dist_validate(A, B),
    if not integer(X) then
        P = 0.0
    elseif (X < A ; X > B) then
        P = 0.0
    else
        P = 1.0 / (B - A + 1)
    end.

% ----- CDF -----
discrete_uniform_dist_cdf(A, B, X) = F =>
    discrete_uniform_dist_validate(A, B),
    if X < A then
        F = 0.0
    elseif X >= B then
        F = 1.0
    else
        % floor(X) handles non-integer inputs
        F = (floor(X) - A + 1) / (B - A + 1)
    end.

% ----- Quantile (left-continuous) -----
discrete_uniform_dist_quantile(A, B, U) = X =>
    discrete_uniform_dist_validate(A, B),
    if (U < 0.0 ; U > 1.0) then
        throw('discrete_uniform_dist_quantile: U must be in [0,1]')
    elseif U =:= 1.0 then
        X = B
    else
        X = A + floor(U * (B - A + 1))
    end.


% ----- Mean and Variance -----
discrete_uniform_dist_mean(A, B) = M =>
    discrete_uniform_dist_validate(A, B),
    M = 0.5 * (A + B).

discrete_uniform_dist_variance(A, B) = V =>
    discrete_uniform_dist_validate(A, B),
    N = (B - A + 1) * 1.0,
    V = (N * N - 1.0) / 12.0.



% Variants of discrete uniform 

% Draw a random integer from 0..N-1
random_integer(N) = random(0,N-1).
random_integer_n(N, Num) = [random_integer(N) : _ in 1..Num].

random_integer_pdf(N,X) = discrete_uniform_dist_pdf(0,N-1,X).
random_integer_cdf(N,X) = discrete_uniform_dist_cdf(0,N-1,X).
random_integer_quantile(N,X) = discrete_uniform_dist_quantile(0,N-1,X).
random_integer_mean(N,X) = discrete_uniform_dist_mean(0,N-1).
random_integer_variance(N,X) = discrete_uniform_dist_variance(0,N-1).


random_integer_dist(N) = random(0,N-1).
random_integer_dist_n(N, Num) = [random_integer_dist(N) : _ in 1..Num].

random_integer_dist_pdf(N,X) = discrete_uniform_dist_pdf(0,N-1,X).
random_integer_dist_cdf(N,X) = discrete_uniform_dist_cdf(0,N-1,X).
random_integer_dist_quantile(N,X) = discrete_uniform_dist_quantile(0,N-1,X).
random_integer_dist_mean(N,X) = discrete_uniform_dist_mean(0,N-1).
random_integer_dist_variance(N,X) = discrete_uniform_dist_variance(0,N-1).



% Draw a random integer from 1..N
random_integer1(N) = random(1,N).
random_integer1_n(N, Num) = [random_integer1(N) : _ in 1..Num].

random_integer1_pdf(N,X) = discrete_uniform_dist_pdf(1,N,X).
random_integer1_cdf(N,X) = discrete_uniform_dist_cdf(1,N,X).
random_integer1_quantile(N,X) = discrete_uniform_dist_quantile(1,N,X).
random_integer1_mean(N) = discrete_uniform_dist_mean(1,N).
random_integer1_variance(N,X) = discrete_uniform_dist_variance(1,N).

random_integer1_dist(N) = random(1,N).
random_integer1_dist_n(N, Num) = [random_integer1_dist(N) : _ in 1..Num].

random_integer1_dist_pdf(N,X) = discrete_uniform_dist_pdf(1,N,X).
random_integer1_dist_cdf(N,X) = discrete_uniform_dist_cdf(1,N,X).
random_integer1_dist_quantile(N,X) = discrete_uniform_dist_quantile(1,N,X).
random_integer1_dist_mean(N) = discrete_uniform_dist_mean(1,N).
random_integer1_dist_variance(N,X) = discrete_uniform_dist_variance(1,N).


%
% uniform_draw(L)
% Discrete uniform distribution over a numeric list.
% Example: L = [10,20,30,40].
%
% Draw a random value from the list List)
uniform_draw(List) = List[random(1,List.len)].
uniform_draw_n(List,N) = [uniform_draw(List) : _ in 1..N].

% PDF: Probability of value X.
uniform_draw_pdf(L, X) = cond( membchk(X, L), 1.0 / length(L), 0.0).

% CDF: P(Xi <= X)
%   Assumes numeric values in L.
uniform_draw_cdf(L, X) = P =>
    Sorted = sort(L),
    N = length(L),
    Count = length([V : V in Sorted, V =< X]),
    P = Count / N.

% Quantile: smallest value q in L such that CDF(q) >= P.
%   For P in [0,1].
uniform_draw_quantile(L, P) = Q =>
    Sorted = sort(L),
    N = length(L),
    Index = ceiling(P * N),
    I = max(1, min(N, Index)),
    Q = Sorted[I].

% Mean
uniform_draw_mean(L) = Mean =>
    Mean = sum(L) / length(L).

% Variance
uniform_draw_variance(L) = Var =>
    M = uniform_draw_mean(L),
    Var = sum([(X - M)**2 : X in L]) / length(L).




% Draw a random value from 1..N
dice() = dice(6).
dice_n(N) = dice_n(6,N).
dice(N) = uniform_draw(1..N).
dice_n(N,Num) = [dice(N) : _ in 1..Num].

/*

  Categorical distribution

  A discrete distribution over a finite set of values with specified probabilities.

  Parameters:
  Probs  - list of nonnegative probabilities summing to 1 (within 1.0e-12)
  Values - list of equally sized outcomes (labels, atoms, strings, or numbers)

  PMF over values:
  P(X = v) = Probs[i] if v == Values[i], else 0.

  CDF and Quantile:
  Since Values need not be ordered, the CDF/Quantile are defined over the
  INDEX 1..n. Use:
    categorical_dist_cdf_index(Probs, K)
    categorical_dist_quantile_index(Probs, U)
  If you want the quantile as a VALUE, use:
    categorical_dist_quantile_value(Probs, Values, U)

  Random generator:
  categorical_dist(Probs, Values) -> one sample in Values

  Mean and Variance:
  Defined only if Values are numeric. Otherwise these functions throw.

  Functions:
  categorical_dist_pdf(Probs, Values, V)
  categorical_dist_cdf_index(Probs, K)
  categorical_dist_quantile_index(Probs, U)
  categorical_dist_quantile_value(Probs, Values, U)
  categorical_dist(Probs, Values)
  categorical_dist_mean(Probs, Values)
  categorical_dist_variance(Probs, Values)

  Usage examples:
  % Labeled outcomes
  Probs  = [0.7, 0.2, 0.1],
  Values = ["A","B","C"],
  println(categorical_dist_pdf(Probs, Values, "B")).         % -> 0.2
  X = categorical_dist(Probs, Values), println(X).            % -> "A" or "B" or "C"
  println(categorical_dist_quantile_value(Probs, Values, 0.75)).

  % Numeric outcomes -> mean/variance available
  Probs2  = [0.5, 0.3, 0.2],
  Values2 = [10, 20, 30],
  println(categorical_dist_mean(Probs2, Values2)).
  println(categorical_dist_variance(Probs2, Values2)).

*/

%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%
% categorical(Probabilities, Values)
% ------------------------------------------------------------
% Draws a random element from Values according to Probabilities.
% Probabilities need not sum to 1; they are automatically normalized.
%
% Example:
%   categorical([0.1, 0.3, 0.6], ["red","green","blue"])
%   -> "blue"   (about 60% of the time)
%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%
categorical(Probs, Values) = V =>
    Len = length(Probs),
    if Len != length(Values) then
      throw($error('categorical: Probabilities and Values must have same length'))
    end,
    % Normalize probabilities
    Total = sum(Probs),
    if Total =< 0.0 then
        throw($error('categorical: probabilities must be positive'))
    else 
      Cum = [],
      S = 0.0,
      foreach(P in Probs)
          S := S + P / Total,
          Cum := Cum ++ [S]
      end,
      U = uniform(0,1),
      I = 1,
      while (I < Len, U > Cum[I])
         I := I + 1
      end,
      V = Values[I]
    end.

categorical_n(Probs, Values, N) = [categorical(Probs,Values) : _ in 1..N].

categorical_pdf(Probs,Values,X) = categorical_dist_pdf(Probs,Values,X).
categorical_cdf(Probs,Values,X) = categorical_dist_cdf(Probs,Values,X).
categorical_quantile_index(Probs,X) = categorical_dist_quantile_index(Probs,X).
categorical_quantile_value(Probs,Values,X) = categorical_dist_quantile_value(Probs,Values,X).
categorical_mean(Probs,Values) = categorical_dist_mean(Probs,Values).
categorical_variance(Probs,Values) = categorical_dist_variance(Probs,Values).

/* ========= Categorical distribution =========
   Support: X ∈ {Values[1], ..., Values[n]} with probabilities Probs[1..n]
   Requirements:
     - n >= 1
     - each Probs[i] >= 0
     - sum(Probs) == 1 within tolerance (no hidden renormalization)
     - Values need not be numeric; may be strings, atoms, numbers, etc.

   PMF over values:
     P(X = v) = Probs[i] if v = Values[i] for some i, else 0

   CDF/Quantile are defined over the INDEX (1..n), since Values may not be ordered.

   Random generator:
     categorical_dist(Probs, Values) -> one sample from Values
*/

% ----- validation -----
categorical_dist_validate(Probs, Values) =>
    N = Probs.length,
    if N =:= 0 then
        throw('categorical_dist: Probs must be non-empty')
    elseif N =\= Values.length then
        throw('categorical_dist: Probs and Values must have the same length')
    end,
    S = 0.0,
    foreach (P in Probs)
        if P < 0.0 then
            throw('categorical_dist: all probabilities must be >= 0')
        end,
        S := S + P
    end,
    if abs(S - 1.0) > 1.0e-12 then
        throw('categorical_dist: probabilities must sum to 1 (within 1.0e-12)')
    end.

% find index of a value in Values (first match), or 0 if not found
categorical_dist_index_of(Values, V) = I =>
    N = Values.length,
    I = 0,
    J = 1,
    Found = false,
    while (J =< N, not Found) do
        if Values[J] == V then
            I := J,
            Found := true
        else
            J := J + 1
        end
    end.

% ----- Random generator -----
% Draws a single value from Values according to Probs (uses u01())
categorical_dist(Probs, Values) = V =>
    categorical_dist_validate(Probs, Values),
    U = u01(),
    K = categorical_dist_quantile_index(Probs, U),
    V = Values[K].

categorical_dist_n(Probs, Values, N) = [categorical_dist(Probs, Values) : _ in 1..N].

% ----- PMF over values -----
categorical_dist_pdf(Probs, Values, V) = P =>
    categorical_dist_validate(Probs, Values),
    I = categorical_dist_index_of(Values, V),
    if I =:= 0 then
        P = 0.0
    else
        P = Probs[I]
    end.

% ----- CDF over index -----
% F(K) = sum_{i=1..floor(K)} Probs[i]
% K is allowed to be non-integer; floor is used.
categorical_dist_cdf_index(Probs, K) = F =>
    categorical_dist_validate(Probs, new_list(Probs.length)), % only length matters here
    if K < 1.0 then
        F = 0.0
    else
        N = Probs.length,
        if K >= N * 1.0 then
            F = 1.0
        else
            KK = floor(K),
            S = 0.0,
            I = 1,
            while (I =< KK) do
                S := S + Probs[I],
                I := I + 1
            end,
            F = S
        end
    end.

categorical_dist_cdf(Probs, Values, V) = F =>
    categorical_dist_validate(Probs, Values),
    I = categorical_dist_index_of(Values, V),
    if I =:= 0 then
        throw('categorical_dist_cdf: value not found in Values')
    else
        F = categorical_dist_cdf_index(Probs, I)
    end.

% ----- Quantile over index (left-continuous) -----
% Returns the smallest integer k in 1..N with CDF_index(k) >= U
categorical_dist_quantile_index(Probs, U) = K =>
    categorical_dist_validate(Probs, new_list(Probs.length)),
    if (U < 0.0 ; U > 1.0) then
        throw('categorical_dist_quantile_index: U must be in [0,1]')
    else
        K = _,
        N = Probs.length,
        if U =:= 0.0 then
            K = 1
        elseif U =:= 1.0 then
            K = N
        else
            Cum = 0.0,
            I = 1,
            Found = false,
            while (I =< N, not Found) do
                Cum := Cum + Probs[I],
                if U =< Cum then
                    K = I,
                    Found := true
                else
                    I := I + 1
                end
            end,
            if not Found then
                K = N
            end
        end
    end.

% Convenience: quantile returning the VALUE (not index)
categorical_dist_quantile_value(Probs, Values, U) = V =>
    categorical_dist_validate(Probs, Values),
    K = categorical_dist_quantile_index(Probs, U),
    V = Values[K].
    
categorical_dist_quantile(Probs, Values, U) = categorical_dist_quantile_value(Probs, Values, U).

% ----- Mean and Variance (only if Values are numeric) -----
categorical_dist_mean(Probs, Values) = M =>
    categorical_dist_validate(Probs, Values),
    N = Probs.length,
    % check numerics
    I = 1,
    while (I =< N) do
        Vi = Values[I],
        if not (integer(Vi) ; float(Vi)) then
            throw('categorical_dist_mean: Values must be numeric to compute mean')
        end,
        I := I + 1
    end,
    S = 0.0,
    J = 1,
    while (J =< N) do
        S := S + Probs[J] * (Values[J] * 1.0),
        J := J + 1
    end,
    M = S.

categorical_dist_variance(Probs, Values) = V =>
    M = categorical_dist_mean(Probs, Values),
    N = Probs.length,
    S2 = 0.0,
    I = 1,
    while (I =< N) do
        Xi = Values[I] * 1.0,
        S2 := S2 + Probs[I] * (Xi - M) * (Xi - M),
        I := I + 1
    end,
    V = S2.


/*
  Negative Hypergeometric distribution

  NOTE: This version has been replaced with the simpler
        Mathematica compliant version below.

  Interpretation:
  X = number of failures observed before the R-th success,
  when sampling without replacement from a population of size NTot that
  contains NSucc successes.

  Parameters:
  R      - target number of successes (integer, R >= 1 and R <= NSucc)
  NSucc  - number of successes in the population (integer, 0 <= NSucc <= NTot)
  NTot   - population size (integer, NTot >= 0)

  Support:
  X = 0, 1, ..., NTot - NSucc   (number of failures available)

  PMF:
  P(X = x) = [ C(x+R-1, x) * C(NTot - R - x, NSucc - R) ] / C(NTot, NSucc)

  CDF:
  F(k) = P(X <= k) computed by summing PMF from x=0 to floor(k).
  A stable recurrence is used: P(x+1)/P(x) =
      ((x + R)/(x + 1)) * ((NTot - NSucc - x)/(NTot - (x + R))).

  Quantile (left-continuous):
  The smallest integer x with F(x) >= u for u in [0,1].

  Mean and Variance:
  E[X] = R * (NTot - NSucc) / (NSucc + 1)
  Var[X] = R * (NTot - NSucc) * (NTot + 1) * (NSucc - R + 1)
           / ((NSucc + 1)^2 * (NSucc + 2))

  Random generation:
  Draw sequentially without replacement, counting successes and failures
  until R successes have been seen; return the number of failures.

  Usage examples:
  % Urn with NSucc=10 successes, NTot=30 total, stop at R=3 successes
  println(negative_hypergeometric_dist_pdf(3, 10, 30, 5)).
  println(negative_hypergeometric_dist_cdf(3, 10, 30, 5)).
  println(negative_hypergeometric_dist_quantile(3, 10, 30, 0.9)).
  X = negative_hypergeometric_dist(3, 10, 30), printf("Sample: %w\n", X).
  println(negative_hypergeometric_dist_mean(3, 10, 30)).
  println(negative_hypergeometric_dist_variance(3, 10, 30)).

  Notes:
  - This distribution is the "without-replacement" analog of the Negative
    Binomial (which assumes independent draws with replacement).
  - When NTot is large relative to NSucc and R, Negative Hypergeometric
    approximates Negative Binomial with p = NSucc/NTot.



*/
/* ========= Negative Hypergeometric distribution =========
   Parameters:
     R      - target number of successes (integer, R >= 1)
     NSucc  - number of successes in the population (integer, 0 <= NSucc <= NTot)
     NTot   - total population size (integer, NTot >= 0)
   Variable:
     X = number of failures before the R-th success

   Support:
     X = 0, 1, ..., (NTot - NSucc)     (requires R <= NSucc)

   PMF:
     P(X = x) = [ C(x+R-1, x) * C(NTot - R - x, NSucc - R) ] / C(NTot, NSucc)

   Recurrence (for summation/inversion):
     P(x+1)/P(x) = ((x + R) / (x + 1)) * ((NTot - NSucc - x) / (NTot - (x + R)))

   Mean:
     E[X] = R * (NTot - NSucc) / (NSucc + 1)

   Variance:
     Var[X] = R * (NTot - NSucc) * (NTot + 1) * (NSucc - R + 1)
              / ((NSucc + 1)^2 * (NSucc + 2))
*/

% ----- validation -----
% negative_hypergeometric_dist_validate(R, NSucc, NTot) =>
%     if not integer(R) then
%         throw($error('negative_hypergeometric_dist: R must be an integer'))
%     elseif not integer(NSucc) then
%         throw($error('negative_hypergeometric_dist: NSucc must be an integer'))
%     elseif not integer(NTot) then
%         throw($error('negative_hypergeometric_dist: NTot must be an integer'))
%     elseif R < 1 then
%         throw($error('negative_hypergeometric_dist: R must be >= 1'))
%     elseif NSucc < 0 then
%         throw($error('negative_hypergeometric_dist: NSucc must be >= 0'))
%     elseif NTot < 0 then
%         throw($error('negative_hypergeometric_dist: NTot must be >= 0'))
%     elseif NSucc > NTot then
%         throw($error('negative_hypergeometric_dist: NSucc must be <= NTot'))
%     elseif R > NSucc then
%         throw($error('negative_hypergeometric_dist: require R <= NSucc'))
%     end.

% % support upper bound (number of failures in population)
% negative_hypergeometric_dist_support_max(NSucc, NTot) = U =>
%    U = NTot - NSucc.

% log-choose via lgamma
% logchoose(A, B) = LG =>
%    LG = lgamma(A + 1.0) - lgamma(B + 1.0) - lgamma((A - B) + 1.0).

% ratio P(x+1)/P(x)
% negative_hypergeometric_dist_ratio(R, NSucc, NTot, X) = Ratio =>
%     Num1 = (X + R) * 1.0,
%     Den1 = (X + 1) * 1.0,
%     Num2 = (NTot - NSucc - X) * 1.0,
%     Den2 = (NTot - (X + R)) * 1.0,
%     Ratio = (Num1 / Den1) * (Num2 / Den2).

% % ----- Random generator: draw without replacement until R successes; return failures count -----
% negative_hypergeometric_dist(R, NSucc, NTot) = X =>
%     negative_hypergeometric_dist_validate(R, NSucc, NTot),
%     RemSucc = NSucc,
%     RemFail = NTot - NSucc,
%     GotSucc = 0,
%     Failures = 0,
%     while (GotSucc < R) do
%         Psucc = (RemSucc * 1.0) / ((RemSucc + RemFail) * 1.0),
%         U = u01(),
%         if U < Psucc then
%             GotSucc := GotSucc + 1,
%             RemSucc := RemSucc - 1
%         else
%             Failures := Failures + 1,
%             RemFail := RemFail - 1
%         end
%     end,
%     X = Failures.

% negative_hypergeometric_dist_n(R, NSucc, NTot, N) = [negative_hypergeometric_dist(R, NSucc, NTot) : _ in 1..N].

% % ----- PMF -----
% negative_hypergeometric_dist_pdf(R, NSucc, NTot, X) = P =>
%     negative_hypergeometric_dist_validate(R, NSucc, NTot),
%     if not integer(X) then
%         P = 0.0
%     else
%         U = negative_hypergeometric_dist_support_max(NSucc, NTot),
%         if (X < 0 ; X > U) then
%             P = 0.0
%         else
%             % log PMF: log C(X+R-1, X) + log C(NTot - R - X, NSucc - R) - log C(NTot, NSucc)
%             LP = logchoose((X + R - 1) * 1.0, X * 1.0)
%                + logchoose((NTot - R - X) * 1.0, (NSucc - R) * 1.0)
%                - logchoose(NTot * 1.0, NSucc * 1.0),
%             P = exp(LP)
%         end
%     end.

% % ----- CDF: F(K) = P(X <= K) -----
% negative_hypergeometric_dist_cdf(R, NSucc, NTot, K) = F =>
%     negative_hypergeometric_dist_validate(R, NSucc, NTot),
%     U = negative_hypergeometric_dist_support_max(NSucc, NTot),
%     if K < 0.0 then
%         F = 0.0
%     elseif K >= U * 1.0 then
%         F = 1.0
%     else
%         Kf = floor(K),
%         % start at x=0
%         P0 = negative_hypergeometric_dist_pdf(R, NSucc, NTot, 0),
%         if Kf =:= 0 then
%             F = P0
%         else
%             S = P0,
%             P = P0,
%             X = 0,
%             while (X < Kf) do
%                 RATIO = negative_hypergeometric_dist_ratio(R, NSucc, NTot, X),
%                 P := P * RATIO,
%                 S := S + P,
%                 X := X + 1
%             end,
%             F = S
%         end
%     end.

% % ----- Quantile (left-continuous): smallest x with F(x) >= U, U in [0,1] -----
% negative_hypergeometric_dist_quantile(R, NSucc, NTot, UProb) = Q =>
%     negative_hypergeometric_dist_validate(R, NSucc, NTot),
%     if (UProb < 0.0 ; UProb > 1.0) then
%         throw('negative_hypergeometric_dist_quantile: U must be in [0,1]')
%     else
%         Q = _,
%         U = negative_hypergeometric_dist_support_max(NSucc, NTot),
%         P = negative_hypergeometric_dist_pdf(R, NSucc, NTot, 0),
%         C = P,
%         if UProb =< C then
%             Q = 0
%         else
%             X = 0,
%             Found = false,
%             while (X < U, not Found) do
%                 RATIO = negative_hypergeometric_dist_ratio(R, NSucc, NTot, X),
%                 P := P * RATIO,
%                 C := C + P,
%                 X := X + 1,
%                 if UProb =< C then
%                     Q = X,
%                     Found := true
%                 end
%             end,
%             if not Found then
%                 Q = U
%             end
%         end
%     end.


% % ----- Mean and Variance -----
% negative_hypergeometric_dist_mean(R, NSucc, NTot) = M =>
%     negative_hypergeometric_dist_validate(R, NSucc, NTot),
%     M = (R * 1.0) * ((NTot - NSucc) * 1.0) / ((NSucc + 1) * 1.0).

% negative_hypergeometric_dist_variance(R, NSucc, NTot) = V =>
%     negative_hypergeometric_dist_validate(R, NSucc, NTot),
%     Num = (R * 1.0) * ((NTot - NSucc) * 1.0) * (NTot + 1.0) * ((NSucc - R + 1) * 1.0),
%     Den = ((NSucc + 1.0) * (NSucc + 1.0)) * (NSucc + 2.0),
%     V = Num / Den.


/*

  From Mathematica BetaBinomialDistribution
  """
  Define a negative hypergeometric distribution:
  NegativeHypergeometricDistribution(w_, wtot_, btot_) := BetaBinomialDistribution(w, wtot - w + 1, btot)

  See https://en.wikipedia.org/wiki/Negative_hypergeometric_distribution

  Note: This is just piggy backing on the beta_binomial distribution functions.

*/
negative_hypergeometric_dist(W,WTot,BTot) = beta_binomial_dist(BTot,W,WTot-W+1).
negative_hypergeometric_dist_n(W,WTot,BTot,N) = [negative_hypergeometric_dist(W,WTot,BTot) : _ in 1..N].

negative_hypergeometric_dist_pdf(W,WTot,BTot,X) = beta_binomial_dist_pdf(BTot,W,WTot-W+1,X).
negative_hypergeometric_dist_cdf(W,WTot,BTot,X) = beta_binomial_dist_cdf(BTot,W,WTot-W+1,X).
negative_hypergeometric_dist_quantile(W,WTot,BTot,X) = beta_binomial_dist_quantile(BTot,W,WTot-W+1,X).
negative_hypergeometric_dist_mean(W,WTot,BTot) = beta_binomial_dist_mean(BTot,W,WTot-W+1).

/*
  NOTE: This is not used since it relies on the now replaced previous version
        of negative_hypergeometric_dist.

  Negative Hypergeometric Draws (draws-until-R-th-success)
  A.k.a. Inverse Hypergeometric Distribution
  (also known as Stopped Negative Hypergeometric


  Variable:
  Y = number of draws when the R-th success occurs, sampling without replacement
  from a population of size NTot with NSucc successes.

  Relation to the "failures-before" parameterization:
  Let X = number of failures before the R-th success (the previous implementation).
  Then Y = X + R.

  Support:
  Y = R, R+1, ..., NTot - NSucc + R

  PMF, CDF, Quantile:
  PDF_Y(y) = PDF_X(y - R)
  CDF_Y(k) = CDF_X(k - R)
  Q_Y(u)   = Q_X(u) + R
  where PDF_X, CDF_X, Q_X are from negative_hypergeometric_dist_*.

  Mean and Variance:
  E[Y]   = E[X] + R = R * (NTot - NSucc) / (NSucc + 1) + R
  Var[Y] = Var[X]   = R * (NTot - NSucc) * (NTot + 1) * (NSucc - R + 1)
                      / ((NSucc + 1)^2 * (NSucc + 2))

  Functions:
  negative_hypergeometric_draws_dist_pdf(R, NSucc, NTot, Y)
  negative_hypergeometric_draws_dist_cdf(R, NSucc, NTot, K)
  negative_hypergeometric_draws_dist_quantile(R, NSucc, NTot, U)
  negative_hypergeometric_draws_dist(R, NSucc, NTot)              % RNG
  negative_hypergeometric_draws_dist_mean(R, NSucc, NTot)
  negative_hypergeometric_draws_dist_variance(R, NSucc, NTot)

  Usage examples:
  % Population: NTot=30, successes NSucc=10, stop at R=3rd success
  println(negative_hypergeometric_draws_dist_pdf(3, 10, 30, 7)).
  println(negative_hypergeometric_draws_dist_cdf(3, 10, 30, 7)).
  println(negative_hypergeometric_draws_dist_quantile(3, 10, 30, 0.9)).
  Y = negative_hypergeometric_draws_dist(3, 10, 30), printf("Sample Y: %w\n", Y).
  println(negative_hypergeometric_draws_dist_mean(3, 10, 30)).
  println(negative_hypergeometric_draws_dist_variance(3, 10, 30)).
*/

/* ========= Negative Hypergeometric (draws-until-R-th-success) =========
   Variable:
     Y = number of draws when the R-th success occurs
   Relation to "failures-before" X:
     Y = X + R

   Support:
     Y = R, R+1, ..., R + (NTot - NSucc)   (i.e., R .. NTot - NSucc + R)

   Mapping:
     PDF_Y(y) = PDF_X(y - R)
     CDF_Y(k) = CDF_X(k - R)
     Q_Y(u)   = Q_X(u) + R
     E[Y]     = E[X] + R
     Var[Y]   = Var[X]
*/

% ----- validation (same constraints as the X-parameterization) -----
% negative_hypergeometric_draws_dist_validate(R, NSucc, NTot) =>
%     negative_hypergeometric_dist_validate(R, NSucc, NTot).

% % support bounds for Y
% negative_hypergeometric_draws_dist_support(R, NSucc, NTot) = [L, U] =>
%     L = R,
%     U = R + (NTot - NSucc).

% % ----- Random generator: simulate X and add R -----
% negative_hypergeometric_draws_dist(R, NSucc, NTot) = Y =>
%     negative_hypergeometric_draws_dist_validate(R, NSucc, NTot),
%     X = negative_hypergeometric_dist(R, NSucc, NTot),
%     Y = X + R.

% negative_hypergeometric_draws_dist_n(R, NSucc, NTot, N) = [negative_hypergeometric_draws_dist(R, NSucc, NTot) : _ in 1..N].

% % ----- PMF -----
% negative_hypergeometric_draws_dist_pdf(R, NSucc, NTot, Y) = P =>
%     negative_hypergeometric_draws_dist_validate(R, NSucc, NTot),
%     if not integer(Y) then
%         P = 0.0
%     else
%         [L, U] = negative_hypergeometric_draws_dist_support(R, NSucc, NTot),
%         if (Y < L ; Y > U) then
%             P = 0.0
%         else
%             X = Y - R,
%             P = negative_hypergeometric_dist_pdf(R, NSucc, NTot, X)
%         end
%     end.

% % ----- CDF: F(k) = P(Y <= k) = P(X <= k - R) -----
% negative_hypergeometric_draws_dist_cdf(R, NSucc, NTot, K) = F =>
%     negative_hypergeometric_draws_dist_validate(R, NSucc, NTot),
%     [L, U] = negative_hypergeometric_draws_dist_support(R, NSucc, NTot),
%     if K < L * 1.0 then
%         F = 0.0
%     elseif K >= U * 1.0 then
%         F = 1.0
%     else
%         F = negative_hypergeometric_dist_cdf(R, NSucc, NTot, floor(K) - R)
%     end.

% % ----- Quantile (left-continuous): smallest y with F(y) >= U -----
% negative_hypergeometric_draws_dist_quantile(R, NSucc, NTot, U) = Q =>
%     negative_hypergeometric_draws_dist_validate(R, NSucc, NTot),
%     if (U < 0.0 ; U > 1.0) then
%         throw('negative_hypergeometric_draws_dist_quantile: U must be in [0,1]')
%     else
%         Qx = negative_hypergeometric_dist_quantile(R, NSucc, NTot, U),
%         Q  = Qx + R
%     end.

% % ----- Mean and Variance -----
% negative_hypergeometric_draws_dist_mean(R, NSucc, NTot) = M =>
%     negative_hypergeometric_draws_dist_validate(R, NSucc, NTot),
%     Mx = negative_hypergeometric_dist_mean(R, NSucc, NTot),
%     M = Mx + (R * 1.0).

% negative_hypergeometric_draws_dist_variance(R, NSucc, NTot) = V =>
%     negative_hypergeometric_draws_dist_validate(R, NSucc, NTot),
%     V = negative_hypergeometric_dist_variance(R, NSucc, NTot).



/*
   Gaussian (Normal) distribution

   From Handbook on probability distributions
   page 49
   """
   The Box-Muller algorithm produces normal random variates:
   * generate U, V from a uniform U(0, 1) distribution,
   * compute X = sqrt(−2*log(U))*cos(2*π*V) and Y = sqrt(−2*log(U))*sin(2*π*V ).
   In outputs, X and Y follow a standard normal distribution (independently).
   ...

   But there appears that this algorithm under estimates the tail
   of the distribution (called the Neave effect, cf. Patard (2007)),
   most softwares use the inversion function method, consist in
   computing the quantile function Φ-1 of a uniform variate. 
   """
*/
normal01() = Y =>
    U = frand(0,1),
    V = frand(0,1),
    % X = sqrt(-2*log(U))*cos(2*math.pi*V),
    Y = sqrt(-2*log(U))*sin(2*math.pi*V).

normal_dist(Mean,Stdev) = Mean + (normal01() * Stdev).
normal_dist_n(Mean,Stdev,N) = [normal_dist(Mean,Stdev) : _ in 1..N].

%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%
% normal_dist_rand(Mean, StdDev)
% ------------------------------------------------------------
% Draws a random sample from Normal(Mean, StdDev)
% using the Box–Muller transform.
%
% Example:
%   normal_dist_rand(0.0, 1.0)  -> standard normal
%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%
normal_dist2(Mu, Sigma) = X =>
    if Sigma =< 0.0 then
        throw($error('normal_dist_rand: standard deviation must be positive'))
    else
        U1 = uniform(0.0, 1.0),
        U2 = uniform(0.0, 1.0),
        % ensure U1 > 0 to avoid log(0)
        U1a = cond(U1 =< 0.0, 1.0e-12, U1),
        Z0 = sqrt(-2.0 * log(U1a)) * cos(2.0 * 3.141592653589793 * U2),
        X = Mu + Sigma * Z0
    end.

% %%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%
% % Normal log-pdf (μ, σ)
% %%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%
% gaussian_dist_logpdf(Mu, Sigma, X) = L =>
%     if Sigma =< 0.0 then
%         L = log_zero()
%     else
%         Z = (X - Mu) / Sigma,
%         L = -0.5*Z*Z - log(Sigma) - 0.5*log(2.0*3.141592653589793)
%     end.

%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%
% Standard-normal PDF φ(z)
%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%
stdnormal_pdf(Z) = Pdf =>
    Pdf = exp(-0.5 * Z * Z) / sqrt(2.0 * 3.141592653589793).


%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%
% normal_dist_pdf(Mean, StdDev, X)
% ------------------------------------------------------------
% Returns the probability density f(X; Mean, StdDev)
% for the Normal (Gaussian) distribution.
%
% Example:
%   normal_dist_pdf(0.0, 1.0, 0.0)  -> 0.3989422804014327
%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%
normal_dist_pdf(Mu, Sigma, X) = F =>
    if Sigma =< 0.0 then
        throw($error('normal_dist_pdf: standard deviation must be positive'))
    else
        Z = (X - Mu) / Sigma,
        F = (1.0 / (Sigma * sqrt(2.0 * 3.141592653589793))) * exp(-0.5 * Z * Z)
    end.


%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%
% normal_dist_cdf(Mu, Sigma, X)
% High-precision Normal CDF (≈ 1e-16 absolute error).
% Based on G. Marsaglia, “Evaluating the Normal Distribution”.
% For numerical safety, we clamp for |Z| > 8.
%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%
normal_dist_cdf(Mu, Sigma, X) = F =>
    if Sigma =< 0.0 then
        throw($error('normal_dist_cdf: standard deviation must be positive'))
    else
        Z = (X - Mu) / Sigma,
        F = phi_marsaglia(Z)
    end.


%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%
% normal_dist_quantile(Mu, Sigma, P)
% --------------------------------------------------------------
% Returns x such that P(X <= x) = P for X ~ Normal(Mu, Sigma^2).
% Accuracy: ~1e-16 absolute (double-precision limit).
%
% Domain:
%   0 < P < 1, Sigma > 0
% Special cases:
%   P = 0 -> -Huge, P = 1 -> +Huge (numerical infinities)
%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%
% 
% Note that this is not a full precision of 10^-17 for low/high tails.
% 
normal_dist_quantile1(Mu, Sigma, P) = X =>
    if Sigma =< 0.0 then
        throw($error('normal_dist_quantile: Sigma must be positive'))
    elseif P =< 0.0 then
        X = Mu + Sigma * (-1.0e300)            % represent -∞
    elseif P >= 1.0 then
        X = Mu + Sigma * ( 1.0e300)            % represent +∞
    else
        Z0 = invnorm_acklam(P),                 % initial z
        % One Newton step: z1 = z0 - (Phi(z0) - P)/phi(z0)
        Phi0 = phi_marsaglia(Z0),
        Pdf0 = stdnormal_pdf(Z0),
        Z1 = Z0 - (Phi0 - P)/Pdf0,
        X = Mu + Sigma * Z1
    end.


%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%
% normal_dist_quantile(Mu, Sigma, P)
% --------------------------------------------------------------
% Returns X such that P(X' <= X) = P for X' ~ Normal(Mu, Sigma^2).
% Acklam (2003) piecewise rational approximation in symmetric form.
% Accuracy: ~1e-9..1e-12. Stable in tails.
%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%
normal_dist_quantile2(Mu, Sigma, P) = X =>
    if Sigma =< 0.0 then
        throw($error('normal_dist_quantile: Sigma must be positive'))
    elseif P =< 0.0 then
        X = Mu + Sigma * (-1.0e300)   % represent -∞
    elseif P >= 1.0 then
        X = Mu + Sigma * ( 1.0e300)   % represent +∞
    else
        Z = invnorm_acklam_symmetric(1-P), % hakank: note that it's 1-P
        X = Mu + Sigma * Z
    end.

%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%
% normal_dist_quantile(Mu, Sigma, P)
% Acklam-based quantile with enforced tail signs (+ optional Newton)
% Accuracy: ~1e-12 from Acklam; with one Newton step ~1e-15 (|z|<=8)
%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%
normal_dist_quantile(Mu, Sigma, P) = X =>
    if Sigma =< 0.0 then
        throw($error('normal_dist_quantile: Sigma must be positive'))
    elseif P =< 0.0 then
        X = Mu + Sigma * (-1.0e300)        % -infinity sentinel
    elseif P >= 1.0 then
        X = Mu + Sigma * ( 1.0e300)        % +infinity sentinel
    else
        Z0 = invnorm_acklam_enforced(P),

        % --- Optional: single Newton refinement using your high-precision CDF ---
        % Safe when |Z0| <= 8 (covers p in [~6e-16, 1-6e-16])
        if abs(Z0) =< 8.0 then
          Phi0 = normal_std_cdf(Z0),      % standard-normal CDF you already have
          Pdf0 = stdnormal_pdf(Z0),
          Z = Z0 - (Phi0 - P)/Pdf0            % Newton step
        else
          Z = Z0
        end,

        X = Mu + Sigma * Z
    end.

%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%
% High-precision standard-normal CDF Φ(z) you already have
% (Rename your phi_marsaglia/1 to normal_std_cdf/1 or wrap it.)
%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%
normal_std_cdf(Z) = P =>
    % If you've named it differently, just call your function here.
    % This wrapper assumes you have: normal_dist_cdf(0.0, 1.0, Z)
    P = normal_dist_cdf(0.0, 1.0, Z).


normal_dist_mean(Mu,Sigma) = Mu.

normal_dist_variance(Mu,Sigma) = Sigma*Sigma.

/* 
  Geometric distribution

  The geometric distribution is either one of two 
  distributions:
  - number of failures until success  <--
  - number of trials until success

  In PPL, the geometric distribution refers to 
  the latter: the number of failures until success
  were the probability of success is the parameter P.

  The second variant is defined below as shifted_geometric_dist.
  
*/
geometric1(P,X) = G => 
    U = frand(0,1),
    if U > P then  G = geometric1(P,X+1) else G = X end.

geometric_dist1(P) = geometric1(P,0).

% RNG: failures before first success (0-based)
geometric_dist(P) = K =>
    if P =< 0.0 then
        throw($error('geometric_dist: P must be in (0,1]'))
    elseif P >= 1.0 then
        K = 0                              % degenerate at 0 when p=1
    else
        U = u01(),
        K = floor( log(U) / log(1.0 - P) ).to_int()
    end.

geometric_dist_n(P, N) = [geometric_dist(P) : _ in 1..N].

% PMF
geometric_dist_pdf(P, K) = Pr =>
    if P =< 0.0 then
        throw($error('geometric_dist_pdf: P must be in (0,1]'))
    elseif K < 0 ; not integer(K) then
        Pr = 0.0
    elseif P >= 1.0 then
        Pr = cond(K == 0, 1.0, 0.0)        % all mass at 0
    else
        Pr = pow(1.0 - P, K) * P
    end.


% CDF
geometric_dist_cdf(P, K) = Cdf =>
    if P =< 0.0 then
        throw($error('geometric_dist_cdf: P must be in (0,1]'))
    elseif K < 0 then
        Cdf = 0.0
    elseif P >= 1.0 then
        Cdf = 1.0                          % F(k)=1 for all k>=0 when p=1
    else
        Cdf = 1.0 - pow(1.0 - P, K + 1)
    end.

% Quantile: smallest K s.t. CDF >= QProb
geometric_dist_quantile(P, QProb) = K =>
    if P =< 0.0 then
        throw($error('geometric_dist_quantile: P must be in (0,1]'))
    elseif QProb =< 0.0 then
        K = 0
    elseif P >= 1.0 then
        K = 0                              % degenerate at 0 for any q∈(0,1]
    elseif QProb >= 1.0 then
        throw($error('geometric_dist_quantile: QProb must be < 1 when 0<P<1'))
    else
        Num = log(1.0 - QProb),
        Den = log(1.0 - P),
        K = (ceiling(Num / Den) - 1).to_int()
    end.

% Mean
geometric_dist_mean(P) = M =>
    if P =< 0.0 then
        throw($error('geometric_dist_mean: P must be in (0,1]'))
    elseif P >= 1.0 then
        M = 0.0
    else
        M = (1.0 - P) / P
    end.

% Variance
geometric_dist_variance(P) = V =>
    if P =< 0.0 then
        throw($error('geometric_dist_variance: P must be in (0,1]'))
    elseif P >= 1.0 then
        V = 0.0
    else
        V = (1.0 - P) / (P * P)
    end.


/*
----------------------------------------------------------------------
shifted_geometric_dist(P)
----------------------------------------------------------------------
Random generation from the shifted geometric distribution
(trials until first success).

Support: K = 1, 2, 3, ...

Method:
  Use inverse CDF sampling.

  If U ~ Uniform(0,1),
  then K = ceil( ln(1-U) / ln(1-P) )

Special cases:
  - If P = 1.0, always return 1.
  - Throws domain_error for invalid P.

Returns an integer >= 1.
----------------------------------------------------------------------
*/
shifted_geometric_dist(P) = K =>
    if P =< 0.0 ; P > 1.0 then
        throw($error('domain_error(probability,P),
                     shifted_geometric_dist/1'))
    end,

    % P = 1 => immediate success
    if P =:= 1.0 then
        K = 1
    else
        U = uniform(0,1),                % uniform in [0,1)
        Num = log(1.0 - U),
        Den = log(1.0 - P),
        RawK = Num / Den,
        K0 = ceiling(RawK),
        if K0 < 1 then
            K = 1
        else
            K = K0
        end
    end.
shifted_geometric_dist_n(P,N) = [shifted_geometric_dist(P) : _ in 1..N].

/*
----------------------------------------------------------------------
shifted_geometric_dist_pdf(P, K)
----------------------------------------------------------------------
PDF of the shifted geometric distribution.

Definition:
  Counts the number of *trials* until first success.
  Support: K = 1, 2, 3, ...

PMF:
  P(K) = (1 - P)^(K-1) * P

Parameters:
  P : success probability, 0 < P =< 1
  K : number of trials (integer >= 1)

Notes:
  - This is the "shifted" geometric distribution, also known as
    the "trials until success" variant.
  - It is related to the standard geometric (failures-based) by
        shifted = failures + 1
----------------------------------------------------------------------
*/

shifted_geometric_dist_pdf(P, K) = Value =>
    % Validate P
    if P =< 0.0 ; P > 1.0 then
        throw($error('domain_error(probability,P),shifted_geometric_dist_pdf/2'))
    end,

    % Validate K
    if K < 1 then
       throw($error('shifted_geometric_dist_pdf/2: K >= 1'))
    end,

    % Special case: P = 1 -> always success at trial 1
    if P =:= 1.0 then
        if K =:= 1 then
            Value = 1.0
        else
            Value = 0.0
        end
    else
        % Core formula: (1-P)^(K-1) * P
        Value = pow(1.0 - P, K - 1) * P
    end.

/*
----------------------------------------------------------------------
shifted_geometric_dist_cdf(P, K)
----------------------------------------------------------------------
CDF of the shifted geometric distribution (trials until success).

Definition:
  Counts the number of *trials* K = 1,2,3,... until first success.

CDF:
  F(K) = 1 - (1 - P)^K

Parameters:
  P : success probability, 0 < P =< 1
  K : integer >= 1

Special cases:
  - If P = 1.0 then F(K) = 1.0 for all K >= 1.
  - If K < 1, domain error.
----------------------------------------------------------------------
*/

shifted_geometric_dist_cdf(P, K) = Value =>
    % Validate P
    if P =< 0.0 ; P > 1.0 then
        throw($error('domain_error(probability,P),shifted_geometric_dist_cdf/2'))
    end,

    % Validate K
    if K < 1 then
       throw($error('shifted_geometric_dist_cdf/2: K >= 1'))
    end,

    % Handle trivial success probability
    if P =:= 1.0 then
        Value = 1.0
    else
        % Closed-form CDF: 1 - (1-P)^K
        Value = 1.0 - pow(1.0 - P, K)
    end.


/*
----------------------------------------------------------------------
shifted_geometric_dist_quantile(P, Q)
----------------------------------------------------------------------
Quantile function for the shifted geometric distribution
(trials until first success).

Definition:
  K = min integer >= 1 such that:
        1 - (1 - P)^K >= Q

Closed form:
  K = ceil( ln(1-Q) / ln(1-P) )

Parameters:
  P : success probability, 0 < P =< 1
  Q : quantile, 0 =< Q =< 1

Special cases:
  - If P = 1.0, always succeed on first trial => return 1.
  - If Q = 0.0, quantile is 1.
  - If Q = 1.0, quantile is +inf (but practically return a big K or
    use the formula; (1-P)^K -> 0 as K -> infinity).

Notes:
  - This is the "trials until success" geometric.
  - Quantile always returns an integer >= 1.
----------------------------------------------------------------------
*/

shifted_geometric_dist_quantile(P, Q) = K =>
    % Validate P
    if P =< 0.0 ; P > 1.0 then
        throw($error('domain_error(probability,P),
                     shifted_geometric_dist_quantile/2'))
    end,

    % Validate Q
    if Q < 0.0 ; Q > 1.0 then
        throw($error('domain_error(probability,Q),
                     shifted_geometric_dist_quantile/2'))
    end,

    % Special cases
    if P =:= 1.0 then
        K = 1
    elseif Q =:= 0.0 then
        K = 1
    elseif Q =:= 1.0 then
        % No finite K satisfies CDF(K) = 1 for P<1
        % Use a large K or approximate:
        % (1-P)^K <= eps, so K ~ ln(eps)/ln(1-P)
        K = 1000000000    % practical infinity
    else
        % Core formula:
        % K = ceil( ln(1-Q) / ln(1-P) )
        Num = log(1.0 - Q),
        Den = log(1.0 - P),
        RawK = Num / Den,
        K0 = ceiling(RawK),
        % Ensure at least 1
        if K0 < 1 then
            K = 1
        else
            K = K0
        end
    end.


/*
----------------------------------------------------------------------
shifted_geometric_dist_mean(P)
----------------------------------------------------------------------
Mean of the shifted geometric distribution (trials until success).

Formula:
  Mean = 1 / P

Valid range:
  0 < P =< 1
----------------------------------------------------------------------
*/
shifted_geometric_dist_mean(P) = Mean =>
    if P =< 0.0 ; P > 1.0 then
        throw($error('domain_error(probability,P),
                     shifted_geometric_dist_mean/1'))
    end,
    Mean = 1 / P.

/*
----------------------------------------------------------------------
shifted_geometric_dist_variance(P)
----------------------------------------------------------------------
Variance of the shifted geometric distribution.

Formula:
  Variance = (1 - P) / P^2

Valid range:
  0 < P =< 1
----------------------------------------------------------------------
*/
shifted_geometric_dist_variance(P) = Var =>
    if P =< 0.0 ; P > 1.0 then
        throw($error('domain_error(probability,P),
                     shifted_geometric_dist_variance/1'))
    end,
    Var = (1 - P) / (P * P).



/*
  flip(P) 
  returns true/false depending on probability P 

  See bernouilli_dist for the version that returns 0 or 1.

*/
flip(P) = cond(bernoulli_dist(P)==1,true,false). 
flip() = flip(1/2).
flip_n(N) = [flip(1/2) : _ in 1..N].
flip_n(P, N) = [flip(P) : _ in 1..N].

/*
  Bernoulli distribution

  Returns 0/1 depending on probability P 

*/
bernoulli_dist(P) = cond(U <= P,1,0) => U = uniform(0,1).
bernoulli_dist_n(P,N) = [bernoulli_dist(P) : _ in 1..N].

% Poor man's binomial
bernoulli_dist_sum(P,N) = bernoulli_dist_n(P,N).sum.

% alias
bern(P) = bernoulli_dist(P).
bern_n(P,N) = bernoulli_dist_n(P,N).
bern_sum(P,N) = bernoulli_dist_n(P,N).sum.


% ========= Bernoulli distribution (X ∈ {0,1}, parameter P ∈ [0,1]) =========

% PDF: P(X = x)
bernoulli_dist_pdf(P, X) = R =>
    if (P < 0.0 ; P > 1.0) then
        throw('bernoulli_dist_pdf: P must be in [0,1]')
    elseif X =:= 0 then
        R = 1.0 - P
    elseif X =:= 1 then
        R = P
    else
        R = 0.0
    end.

% CDF: F(k) = P(X ≤ k)
bernoulli_dist_cdf(P, K) = R =>
    if (P < 0.0 ; P > 1.0) then
        throw('bernoulli_dist_cdf: P must be in [0,1]')
    elseif K < 0.0 then
        R = 0.0
    elseif K < 1.0 then
        R = 1.0 - P
    else
        R = 1.0
    end.

% Quantile (left-continuous): smallest x ∈ {0,1} with F(x) ≥ U
% U ∈ [0,1]
bernoulli_dist_quantile(P, U) = Q =>
    if (P < 0.0 ; P > 1.0) then
        throw('bernoulli_dist_quantile: P must be in [0,1]')
    elseif (U < 0.0 ; U > 1.0) then
        throw('bernoulli_dist_quantile: U must be in [0,1]')
    else
        Q = cond(U =< 1.0 - P, 0, 1)
    end.

% Random generator: returns 1 with probability P, else 0
% Uses u01() that avoids endpoints {0,1}
bernoulli_dist_rand(P) = X =>
    if (P < 0.0 ; P > 1.0) then
        throw('bernoulli_dist_rand: P must be in [0,1]')
    else
        U = u01(),
        X = cond(U < P, 1, 0)
    end.

% Mean = P
bernoulli_dist_mean(P) = M =>
    if (P < 0.0 ; P > 1.0) then
        throw('bernoulli_dist_mean: P must be in [0,1]')
    else
        M = P
    end.

% Variance = P*(1-P)
bernoulli_dist_variance(P) = V =>
    if (P < 0.0 ; P > 1.0) then
        throw('bernoulli_dist_variance: P must be in [0,1]')
    else
        V = P * (1.0 - P)
    end.




% BERNOULLI log-pmf:
% X ∈ {0,1}, P ∈ [0,1]
% log P(X=x) = x*log(P) + (1-x)*log(1-P)
% bernoulli_dist_logpdf(P, X) = L =>
%     if (X != 0, X != 1) ; P < 0.0 ; P > 1.0 then
%         L = log_zero()
%     elseif P == 0.0 then
%         L = cond(X == 0, 0.0, log_zero())     % deterministic at 0
%     elseif P == 1.0 then
%         L = cond(X == 1, 0.0, log_zero())     % deterministic at 1
%     else
%         L = cond(X == 1, log(P), log(1.0 - P))
%     end.

/*
  Binomial distribution

  From Handbook on probability distributions
   page 8
   """
   It is easy to simulate Bernoulli distribution with the following heuristic:
   * generate U from a uniform distribution,
   * compute X as 1 if U <= p and 0 otherwise.

   The binomial distribution is obtained by summing n i.i.d. Bernoulli random variates.
   """

*/

binomial_dist1(N,P) = sum([bernoulli_dist(P) : _ in 1..N]).
binomial_dist_n(N,P,Num) = [binomial_dist(N,P) : _ in 1..Num].

binomial_dist(N,P) = K =>
    if (P < 0.0 ; P > 1.0) then
        throw('binomial_dist: P must be in [0,1]')
    elseif N < 0 then
        throw('binomial_dist: N must be >= 0')
    elseif N =:= 0 then
        K = 0
    else
        K = 0,
        I = 1,
        while (I =< N) do
            if u01() < P then
                K := K + 1
            end,
            I := I + 1
        end
    end.

% A smarter/faster version for larger N
binomial_dist_smart(N,P) = K =>
    if (P < 0.0 ; P > 1.0) then
        throw('binomial_dist: P must be in [0,1]')
    elseif N < 0 then
        throw('binomial_dist: N must be >= 0')
    elseif N =:= 0 then
        K = 0
    elseif N > 1000 then
        % use normal approximation
        Mu = N * P,
        Sigma = sqrt(N * P * (1.0 - P)),
        K = round(normal_dist(Mu, Sigma))
    elseif N * P < 10.0 then
        % use Poisson approximation
        K = poisson_dist(N * P)
    else
        % exact sampling
        K = 0,
        I = 1,
        while (I =< N) do
            if u01() < P then
                K := K + 1
            end,
            I := I + 1
        end
    end.
binomial_dist_smart_n(N,P,Num) = [binomial_dist_smart(N,P) : _ in 1..Num].


/* Binomial BTPE method
   Reference: Kachitvichyanukul & Schmeiser (1988)
   Parameters:
      N - number of trials (integer, >= 0)
      P - success probability (0 <= P <= 1)
   Returns: integer K in 0..N
*/

/* Binomial BTPE method
   Reference: Kachitvichyanukul & Schmeiser (1988)
   Parameters:
      N - number of trials (integer, >= 0)
      P - success probability (0 <= P <= 1)
   Returns: integer K in 0..N
*/
% Fast exact Binomial sampler by mode-centered expansion (no break/continue).
% Draws K ~ Binomial(N,P).
% Exact: uses PMF recurrences; no approximations.
% Efficient when variance is moderate/large (compared to summing N Bernoullis).
binomial_dist_btpe(N, P) = K =>
    % Validate
    if (P < 0.0 ; P > 1.0) then
        throw('binomial_dist: P must be in [0,1]')
    elseif N < 0 then
        throw('binomial_dist: N must be >= 0')
    elseif N =:= 0 then
        K = 0
    elseif P =:= 0.0 then
        K = 0
    elseif P =:= 1.0 then
        K = N
    else
        % Use symmetry to keep P <= 0.5 for stability; flip result if needed.
        Flip = (P > 0.5),
        PP = cond(Flip, 1.0 - P, P),
        QQ = 1.0 - PP,

        % For very small mean, direct loop is OK (keeps tails cheap).
        if N * PP < 30.0 then
            KK = 0,
            I = 1,
            while (I =< N) do
                if u01() < PP then
                    KK := KK + 1
                end,
                I := I + 1
            end,
            K0 = KK
        else
            % Mode m = floor((N+1)*p)
            M = floor((N + 1.0) * PP),

            % Compute pmf at mode using log to avoid overflow, then exp.
            % pmf(m) = C(N,m) * p^m * q^(N-m)
            LnPM = lgamma(N + 1.0) - lgamma(M + 1.0) - lgamma((N - M) + 1.0)
                   + M * log(PP) + (N - M) * log(QQ),
            PM = exp(LnPM),

            % Draw U in (0,1) and accumulate from the mode outward.
            U = u01(),
            Cum = PM,
            if U =< Cum then
                K0 = M
            else
                % Prepare left and right runners
                L = M - 1,
                R = M + 1,

                % Next left mass: P(L) from PM via one-step ratio
                PL = cond(L >= 0,
                          PM * ((M * 1.0) / ((N - M + 1.0))) * (QQ / PP),
                          0.0),
                % Next right mass: P(R) from PM via one-step ratio
                PR = cond(R =< N,
                          PM * (((N - M) * 1.0) / (M + 1.0)) * (PP / QQ),
                          0.0),

                Found = false,
                Ktmp = -1,

                % Expand until we cover U
                while (not Found) do
                    % Choose the side with larger next mass (speeds convergence)
                    TakeLeft = (PL >= PR),

                    if TakeLeft then
                        if PL =:= 0.0 then
                            % No left mass; must go right
                            TakeLeft := false
                        end
                    end,

                    if TakeLeft then
                        Cum := Cum + PL,
                        if U =< Cum then
                            Ktmp := L,
                            Found := true
                        else
                            % Move further left: update PL <- P(L-1)
                            L := L - 1,
                            if L >= 0 then
                                % P(k-1) = P(k) * [k / (N-k+1)] * (q/p)
                                PL := PL * ((L + 1.0) / (N - L * 1.0)) * (QQ / PP)
                            else
                                PL := 0.0
                            end
                        end
                    else
                        % Take right
                        Cum := Cum + PR,
                        if U =< Cum then
                            Ktmp := R,
                            Found := true
                        else
                            % Move further right: update PR <- P(R+1)
                            R := R + 1,
                            if R =< N then
                                % P(k+1) = P(k) * [(N-k) / (k+1)] * (p/q)
                                PR := PR * (((N - (R - 1)) * 1.0) / (R * 1.0)) * (PP / QQ)
                            else
                                PR := 0.0
                            end
                        end
                    end
                end,
                K0 = Ktmp
            end
        end,

        % Flip back if we sampled with 1-p
        K = cond(Flip, N - K0, K0)
    end.
binomial_dist_btpe_n(N, P, Num) = [binomial_dist_btpe(N, P) : _ in 1..Num].
%                       (* (binomialf n k) (expt p k) (expt (- 1 p) (- n k))))

/*
% table
% binomial_dist_pdf(N,P,K) = 0 => K < 0.
binomial_dist_pdf(N,P,K) =
  % cond(K<0, 0, binomialf(N,K) * P**K * (1 -P)**(N-K)).
  cond(K<0, 0, binomialf(N,K) * safe_pow(P,K) * safe_pow((1 -P),(N-K))).
*/
/*
------------------------------------------------------------
Binomial distribution PDF
------------------------------------------------------------
PDF(X; N, P) = C(N, X) * P^X * (1-P)^(N-X)
where
  C(N,X) = N! / (X! * (N-X)!)

Parameter constraints:
  N : nonnegative integer
  P : 0.0 =< P =< 1.0
  X : integer, 0 =< X =< N

Throws domain_error for invalid parameters.
------------------------------------------------------------
*/

binomial_dist_pdf(N, P, X) = PDF =>
    if not integer(N) ; N < 0 then
        throw($error(domain_error(number,N),binomial_dist_pdf/3))
    elseif P < 0.0 ; P > 1.0 then
        throw($error(domain_error(number,P),binomial_dist_pdf/3))
    elseif not integer(X) ; X < 0 ; X > N then
        PDF = 0.0
    else
        % Use logarithmic form for numerical stability
        LogPDF = log_binom(N, X) + X*log(P) + (N - X)*log(1.0 - P),
        PDF = exp(LogPDF)
    end.

/*------------------------------------------------------------
Helper: log_binom(N,K) = ln( N choose K )
------------------------------------------------------------*/
log_binom(N, K) = Res =>
    if K < 0 ; K > N then
        Res = -1.0e308
    else
        Res = lgamma(N + 1.0) - lgamma(K + 1.0) - lgamma(N - K + 1.0)
    end.


% table
% binomial_distf_pdf(N,P,K) = binomialf_float(N,K) * P**K * (1 -P)**(N-K).
binomial_distf_pdf(N,P,K) = 
    cond( (K < 0 ; K > N),
         0.0,
         exp(lgamma(N+1.0) - lgamma(K+1.0) - lgamma(N-K+1.0)
         + K*log(P) + (N-K)*log(1.0-P))).


/*
> (map (lambda (v) (list v (* 1.0 (binomial_dist_cdf 10 1/10 v)))) (range 11))
'((0 0.3486784401)
  (1 0.7360989291)
  (2 0.9298091736)
  (3 0.9872048016)
  (4 0.9983650626)
  (5 0.9998530974)
  (6 0.9999908784)
  (7 0.9999996264)
  (8 0.9999999909)
  (9 0.9999999999)
  (10 1.0))

  Picat> X=[binomial_dist_cdf(1/10,10,I)  : I in 0..10]
  X = [0.3486784401,0.7360989291,0.9298091736,0.9872048016,0.9983650626,0.9998530974,0.9999908784,0.9999996264,0.9999999909,0.9999999999,1.0]
*/
% table
binomial_dist_cdf(N,P,K) = sum([binomial_dist_pdf(N,P,I) : I in 0..K]).

% table
% binomial_distf_cdf(N,P,K) = sum([binomial_distf_pdf(N,P,I) : I in 0..K]).
binomial_distf_cdf(N,P,K) = sum([binomial_distf_pdf(N,P,I) : I in 0..K]).

/*
  Reversing the CDF.

> (map (lambda (v) (list v (* 1.0 (binomial_dist_quantile 10 1/10 v)))) (list 0.1 0.5 0.75 0.84 0.9 0.95 0.99 0.99999))
'((0.1 0)
  (0.5 1.0)
  (0.75 2.0)
  (0.84 2.0)
  (0.9 2.0)
  (0.95 3.0)
  (0.99 4.0)
  (0.99999 6.0))

  Picat> X=[I=binomial_dist_quantile(1/10,10,I)  : I in [0.1,0.5,0.75,0.84,0.9,0.95,0.99,0.99999]]
  X = [0.1 = 0,0.5 = 1,0.75 = 2,0.84 = 2,0.9 = 2,0.95 = 3,0.99 = 4,0.99999 = 6]
*/

% table
binomial_dist_quantile(N,P,Q) = Res =>
  Res = 0,
  OK = true,
  foreach(I in 0..N, break(OK == false))
    V := binomial_dist_cdf(N,P,I),
    if V > Q then
      Res := I,
      OK := false
    end,
  end.

% Faster float version
% table
binomial_distf_quantile(N,P,Q) = Res =>
  Res = 0,
  OK = true,
  foreach(I in 0..N, break(OK == false))
    V := binomial_distf_cdf(N,P,I),
    if V > Q then
      Res := I,
      OK := false
    end,
  end.


binomial_dist_mean(N,P) = P * N.

binomial_dist_variance(N,P) = N * (1-P) * P.



/* 
   Negative binomial distribtion

   From Handbook on probability distributions
   page 23
   Expectation: m*(1-p)/p
   """
   The algorithm to simulate a negative binomial distribution NB(m,p) is simply
   to generate m random variables geometrically distributed and to sum them.
   """
*/
negative_binomial_dist(M,P) = sum([geometric_dist(P) : _ in 1..M]).
negative_binomial_dist_n(M,P,N) = [negative_binomial_dist(M,P) : _ in 1..N].

%============================================================
% NEGATIVE BINOMIAL (failures until r successes)
% support X ∈ {0,1,2,...}, with X = K - R
% pmf: P(X=x) = C(x+r-1, r-1) p^r (1-p)^x
%============================================================
% negative_binomial_dist_logpdf(R, P, X) = L =>
%     if R < 1 ; P < 0.0 ; P > 1.0 ; X < 0 then
%         L = log_zero()
%     elseif P == 0.0 then
%         L = log_zero()                        % never get r successes
%     elseif P == 1.0 then
%         L = cond(X == 0, 0.0, log_zero())       % deterministic: 0 failures
%     else
%         % log C(x+r-1, r-1) + r*log p + x*log(1-p)
%         L =  (lgamma(X+R) - lgamma(R) - lgamma(X+1.0))
%            + R*log(P) + X*log(1.0-P)
%     end.


/*
> (map (lambda (v) (list  v (* 1.0 (negative_binomial_pdf 10 1/2 v)))) (range 11))
'((0 0.0009765625)
  (1 0.0048828125)
  (2 0.013427734375)
  (3 0.02685546875)
  (4 0.04364013671875)
  (5 0.06109619140625)
  (6 0.0763702392578125)
  (7 0.0872802734375)
  (8 0.09273529052734375)
  (9 0.09273529052734375)
  (10 0.08809852600097656))

  Picat> X=[I=negative_binomial_dist_pdf(10,1/2,I)  : I in 0..10]
  X = [0 = 0.0009765625,1 = 0.0048828125,2 = 0.013427734375,3 = 0.02685546875,4 = 0.04364013671875,5 = 0.06109619140625,6 = 0.076370239257812,7 = 0.0872802734375,8 = 0.092735290527344,9 = 0.092735290527344,10 = 0.088098526000977]

*/
% table
negative_binomial_dist_pdf(N,P,K) = binomialf( (K + N)- 1, K) * P**N * (1-P)** K.


/*
> (map (lambda (v) (list  v (* 1.0 (negative_binomial_cdf 10 1/2 v)))) (range 31))
'((0 0.0009765625)
  (1 0.005859375)
  (2 0.019287109375)
  (3 0.046142578125)
  (4 0.08978271484375)
  (5 0.15087890625)
  (6 0.2272491455078125)
  (7 0.3145294189453125)
  (8 0.40726470947265625)
  (9 0.5)
  (10 0.5880985260009766)
  (11 0.6681880950927734)
  (12 0.7382664680480957)
  (13 0.7975635528564453)
  (14 0.8462718725204468)
  (15 0.885238528251648)
  (16 0.9156812280416489)
  (17 0.9389609396457672)
  (18 0.956420723348856)
  (19 0.9692858271300793)
  (20 0.9786130273714662)
  (21 0.9852753132581711)
  (22 0.9899691964965314)
  (23 0.9932345065753907)
  (24 0.9954794072546065)
  (25 0.9970059397164732)
  (26 0.9980334134888835)
  (27 0.9987183960038237)
  (28 0.9991709737369092)
  (29 0.9994674901827239)
  (30 0.9996602258725034))

  Picat> X=[I=negative_binomial_dist_cdf(10,1/2,I)  : I in 0..30]
  X = [0 = 0.0009765625,1 = 0.005859375,2 = 0.019287109375,3 = 0.046142578125,4 = 0.08978271484375,5 = 0.15087890625,6 = 0.227249145507812,7 = 0.314529418945312,8 = 0.407264709472656,9 = 0.5,10 = 0.588098526000977,11 = 0.668188095092773,12 = 0.738266468048096,13 = 0.797563552856445,14 = 0.846271872520447,15 = 0.885238528251648,16 = 0.915681228041649,17 = 0.938960939645767,18 = 0.956420723348856,19 = 0.969285827130079,20 = 0.978613027371466,21 = 0.985275313258171,22 = 0.989969196496531,23 = 0.993234506575391,24 = 0.995479407254606,25 = 0.997005939716473,26 = 0.998033413488884,27 = 0.998718396003824,28 = 0.999170973736909,29 = 0.999467490182724,30 = 0.999660225872503]

*/
% table
negative_binomial_dist_cdf(N,P,K) = sum([negative_binomial_dist_pdf(N,P,I) : I in 0..K]).

% table
negative_binomial_dist_quantile(N,P,Q) = Res =>
  Res = 0,
  OK = true,
  foreach(I in 0..N*100, break(OK == false))
    V := negative_binomial_dist_cdf(N,P,I),
    if V > Q then
      Res := I,
      OK := false
    end,
  end.

negative_binomial_dist_mean(N,P) =  (N * (1 - P)) / P.  %  (/ (* n (- 1 p)) p)


/*

  Zero truncated Geometric distribution

  From Handbook on probability distributions, page 21
  Zero truncated Geometric distribution is a Geometric distribution 
  but zero is not a possible value.
   
  It's used for generating a Pascal distribution, see pascal_dist

*/
geometric_zero_truncated1(P,N) = cond( U > P, geometric_zero_truncated1(P,N+1), N) =>
  U = uniform(0,1).


geometric_zero_truncated_dist(P) = geometric_zero_truncated1(P,1).
geometric_zero_truncated_dist_n(P,N) = [geometric_zero_truncated_dist(P)  : _ in 1..N].
 
/*
  Pascal distribution

  https://en.wikipedia.org/wiki/Negative_binomial_distribution

  

*/
pascal_dist(N,P) = sum([geometric_zero_truncated_dist(P) : _ in 1..N]).
pascal_dist_n(N,P,Num) = [pascal_dist(N,P) : _ in 1..Num]. 


%============================================================
% PASCAL / NEGATIVE BINOMIAL (trials until r successes)
% support K ∈ {R, R+1, ...}
% pmf: P(K=k) = C(k-1, r-1) p^r (1-p)^(k-r)
%============================================================
% pascal_dist_logpdf(R, P, K) = L =>
%     if R < 1 ; P < 0.0 ; P > 1.0 ; K < R then
%         L = log_zero()
%     elseif P == 0.0 then
%         L = log_zero()                        % never reach r successes
%     elseif P == 1.0 then
%         L = cond(K == R,0.0,log_zero())       % deterministic at K=R
%     else
%         % log C(k-1, r-1) + r*log p + (k-r)*log(1-p)
%         L =  (lgamma(K) - lgamma(R) - lgamma(K-R+1.0))
%            + R*log(P) + (K-R)*log(1.0-P)
%     end.

% Slower
% pascal_dist2(N,P) = Res =>
%   U = uniform(0,1),
%   Res = pascal_dist_quantile(N,P,U).

pascal_dist_pdf(N,P,X) = Res =>
  Res = cond(X >= N,
             (1-P)**(X-N) * P**N * binomialf(X-1,N-1),
             0).
             

% Pascal CDF: F(k; r, p) = I_p(r, k-r+1),  support k>=r
pascal_dist_cdf(R,P,K) = F =>
    if R < 1 ; P =< 0.0 ; P > 1.0 then
        throw(pascal_dist_cdf_out_of_range)
    elseif K < R then
        F = 0.0
    elseif P == 1.0 then
        % Deterministic at k = R
        F = 1.0
    else
        A = R * 1.0,
        B = (K - R + 1) * 1.0,
        F = incbeta_reg(A, B, P)
    end.

%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%
% Pascal quantile (inverse CDF)
% Parameters:
%   QProb ∈ [0,1]
%   R = number of successes (>=1)
%   P = success probability (0<p<=1)
% Returns smallest integer K such that CDF(K) >= QProb
%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%
pascal_dist_quantile(R, P, QProb) = K =>
    if QProb =< 0.0 then
        K = R
    elseif QProb >= 1.0 then
        K = R + 1000000    % upper bound for "infinity"
    else
        % --- Step 1: find upper bound where CDF >= QProb ---
        Low = R,
        High = R,
        CdfHigh = pascal_dist_cdf(R, P, High),
        % guard against infinite loops if P=1.0 etc
        if P == 1.0 then
            K = R
        else
            OK = true,
            while (CdfHigh < QProb, OK == true)
                High := High * 2,
                CdfHigh := pascal_dist_cdf(R, P, High),
                if (High > 1000000) then
                    OK := false
                end
            end,

            % --- Step 2: binary search between Low..High ---
            while (Low < High)
                Mid = (Low + High) // 2,
                CdfMid = pascal_dist_cdf(R, P, Mid),
                if CdfMid >= QProb then
                    High := Mid
                else
                    Low := Mid + 1
                end
            end,
            K = Low
        end
    end.


% Survival / upper tail: S(k)=Pr(X>k)=1-CDF(k)
pascal_sf(K, R, P) = 1.0 - pascal_dist_cdf(K, R, P).

pascal_dist_mean(N,P) = N/P.

pascal_dist_variance(N,P) = (N * (1 - P)) / P**2.

/*
  Poisson distribution

  The Poisson distribution models the number of events occurring in a fixed
  interval of time or space, given that these events occur independently
  and with a constant mean rate Lambda.

  PMF(k) = e^(-Lambda) * Lambda^k / k!
  CDF(k) = sum_{i=0}^k PMF(i)
  Quantile(u) = smallest integer k such that CDF(k) >= u
  Mean = Lambda
  Variance = Lambda

  Parameters:
  Lambda - rate parameter (average number of events), must be > 0

  Functions:
  poisson_dist_pdf(Lambda, K)
      Probability mass at integer K.

  poisson_dist_cdf(Lambda, K)
      Cumulative probability P(X <= K).

  poisson_dist_quantile(Lambda, U)
      Smallest integer k such that CDF(k) >= U (left-continuous).

  poisson_dist(Lambda)
      Random draw using Knuth’s direct algorithm. Accurate for Lambda < ~30;
      for larger Lambda, an acceptance–rejection or normal approximation
      could be used instead.

  poisson_dist_mean(Lambda)
      Expected value (Lambda).

  poisson_dist_variance(Lambda)
      Variance (Lambda).

  Usage examples:
  println(poisson_dist_pdf(3.0, 2)).
      % Probability of exactly 2 events when Lambda=3.

  println(poisson_dist_cdf(3.0, 4)).
      % Probability of observing <= 4 events.

  println(poisson_dist_quantile(3.0, 0.9)).
      % 90th percentile (number of events).

  X = poisson_dist(3.0),
  printf("Random draw: %w\n", X).

  println(poisson_dist_mean(3.0)).
  println(poisson_dist_variance(3.0)).

  Conceptual notes:
  - The Poisson distribution describes counts of rare, independent events.
  - It is the discrete analog of the Exponential waiting-time model.
  - Approximations:
      * Binomial(n,p) with large n and small p → Poisson(np)
      * For large Lambda, Normal(Lambda, Lambda) is a good approximation.
  - Common applications: number of emails per hour, decay events per second,
    or web hits per minute.

*/ 
% Note: n is the counter, p2 is the acculumated value
% poisson_dist1(Lambda,N,P2) = cond(P >= L, poisson_dist1(Lambda,N+1,P), N) =>
%   L = exp(-Lambda),
%   U = uniform(0,1),
%   P = P2*U.

% poisson_dist(Lambda) = poisson_dist1(Lambda,0,1).

% ----- Random generator -----
% Knuth’s direct method (good for Lambda < ~30)
poisson_dist(Lambda) = K =>
    poisson_dist_validate(Lambda),
    L = exp(-Lambda),
    P = 1.0,
    K = 0,
    while (P > L) do
        K := K + 1,
        P := P * u01()
    end,
    K := K - 1.

poisson_dist_n(Lambda,N)  = [poisson_dist(Lambda) : _ in 1..N].

%
% Special fix for generating samples
% If the generated value is 0, try again
% 
poisson_gt0_dist(Lambda) = K =>
  OK = false,
  K = poisson_dist(Lambda),
  while (OK == false)
    if K > 0 then
      OK := true
    else
      K := poisson_dist(Lambda)
    end
  end.

poisson_gt0_dist_n(Lambda,N) = [poisson_gt0_dist(Lambda) : _ in 1..N].

/* ========= Poisson distribution =========
   Support: X = 0, 1, 2, ...
   Parameter: Lambda > 0

   PMF(x) = e^(-Lambda) * Lambda^x / x!
   CDF(x) = sum_{k=0..x} PMF(k)
   Quantile(u) = smallest integer x such that CDF(x) >= u
   Mean = Lambda
   Variance = Lambda
*/

% ----- validation -----
poisson_dist_validate(Lambda) =>
    if Lambda =< 0.0 then
        throw('poisson_dist: Lambda must be > 0')
    end.

% ----- PDF (PMF) -----
poisson_dist_pdf(Lambda, K) = P =>
    poisson_dist_validate(Lambda),
    if not integer(K) then
        P = 0.0
    elseif K < 0 then
        P = 0.0
    else
        % log form for stability
        L = K * log(Lambda) - Lambda - lgamma(K + 1.0),
        P = exp(L)
    end.

% ----- CDF -----
poisson_dist_cdf(Lambda, K) = F =>
    poisson_dist_validate(Lambda),
    if K < 0 then
        F = 0.0
    else
        % iterative sum using recurrence P(k+1) = P(k)*Lambda/(k+1)
        P = exp(-Lambda),  % P(0)
        S = P,
        I = 0,
        while (I < K) do
            P := P * Lambda / (I + 1.0),
            S := S + P,
            I := I + 1
        end,
        F = S
    end.

% ----- Quantile (left-continuous) -----
poisson_dist_quantile(Lambda, U) = Q =>
    poisson_dist_validate(Lambda),
    if (U < 0.0 ; U > 1.0) then
        throw('poisson_dist_quantile: U must be in [0,1]')
    elseif U =:= 0.0 then
        Q = 0
    else
        Q = _,
        P = exp(-Lambda),  % P(0)
        C = P,
        K = 0,
        Found = false,
        while (not Found) do
            if U =< C then
                Q = K,
                Found := true
            else
                K := K + 1,
                P := P * Lambda / K,
                C := C + P,
                if C >= U then
                    Q = K,
                    Found := true
                end
            end
        end
    end.


% ----- Mean and Variance -----
poisson_dist_mean(Lambda) = M =>
    poisson_dist_validate(Lambda),
    M = Lambda.

poisson_dist_variance(Lambda) = V =>
    poisson_dist_validate(Lambda),
    V = Lambda.



/*
  Hypergeometric distribution.
  This follows the parameter from Mathematica's HypergeometricDistribution[n, n_succ, n_tot]

  Models the probability of obtaining exactly X successes when drawing N items
  *without replacement* from a population of size NTot that contains NSucc successes.

  P(X = x) = (C(NSucc,x) * C(NTot-NSucc, N-x)) / C(NTot,N)

  where C(a,b) = a choose b, and
  x is in { max(0, N - (NTot - NSucc)) .. min(N, NSucc) }.

  Parameters:
  N      - number of draws (sample size), integer, 0 <= N <= NTot
  NSucc  - number of "success" items in the population, integer, 0 <= NSucc <= NTot
  NTot   - total population size, integer, NTot >= 0

  Functions:
  hypergeometric_dist_pdf(N, NSucc, NTot, X)
      Probability that exactly X successes are drawn.

  hypergeometric_dist_cdf(N, NSucc, NTot, K)
      Probability that X <= K.

  hypergeometric_dist_quantile(N, NSucc, NTot, U)
      Smallest x such that CDF(x) >= U (left-continuous).

  hypergeometric_dist_mean(N, NSucc, NTot)
      Expected number of successes: N * (NSucc / NTot).

  hypergeometric_dist_variance(N, NSucc, NTot)
      Variance: N * p * (1 - p) * (NTot - N) / (NTot - 1),
      where p = NSucc / NTot.

  hypergeometric_dist(N, NSucc, NTot)
      Randomly draws N items without replacement and returns
      the number of successes.

  Usage examples:

  println(hypergeometric_dist_pdf(5, 10, 20, 2)).
      % Probability of exactly 2 successes when drawing 5 items
      % from 20, of which 10 are successes.

  println(hypergeometric_dist_cdf(5, 10, 20, 2)).
      % P(X <= 2)

  println(hypergeometric_dist_quantile(5, 10, 20, 0.9)).
      % The smallest x such that CDF(x) >= 0.9

  X = hypergeometric_dist(5, 10, 20),
  println(X).
      % Number of successes in one random draw.

  println(hypergeometric_dist_mean(5, 10, 20)).
  println(hypergeometric_dist_variance(5, 10, 20)).

  Conceptual notes:
  - The Hypergeometric distribution describes sampling without replacement,
    unlike the Binomial distribution, which assumes sampling with replacement.
  - For large NTot relative to N, the Hypergeometric distribution approaches
    the Binomial distribution with p = NSucc / NTot.
  - Common applications: quality control, card draws, and finite-population
    experiments where the sample is not negligible compared to the population.

*/


% ========= Hypergeometric distribution (X = #successes in N draws) =========
% Parameters: N (sample size), NSucc (successes in population), NTot (population size)
% Support: L = max(0, N - (NTot - NSucc)) .. U = min(N, NSucc)

% ----- Helpers -----

% Validate parameters (integers, ranges, and logical constraints)
hypergeometric_dist_validate(N, NSucc, NTot) =>
    if not integer(N) then
        throw('hypergeometric_dist: N must be an integer')
    elseif not integer(NSucc) then
        throw('hypergeometric_dist: NSucc must be an integer')
    elseif not integer(NTot) then
        throw('hypergeometric_dist: NTot must be an integer')
    elseif N < 0 then
        throw('hypergeometric_dist: N must be >= 0')
    elseif NSucc < 0 then
        throw('hypergeometric_dist: NSucc must be >= 0')
    elseif NTot < 0 then
        throw('hypergeometric_dist: NTot must be >= 0')
    elseif NSucc > NTot then
        throw('hypergeometric_dist: NSucc must be <= NTot')
    elseif N > NTot then
        throw('hypergeometric_dist: N must be <= NTot')
    else
        true
    end.

% Support bounds
hypergeometric_dist_support(N, NSucc, NTot) = [L, U] =>
    L = max(0, N - (NTot - NSucc)),
    U = min(N, NSucc).

% log-choose via lgamma
logchoose(A, B) = LG =>
    % assumes 0 =< B =< A
    LG = lgamma(A + 1.0) - lgamma(B + 1.0) - lgamma(A - B + 1.0).

% Safe PMF core using log-choose (expects X within support)
hypergeometric_dist_pmf_core(N, NSucc, NTot, X) = P =>
    % P(X=x) = C(NSucc,x) * C(NTot-NSucc, N-x) / C(NTot,N)
    LN = logchoose(NSucc * 1.0, X * 1.0)
       + logchoose((NTot - NSucc) * 1.0, (N - X) * 1.0)
       - logchoose(NTot * 1.0, N * 1.0),
    P = exp(LN).

% Ratio for recurrence: P(x+1)/P(x)
hypergeometric_dist_ratio(N, NSucc, NTot, X) = R =>
    % R = ((NSucc - X)/(X+1)) * ((N - X)/(NTot - NSucc - (N - X - 1)))
    Num1 = (NSucc - X) * 1.0,
    Num2 = (N - X) * 1.0,
    Den1 = (X + 1) * 1.0,
    Den2 = (NTot - NSucc - (N - X - 1)) * 1.0,
    R = (Num1 / Den1) * (Num2 / Den2).


% ----- Random generator (inverse-CDF via sequential draws without replacement) -----
% Draw N items without replacement from population with NSucc successes.
% Return the count of successes.
hypergeometric_dist(N, NSucc, NTot) = X =>
    hypergeometric_dist_validate(N, NSucc, NTot),
    RemSucc = NSucc,
    RemFail = NTot - NSucc,
    S = 0,
    I = 1,
    while (I =< N) do
        % P(success at this draw) = RemSucc / (RemSucc + RemFail)
        Psucc = (RemSucc * 1.0) / ((RemSucc + RemFail) * 1.0),
        U = u01(),
        if U < Psucc then
            S := S + 1,
            RemSucc := RemSucc - 1
        else
            RemFail := RemFail - 1
        end,
        I := I + 1
    end,
    X = S.

hypergeometric_dist_n(N, NSucc, NTot, Num) = [hypergeometric_dist(N, NSucc, NTot) : _ in 1..Num].

% ----- PDF (PMF) -----
hypergeometric_dist_pdf(N, NSucc, NTot, X) = R =>
    hypergeometric_dist_validate(N, NSucc, NTot),
    % Non-integer X has zero mass for a discrete distribution
    if not integer(X) then
        R = 0.0
    else
        [L, U] = hypergeometric_dist_support(N, NSucc, NTot),
        if X < L ; X > U then
            R = 0.0
        else
            R = hypergeometric_dist_pmf_core(N, NSucc, NTot, X)
        end
    end.

% ----- CDF: F(K) = P(X <= K) -----
hypergeometric_dist_cdf(N, NSucc, NTot, K) = F =>
    hypergeometric_dist_validate(N, NSucc, NTot),
    [L, U] = hypergeometric_dist_support(N, NSucc, NTot),
    if K < L * 1.0 then
        F = 0.0
    elseif K >= U * 1.0 then
        F = 1.0
    else
        Kf = floor(K),
        % Start at x=L using core PMF, then accumulate via stable recurrence
        P0 = hypergeometric_dist_pmf_core(N, NSucc, NTot, L),
        if Kf =:= L then
            F = P0
        else
            S = P0,
            P = P0,
            X = L,
            while (X < Kf) do
                R = hypergeometric_dist_ratio(N, NSucc, NTot, X),
                P := P * R,
                S := S + P,
                X := X + 1
            end,
            F = S
        end
    end.

% ----- Quantile (left-continuous): smallest x with F(x) >= U in [0,1] -----
hypergeometric_dist_quantile(N, NSucc, NTot, U) = Q =>
    hypergeometric_dist_validate(N, NSucc, NTot),
    if (U < 0.0 ; U > 1.0) then
        throw('hypergeometric_dist_quantile: U must be in [0,1]')
    else
        [L, Umax] = hypergeometric_dist_support(N, NSucc, NTot),
        % Accumulate CDF from L upward using recurrence
        P = hypergeometric_dist_pmf_core(N, NSucc, NTot, L),
        C = P,
        if U =< C then
            Q = L
        else
            Q = _,
            X = L,
            Found = false,
            while (X < Umax, not Found) do
                R = hypergeometric_dist_ratio(N, NSucc, NTot, X),
                P := P * R,
                C := C + P,
                X := X + 1,
                if U =< C then
                    Q = X,
                    Found := true
                end
            end,
            if Found == false then
                % Should only happen at U=1.0 due to rounding
                Q = Umax
            end
        end
    end.

% ----- Mean and Variance -----
hypergeometric_dist_mean(N, NSucc, NTot) = M =>
    hypergeometric_dist_validate(N, NSucc, NTot),
    M = (N * 1.0) * (NSucc * 1.0) / (NTot * 1.0).

hypergeometric_dist_variance(N, NSucc, NTot) = V =>
    hypergeometric_dist_validate(N, NSucc, NTot),
    if NTot =:= 0 then
        V = 0.0
    elseif NTot =:= 1 then
        % Only 0 or 1 item in population; if N<=1, variance is 0
        V = 0.0
    else
        P = (NSucc * 1.0) / (NTot * 1.0),
        V = (N * 1.0) * P * (1.0 - P) * ((NTot - N) * 1.0) / ((NTot - 1) * 1.0)
    end.



/*
  Hypergeometric 1 dist 1

  This is an alternative method of Hypergeometric distribution
  to the one above.

  What is the probability that we draw exactly Kk "success" objects
  of the Nn drawn objects of total N objects where there are in total
  K "success" objects

  Kk: number of successes we want to check
  N: total number of objects
  K: total number of success objects
  Nn: number of draws

  Let's call this hypergeometric1_dist/4 instead of hypergeometric_dist/3.

  [To confuse things, in my Gamble probability distribution toolkit 
   (gamble_distributions.rkt), I call this distribution hypergeometric/4,
   and the other approach, shown above hypergeometric2/3.
   But the reason is that I now prefer the one that
   has the same parameters as the Mathematica function HypergeometricDisttribution.
   Sorry about all this.
  ]

*/

% Helper function
hypergeometric1_dist_helper(Kk,N,K,Nn,Count) = Ret =>
  if Nn == 0 ; K <= 0 then
    Ret = Count
  else
    % we have K successes left and N objects left
    P = K/N,
    if flip(P) == true then
       % We drew a success:
       % - decrement the total objects (N)
       % - decrement the number of "success" objects (K)
       % - decrement the number of drawn objects (Nn)
       % - increment the number of successful draws (count)
       Ret = hypergeometric1_dist_helper(Kk, N-1, K-1, Nn-1, Count+1)
    else
       % We drew a failure:
       % - decrement the total objects (N)
       % - decrement the number of drawn objects (n)
       Ret = hypergeometric1_dist_helper(Kk, N-1, K, Nn-1, Count)
    end
  end.


% Random generator: for true/false
hypergeometric1_dist_bool(Kk,N,K,Nn) = cond(Res == K, true, false) => 
  Res = hypergeometric1_dist_helper(Kk,N,K,Nn,0).
hypergeometric1_dist_bool_n(Kk,N,K,Nn, Num) = [hypergeometric1_dist_bool(Kk,N,K,Nn) : _ in 1..Num].

% Random generate the count
% (This was called hypergeometricCount in the Gamble toolkit)
hypergeometric1_dist(Kk,N,K,Nn) = hypergeometric1_dist_helper(Kk,N,K,Nn,0).
hypergeometric1_dist_n(Kk,N,K,Nn,Num) = [hypergeometric1_dist(Kk,N,K,Nn) : _ in 1..Num].

/*
   Hypergeometric 
   From https://en.wikipedia.org/wiki/Hypergeometric_distribution
*/
% Note: It's the same order of arguments as hypergeometric_count (except that k is last)
hypergeometric1_dist_pdf(N, K, Nn, Kk) = binomialf(K,Kk) * binomialf(N-K,Nn-Kk) / binomialf(N,Nn).

hypergeometric1_dist_cdf(N, K, Nn, Kk) =
   sum([hypergeometric1_dist_pdf(N, K, Nn, R) : R in 0..Kk]).

hypergeometric1_dist_quantile(N, K, Nn, Q) = Ret =>
  Ret = 0,
  OK = false,
  foreach(I in 0.. N, break(OK == true))
    if hypergeometric1_dist_cdf(N, K, Nn, I) >= Q then
      Ret := I,
      OK := true
    end
  end.


/*
  Beta-binomial distribution

  The Beta-Binomial distribution is a discrete distribution that models
  the number of successes in N independent Bernoulli trials when the
  probability of success itself is random and follows a Beta(Alpha, Beta)
  distribution.

  It is commonly used in Bayesian statistics as the posterior predictive
  distribution for a Binomial model with a Beta prior.

  Mathematical definition:
    P(X = k) = C(N,k) * B(k + Alpha, N - k + Beta) / B(Alpha, Beta)

  where:
    - C(N,k) = N choose k = N! / (k! * (N - k)!)
    - B(a,b) = Gamma(a) * Gamma(b) / Gamma(a + b)

  Parameters:
    N      - number of Bernoulli trials (integer, N >= 0)
    Alpha  - Beta prior "success" parameter (float, Alpha > 0)
    Beta   - Beta prior "failure" parameter (float, Beta > 0)

  Support:
    k = 0 .. N  (number of observed successes)

  Mean and Variance:
    mean     = N * Alpha / (Alpha + Beta)
    variance = N * Alpha * Beta * (Alpha + Beta + N)
               / ((Alpha + Beta)^2 * (Alpha + Beta + 1))

  Functions:
    beta_binomial_dist_pdf(N, Alpha, Beta, K)
      Probability that X = K successes.

    beta_binomial_dist_cdf(N, Alpha, Beta, K)
      Cumulative probability P(X <= K).

    beta_binomial_dist_quantile(N, Alpha, Beta, U)
      Smallest integer k such that CDF(k) >= U (left-continuous).

    beta_binomial_dist(N, Alpha, Beta)
      Random draw: sample P ~ Beta(Alpha, Beta), then X ~ Binomial(N, P).
      Returns one sample X.

    beta_binomial_dist_mean(N, Alpha, Beta)
      Expected number of successes.

    beta_binomial_dist_variance(N, Alpha, Beta)
        Variance of number of successes.

  Usage examples:

  println(beta_binomial_dist_pdf(10, 2.0, 3.0, 5)).
      % Probability of observing 5 successes in 10 trials
      % given Beta(2,3) prior on success probability.

  println(beta_binomial_dist_cdf(10, 2.0, 3.0, 5)).
      % Cumulative probability P(X <= 5).

  println(beta_binomial_dist_quantile(10, 2.0, 3.0, 0.9)).
      % The smallest k such that P(X <= k) >= 0.9.

  println(beta_binomial_dist_mean(10, 2.0, 3.0)).
      % Mean number of successes.

  println(beta_binomial_dist_variance(10, 2.0, 3.0)).
      % Variance of number of successes.

  X = beta_binomial_dist(10, 2.0, 3.0),
  printf("Random draw: %w\n", X).

  Conceptual notes:
  - The Beta-Binomial accounts for extra uncertainty in the success
    probability compared to the standard Binomial distribution.
  - When Alpha and Beta are large (e.g., Alpha=Beta=100), it approaches
    a regular Binomial with p = Alpha / (Alpha + Beta).
  - When Alpha and Beta are small (<1), the distribution becomes more
    U-shaped, giving higher probability to extreme outcomes (0 or N).
  - It is often used to model overdispersion in Binomial data or to
    perform Bayesian inference for coin-toss experiments.
*/

% ========= Beta–Binomial distribution =========
% X ~ BetaBinomial(N, Alpha, Beta)
% PMF: P(X = K) = C(N,K) * B(K+Alpha, N-K+Beta) / B(Alpha, Beta)
% Support: K = 0..N, with N integer >= 0, Alpha > 0, Beta > 0

% ----- validation -----
beta_binomial_dist_validate(N, Alpha, Beta) =>
    if not integer(N) then
        throw('beta_binomial_dist: N must be an integer')
    elseif N < 0 then
        throw('beta_binomial_dist: N must be >= 0')
    elseif Alpha =< 0.0 then
        throw('beta_binomial_dist: Alpha must be > 0')
    elseif Beta =< 0.0 then
        throw('beta_binomial_dist: Beta must be > 0')
    end.

% ----- helpers -----
lnbeta(A, B) = (lgamma(A) + lgamma(B) - lgamma(A + B)).

% logchoose(N, K) = LG =>
%     % assumes 0 =< K =< N
%     LG = lgamma(N + 1.0) - lgamma(K + 1.0) - lgamma((N - K) + 1.0).

% ratio R = P(K+1)/P(K) for stable recurrences
beta_binomial_dist_ratio(N, Alpha, Beta, K) = R =>
    % R = ((N-K)/(K+1)) * ((K+Alpha)/(N-K-1+Beta))
    Num1 = (N - K) * 1.0,
    Den1 = (K + 1) * 1.0,
    Num2 = (K + Alpha),
    Den2 = ((N - K - 1) + Beta),
    R = (Num1 / Den1) * (Num2 / Den2).


% ----- Random generator (PPL RNG): sample P ~ Beta(Alpha,Beta), then Binomial(N,P) -----
% Requires: gamma_dist(Shape, Scale) and binomial_dist_quantile/3 or binomial_dist/2 available.
beta_binomial_dist(N, Alpha, Beta) = K =>
    beta_binomial_dist_validate(N, Alpha, Beta),
    % Beta via Gamma ratio
    X = gamma_dist(Alpha, 1.0),
    Y = gamma_dist(Beta, 1.0),
    P = X / (X + Y),
    % Binomial draw using existing binomial facility; here we use quantile(U)
    U = u01(),
    K = binomial_dist_quantile(N, P, U).

beta_binomial_dist_n(N, Alpha, Beta, Num) = [beta_binomial_dist(N, Alpha, Beta)  : _ in 1..Num].

% ----- PDF (PMF) -----
beta_binomial_dist_pdf(N, Alpha, Beta, K) = P =>
    beta_binomial_dist_validate(N, Alpha, Beta),
    if not integer(K) then
        P = 0.0
    elseif (K < 0 ; K > N) then
        P = 0.0
    else
        % log-PMF then exp
        LP = logchoose(N * 1.0, K * 1.0)
           + lnbeta(K + Alpha, (N - K) + Beta)
           - lnbeta(Alpha, Beta),
        P = exp(LP)
    end.

% ----- CDF: F(K) = P(X <= K) -----
beta_binomial_dist_cdf(N, Alpha, Beta, K) = F =>
    beta_binomial_dist_validate(N, Alpha, Beta),
    if K < 0.0 then
        F = 0.0
    elseif K >= N * 1.0 then
        F = 1.0
    else
        Kf = floor(K),
        % start from K=0 using closed-form PMF then recur upward
        P0 = exp( lnbeta(Alpha, N + Beta) - lnbeta(Alpha, Beta) ),  % PMF at 0
        if Kf =:= 0 then
            F = P0
        else
            S = P0,
            P = P0,
            I = 0,
            while (I < Kf) do
                R = beta_binomial_dist_ratio(N, Alpha, Beta, I),
                P := P * R,
                S := S + P,
                I := I + 1
            end,
            F = S
        end
    end.

% ----- Quantile (left-continuous): smallest k with F(k) >= U, U in [0,1] -----
beta_binomial_dist_quantile(N, Alpha, Beta, U) = Q =>
    beta_binomial_dist_validate(N, Alpha, Beta),
    if (U < 0.0 ; U > 1.0) then
        throw('beta_binomial_dist_quantile: U must be in [0,1]')
    else
        % accumulate CDF from 0 upward
        P = exp( lnbeta(Alpha, N + Beta) - lnbeta(Alpha, Beta) ),  % PMF at 0
        C = P,
        if U =< C then
            Q = 0
        else
            Q = _,
            K = 0,
            Found = false,
            while (K < N, not Found) do
                R = beta_binomial_dist_ratio(N, Alpha, Beta, K),
                P := P * R,
                C := C + P,
                K := K + 1,
                if U =< C then
                    Q = K,
                    Found := true
                end
            end,
            if not Found then
                % should only happen due to rounding at U=1.0
                Q = N
            end
        end
    end.

% ----- Mean and Variance -----
beta_binomial_dist_mean(N, Alpha, Beta) = M =>
    beta_binomial_dist_validate(N, Alpha, Beta),
    M = (N * 1.0) * (Alpha / (Alpha + Beta)).

beta_binomial_dist_variance(N, Alpha, Beta) = V =>
    beta_binomial_dist_validate(N, Alpha, Beta),
    S = Alpha + Beta,
    V = (N * 1.0) * (Alpha * Beta) * (S + N) / (S * S * (S + 1.0)).


/*
  Exponentital distribution

  From Handbook on probability distributions
  page 53
  """
  Despite the quantile function is F−1(u) = −1/lambda * log(1-u),
  generally the exponential distribution E(λ) is generated by applying
   −1/lambda * log(U) on a uniform variate U.
  """

*/
exponential_dist(Lambda) = -1*log(U)/Lambda => 
    U = uniform(0,1).

exponential_dist_n(Lambda,N) = [exponential_dist(Lambda) : _ in 1..N].

% From https://en.wikipedia.org/wiki/Inverse_transform_sampling
% exponential_dist2(Lambda) = (-1/Lambda)*log(1-U) =>
%    U = uniform(0,1).



% %%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%
% % Exponential log-pdf (rate λ)
% %%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%
% exponential_dist_logpdf(Lambda, X) = L =>
%     if Lambda =< 0.0 ; X < 0.0 then
%         L = log_zero()
%     else
%         L = log(Lambda) - Lambda*X
%     end.

%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%
% exponential_dist_pdf(Lambda, X)
%
% PDF of Exponential(λ):
%   f(x;λ) = λ * exp(-λx)  for x ≥ 0
%          = 0             for x < 0
%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%
exponential_dist_pdf(Lambda, X) = P =>
    (Lambda =< 0.0) -> throw($errir('exponential_dist_pdf: Lambda must be > 0'))
    ;
    P = cond(X < 0.0, 0.0, Lambda * exp(-Lambda * X)).
    

%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%
% exponential_dist_cdf(Lambda, X)
%
% CDF of Exponential(λ):
%   F(x;λ) = 1 - exp(-λx)  for x ≥ 0
%          = 0             for x < 0
%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%
exponential_dist_cdf(Lambda, X) = CDF =>
    (Lambda =< 0.0) -> throw($error('exponential_dist_cdf: Lambda must be > 0'))
    ;
    CDF = cond(X < 0.0, 0.0, 1.0 - exp(-Lambda * X)).


%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%
% exponential_dist_quantile(Lambda, P)
%
% Quantile (inverse CDF) of Exponential(λ):
%   Q(p;λ) = -ln(1 - p) / λ   for 0 ≤ p < 1
%   Q(1;λ) = inf  (approximated as very large number)
%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%
exponential_dist_quantile(Lambda, P) = X =>
    (Lambda =< 0.0) -> throw($error('exponential_dist_quantile: Lambda must be > 0'))
    ;
    (P < 0.0 ; P > 1.0) -> throw($error('exponential_dist_quantile: P must be in [0,1]'))
    ;
    X = cond(P >= 1.0, 1.0e300, -log(1.0 - P) / Lambda).



exponential_dist_mean(Lambda) = M =>
    (Lambda =< 0.0) -> throw($error('exponential_dist_mean: Lambda must be > 0'))
    ;
    M = 1.0 / Lambda.

exponential_dist_variance(Lambda) = V =>
    (Lambda =< 0.0) -> throw($error('exponential_dist_variance: Lambda must be > 0'))
    ;
    V = 1.0 / (Lambda * Lambda).

%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%
% Survival function S(x) = P(X > x)
%   S(x;λ) = exp(-λx)   for x ≥ 0
%          = 1          for x < 0
%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%
exponential_dist_sf(Lambda, X) = S =>
    (Lambda =< 0.0) -> throw($error('exponential_dist_sf: Lambda must be > 0'))
    ;
    S = cond(X < 0.0, 1.0, exp(-Lambda * X)).

%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%
% Hazard h(x) = f(x)/S(x)
% For Exponential: h(x) = λ   (constant) for x ≥ 0
% Undefined for x < 0  -> we throw.
%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%
exponential_dist_hazard(Lambda, X) = H =>
    (Lambda =< 0.0) -> throw($error('exponential_dist_hazard: Lambda must be > 0'))
    ;
    (X < 0.0) -> throw($error('exponential_dist_hazard: X must be >= 0'))
    ;
    H = Lambda.

%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%
% (Optional) Cumulative hazard H(x) = -ln S(x) = λ x  (for x ≥ 0)
%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%
exponential_dist_cumhazard(Lambda, X) = CH =>
    (Lambda =< 0.0) -> throw($error('exponential_dist_cumhazard: Lambda must be > 0'))
    ;
    (X < 0.0) -> throw($error('exponential_dist_cumhazard: X must be >= 0'))
    ;
    CH = Lambda * X.


/*
  Beta distribution

Continuous probability distribution on the interval (0,1).

PDF(x) = x^(Alpha-1) * (1-x)^(Beta-1) / B(Alpha, Beta)
CDF(x) = I_x(Alpha, Beta)
where B(a,b) = Gamma(a)*Gamma(b)/Gamma(a+b)
and I_x(a,b) is the regularized incomplete beta function.

Parameters:
  Alpha  - shape parameter > 0
  Beta   - shape parameter > 0
Support:
  0 < x < 1

Mean     = Alpha / (Alpha + Beta)
Variance = (Alpha * Beta) / ((Alpha + Beta)^2 * (Alpha + Beta + 1))

Functions:
  beta_dist_pdf(Alpha, Beta, X)
      Probability density at X.
  beta_dist_cdf(Alpha, Beta, X)
      Cumulative probability P(X' <= X).
  beta_dist_quantile(Alpha, Beta, U)
      Value x such that CDF(x) = U.
  beta_dist_mean(Alpha, Beta)
      Expected value.
  beta_dist_variance(Alpha, Beta)
      Variance.
  beta_dist(Alpha, Beta)
      Random sample using Gamma ratio method.

Usage examples:

  println(beta_dist_pdf(2.0, 3.0, 0.5)).
      % Density at x=0.5 for Beta(2,3).

  println(beta_dist_cdf(2.0, 3.0, 0.5)).
      % Cumulative probability at x=0.5.

  println(beta_dist_quantile(2.0, 3.0, 0.9)).
      % 90th percentile of Beta(2,3).

  println(beta_dist_mean(2.0, 3.0)).
  println(beta_dist_variance(2.0, 3.0)).

  X = beta_dist(2.0, 3.0),
  printf("Random draw: %w\n", X).

Conceptual notes:
  - The Beta distribution is the conjugate prior for the Bernoulli
    and Binomial models, making it a cornerstone of Bayesian inference.
  - Small Alpha and Beta (<1) produce U-shaped densities with higher
    probability near 0 and 1.
  - Alpha = Beta = 1 gives the Uniform(0,1) distribution.
  - Large Alpha and Beta produce sharply peaked densities near the mean.
*/

% ========= Beta distribution =========
% Support: 0 < X < 1
% Parameters: Alpha > 0, Beta > 0
%
% PDF(x) = x^(Alpha-1) * (1-x)^(Beta-1) / B(Alpha, Beta)
% CDF(x) = I_x(Alpha, Beta)
% Mean = Alpha / (Alpha + Beta)
% Variance = (Alpha * Beta) / ((Alpha + Beta)^2 * (Alpha + Beta + 1))

% ----- validation -----
beta_dist_validate(Alpha, Beta) =>
    if Alpha =< 0.0 then
        throw('beta_dist: Alpha must be > 0')
    elseif Beta =< 0.0 then
        throw('beta_dist: Beta must be > 0')
    end.

% ----- Random generator -----
% Beta(alpha,beta) = X / (X + Y), with X~Gamma(alpha,1), Y~Gamma(beta,1)
beta_dist(Alpha, Beta) = R =>
    beta_dist_validate(Alpha, Beta),
    X = gamma_dist(Alpha, 1.0),
    Y = gamma_dist(Beta, 1.0),
    R = X / (X + Y).

beta_dist_n(Alpha, Beta, N) = [beta_dist(Alpha, Beta) : _ in 1..N].

% ----- PDF -----
beta_dist_pdf(Alpha, Beta, X) = P =>
    beta_dist_validate(Alpha, Beta),
    if (X =< 0.0 ; X >= 1.0) then
        P = 0.0
    else
        L = (Alpha - 1.0) * log(X) + (Beta - 1.0) * log(1.0 - X)
            - (lgamma(Alpha) + lgamma(Beta) - lgamma(Alpha + Beta)),
        P = exp(L)
    end.
% Regularized incomplete beta I_x(Alpha,Beta)
beta_reg_incomplete(Alpha, Beta, X) = I =>
    I = _,
    if X =< 0.0 then
        I = 0.0
    elseif X >= 1.0 then
        I = 1.0
    elseif X < (Alpha + 1.0) / (Alpha + Beta + 2.0) then
        I = beta_reg_incomplete_cf(Alpha, Beta, X)
    else
        % symmetry for better convergence
        I = 1.0 - beta_reg_incomplete_cf(Beta, Alpha, 1.0 - X)
    end.

% Helper: I_x(a,b) using continued fraction (Lentz). Returns bt*cf/a.
beta_reg_incomplete_cf(A, B, X) = R =>
    EPS   = 1.0e-14,
    FPMIN = 1.0e-30,
    MAXIT = 200,

    % Front factor bt = exp(lnGamma(a+b)-lnGamma(a)-lnGamma(b) + a ln x + b ln(1-x))
    BT = exp(lgamma(A + B) - lgamma(A) - lgamma(B) + A*log(X) + B*log(1.0 - X)),

    % Lentz initialization
    C = 1.0,
    D = 1.0 - (A + B) * X / (A + 1.0),
    if abs(D) < FPMIN then D := FPMIN end,
    D := 1.0 / D,
    H = D,

    M = 1,
    while (M =< MAXIT) do
        M2 = 2 * M,

        % a1 term
        AA = M * (B - M) * X / ((A + M2 - 1.0) * (A + M2)),
        D := 1.0 + AA * D,       if abs(D) < FPMIN then D := FPMIN end,
        C := 1.0 + AA / C,       if abs(C) < FPMIN then C := FPMIN end,
        D := 1.0 / D,
        H := H * D * C,
        % a2 term
        AA := - (A + M) * (A + B + M) * X / ((A + M2) * (A + M2 + 1.0)),
        D := 1.0 + AA * D,       if abs(D) < FPMIN then D := FPMIN end,
        C := 1.0 + AA / C,       if abs(C) < FPMIN then C := FPMIN end,
        D := 1.0 / D,
        DELTA = D * C,
        H := H * DELTA,

        if abs(DELTA - 1.0) < EPS then
            M := MAXIT + 1   % break
        else
            M := M + 1
        end
    end,

    R = BT * H / A.


% ----- CDF -----
beta_dist_cdf(Alpha, Beta, X) = F =>
    beta_dist_validate(Alpha, Beta),
    F = beta_reg_incomplete(Alpha, Beta, X).

% ----- Quantile -----
% Inverse of regularized incomplete beta I_x(Alpha,Beta) = U
beta_dist_quantile(Alpha, Beta, U) = X =>
    beta_dist_validate(Alpha, Beta),
    X = _,
    if (U < 0.0 ; U > 1.0) then
        throw('beta_dist_quantile: U must be in [0,1]')
    elseif U =:= 0.0 then
        X = 0.0
    elseif U =:= 1.0 then
        X = 1.0
    else
        % initial guess via mean
        X0 = Alpha / (Alpha + Beta),
        X := X0,
        EPS = 1.0e-12,
        MAXIT = 50,
        I = 1,
        while (I =< MAXIT) do
            F = beta_dist_cdf(Alpha, Beta, X),
            Pdf = beta_dist_pdf(Alpha, Beta, X),
            if Pdf =:= 0.0 then
                I := MAXIT + 1 % hakank: Is this correct? What is the X here?
            else
                Step = (F - U) / Pdf,
                X := X - Step,
                if X =< 0.0 then X := 1.0e-12 end,
                if X >= 1.0 then X := 1.0 - 1.0e-12 end,
                if abs(Step) < EPS then
                    I := MAXIT + 1
                else
                    I := I + 1
                end
            end
        end
    end.


% ----- Mean and Variance -----
beta_dist_mean(Alpha, Beta) = M =>
    beta_dist_validate(Alpha, Beta),
    M = Alpha / (Alpha + Beta).

beta_dist_variance(Alpha, Beta) = V =>
    beta_dist_validate(Alpha, Beta),
    S = Alpha + Beta,
    V = (Alpha * Beta) / (S * S * (S + 1.0)).



/*
  Gamma distribution

  Note: This is the shape-scale version with the parameters Alpha, Theta

  To get the rate version, use 
  gamma_dist_pdf(Alpha, 1/Theta, X) 
  gamma_dist_cdf(Alpha, 1/Theta, X) 

  The original ("fancy") version is reliable for large values. 
  This is tried first, but if throws an error, the binary search
  version is then tried (fairly reliable even for larger values).

*/

% Gamma(shape=Alpha, scale=Theta) random generator (Marsaglia–Tsang)
gamma_dist(Alpha, Theta) = X =>
    (Alpha =< 0.0 ; Theta =< 0.0) -> throw($error('gamma_dist: Alpha,Theta must be > 0'))
    ;
    if Alpha == 1.0 then
        % Exponential(Theta)
        U = u01(),
        X = -Theta * log(U)
    elseif Alpha > 1.0 then
        % Marsaglia–Tsang for k>=1 (rate=1), then scale by Theta
        D = Alpha - 1.0/3.0,
        C = 1.0 / sqrt(9.0 * D),
        Accepted = false,
        % V = 0.00001,
        V = 0,
        while (not Accepted)
            Z = std_normal_rand(),
            Vt = 1.0 + C * Z,
            if Vt > 0.0 then
                V := Vt * Vt * Vt,             % (1 + c z)^3
                U0 = uniform(0,1), % u01(),
                % Squeeze test
                if U0 < 1.0 - 0.0331 * Z*Z*Z*Z then
                    Accepted := true
                else
                    % Exact acceptance (log form) with V guarded
                    VV = max(1.0e-300, V),
                    if log(U0) < 0.5*Z*Z + D*(1.0 - V + log(VV)) then
                        Accepted := true
                    end
                end
            end
        end,
        X = Theta * (D * V)
    else
        % 0<Alpha<1 : "boost" trick
        Y = gamma_dist(Alpha + 1.0, Theta),
        U = u01(),
        % X = Y * pow(U, 1.0 / Alpha)
        X = Y * safe_pow(U, 1.0 / Alpha)
    end.

gamma_dist_n(Alpha, Theta, N) = [gamma_dist(Alpha, Theta) : _ in 1..N].

%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%
% Gamma (shape–scale) to match Racket/Gamble:
%   pdf:  f(x;α,θ) = x^(α-1) * exp(-x/θ) / (Γ(α) * θ^α),  x>=0
%   cdf:  F(x;α,θ) = P(α, x/θ), where P is regularized lower inc. gamma
%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%
gamma_dist_pdf(Alpha, Theta, X) = P =>
    (Alpha =< 0.0 ; Theta =< 0.0) -> throw($error('gamma_dist_pdf: Alpha,Theta must be > 0'))
    ;
    if X < 0.0 then
        P = 0.0
    else
        P = safe_pow(X, Alpha - 1.0) * exp(-X / Theta) /
            (gamma_func(Alpha) * safe_pow(Theta, Alpha))
    end.

gamma_dist_cdf(Alpha, Theta, X) = CDF =>
    (Alpha =< 0.0 ; Theta =< 0.0) -> throw($error('gamma_dist_cdf: Alpha,Theta must be > 0'))
    ;
    if X =< 0.0 then
        CDF = 0.0
    else
        % regularized lower incomplete gamma at A=Alpha, X' = X/Theta
        CDF = inc_gamma_lower_reg(Alpha, X / Theta)
    end.


%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%
% gamma_dist_quantile(Alpha, Theta, P)
% ------------------------------------------------------------
% Handles integer or float Alpha=1 safely.
%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%

gamma_dist_quantile(Alpha, Theta, P) = X =>
  % Check first with the fancy version
  % which will fail for larger values
  catch(X1 = gamma_dist_quantile1(Alpha, Theta, P),Error,true),
  if nonvar(Error) then
    % If not, try the binary search approach.
    % println(error=Error),
    X = binary_search_quantile($gamma_dist_cdf(Alpha, Theta),
                                P, 
                                1.0e-13, % Tolerance
                                0,       % Lower bound
                                1.0e100  % Upper bound
                                )
  else
    X = X1
  end.

% THIS IS NOT RELIABLE for larger values.
% Use the binary_search_quantile instead:
gamma_dist_quantile1(Alpha, Theta, P) = X =>
    (Alpha =< 0.0 ; Theta =< 0.0) -> throw($error('gamma_dist_quantile: Alpha,Theta must be > 0'))
    ;
    (P < 0.0 ; P > 1.0) -> throw($error('gamma_dist_quantile: P must be in [0,1]'))
    ;
    TolAlpha = 1.0e-12,
    if P == 0.0 then
        X = 0.0
    elseif P == 1.0 then
        X = 1.0e300
    elseif abs(Alpha - 1.0) < TolAlpha then
        % Exponential(Theta)
        U = max(1.0e-300, 1.0 - P),   % guard against log(0)
        X = Theta * (-log(U))
    else
        % ---- Initial guess ----
        Z = cond((P =< 0.0 ; P >= 1.0), 0.0, normal_dist_quantile(0.0, 1.0, P)),
        X0 = cond(Alpha >= 1.0,
                  Theta * Alpha * safe_pow(1.0 - 1.0/(9.0*Alpha) + Z/(3.0*sqrt(Alpha)), 3.0),
                  Theta * safe_pow(-log(1.0 - P), 1.0/Alpha)),
        Xstart = cond(X0 =< 0.0, Theta, X0),

        % ---- Bracket [Lo,Hi] with increasing CDF ----
        Lo = 0.0,
        Hi = Xstart,
        CdfHi = gamma_dist_cdf(Alpha, Theta, Hi),
        while (CdfHi < P)
            Hi := Hi * 2.0,
            CdfHi := gamma_dist_cdf(Alpha, Theta, Hi)
        end,

        % ---- Safeguarded Newton iterations ----
        Tol = 1.0e-12,
        Xn = Xstart,
        CdfX = gamma_dist_cdf(Alpha, Theta, Xn),
        PdfX = max(1.0e-300, gamma_dist_pdf(Alpha, Theta, Xn)),
        Done = false,
        Iters = 0,
        Xres = Xn,
        while (not Done, Iters < 60)
            if CdfX > P then
                Hi := Xn
            else
                Lo := Xn
            end,

            Step = (CdfX - P) / PdfX,
            Xcand = Xn - Step,

            Out = (Xcand =< Lo ; Xcand >= Hi ; Xcand != Xcand ;
                   Xcand =:= 1.0/0.0 ; Xcand =:= -1.0/0.0),
            Xnext = cond(Out, 0.5*(Lo + Hi), Xcand),

            if abs(Xnext - Xn) =< Tol * (1.0 + abs(Xn)) then
                Xres := Xnext,
                Done := true
            else
                Xn := Xnext,
                CdfX := gamma_dist_cdf(Alpha, Theta, Xn),
                PdfX := max(1.0e-300, gamma_dist_pdf(Alpha, Theta, Xn)),
                Iters := Iters + 1
            end
        end,
        X = Xres
    end.

%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%
% Gamma(shape–scale): mean, variance, mode, survival, hazard, cumhazard
% Requires your gamma_dist_pdf/3 and gamma_dist_cdf/3 (shape–scale).
%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%

gamma_dist_mean(Alpha, Theta) = M =>
    (Alpha =< 0.0 ; Theta =< 0.0) -> throw($error('gamma_dist_mean: Alpha,Theta must be > 0'))
    ;
    M = Alpha * Theta.

gamma_dist_variance(Alpha, Theta) = V =>
    (Alpha =< 0.0 ; Theta =< 0.0) -> throw($error('gamma_dist_variance: Alpha,Theta must be > 0'))
    ;
    V = Alpha * Theta * Theta.

% Mode:
%   if Alpha > 1: (Alpha - 1)*Theta
%   if 0 < Alpha <= 1: 0  (mode at boundary)
gamma_dist_mode(Alpha, Theta) = Mode =>
    (Alpha =< 0.0 ; Theta =< 0.0) -> throw($error('gamma_dist_mode: Alpha,Theta must be > 0'))
    ;
    Mode = cond(Alpha > 1.0, (Alpha - 1.0)*Theta, 0.0).

% Survival S(x) = 1 - F(x)  for x >= 0;  S(x)=1 for x<0
gamma_dist_sf(Alpha, Theta, X) = S =>
    (Alpha =< 0.0 ; Theta =< 0.0) -> throw($error('gamma_dist_sf: Alpha,Theta must be > 0'))
    ;
    S = cond(X < 0.0, 1.0, 1.0 - gamma_dist_cdf(Alpha, Theta, X)).

% Hazard h(x) = f(x)/S(x) for x >= 0; undefined for x<0 -> throw.
% Guard S(x) to avoid division by ~0 in extreme upper tail.
gamma_dist_hazard(Alpha, Theta, X) = H =>
    (Alpha =< 0.0 ; Theta =< 0.0) -> throw($error('gamma_dist_hazard: Alpha,Theta must be > 0'))
    ;
    (X < 0.0) -> throw($error('gamma_dist_hazard: X must be >= 0'))
    ;
    Fx = gamma_dist_cdf(Alpha, Theta, X),
    Sx = max(1.0e-300, 1.0 - Fx),
    H = gamma_dist_pdf(Alpha, Theta, X) / Sx.

% Cumulative hazard H(x) = -ln S(x) for x >= 0; undefined for x<0 -> throw.
gamma_dist_cumhazard(Alpha, Theta, X) = CH =>
    (Alpha =< 0.0 ; Theta =< 0.0) -> throw($error('gamma_dist_cumhazard: Alpha,Theta must be > 0'))
    ;
    (X < 0.0) -> throw($error('gamma_dist_cumhazard: X must be >= 0'))
    ;
    Sx = max(1.0e-300, 1.0 - gamma_dist_cdf(Alpha, Theta, X)),
    CH = -log(Sx).


/*
  Multinomial distribution

  It does not have a quantile.

*/

%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%
% multinomial_dist(N, Ps) = Counts
% ------------------------------------------------------------
% Draws Counts ~ Multinomial(N, Ps), where Ps is a list of probabilities.
% Ps need not be normalized; zeros are allowed; all must be >= 0.
% Returns a list Counts of same length as Ps, summing to N.
%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%
multinomial_dist(N, Ps) = Counts =>
    S = sum(Ps),
    if not integer(N) ; N < 0 then
      throw($error('multinomial_dist: N must be a nonnegative integer'))
    elseif Ps == [] then
      throw($error('multinomial_dist: Ps must be non-empty'))
    elseif not all_nonneg(Ps) then
      throw($error('multinomial_dist: all probabilities must be >= 0'))
    elseif S =:= 0.0 then
      throw($error('multinomial_dist: sum(Ps) must be > 0'))
    end,
    % Normalize
    Qs = [P / S : P in Ps],
    K  = length(Qs),
    % Build cumulative CDF
    Cum = new_list(K, 0.0),
    Acc = 0.0,
    foreach(I in 1..K)
        Acc := Acc + Qs[I],
        Cum[I] := Acc
    end,
    % ensure last bin closes to 1 (guard rounding)
    Cum[K] := 1.0,

    % Draw N categorical indices and tally
    Cnts = new_list(K, 0),
    foreach(_ in 1..N)
        U = u01(),
        J = 1,
        Found = false,
        while (J =< K, not Found)
            if U =< Cum[J] then
                Cnts[J] := Cnts[J] + 1,
                Found := true
            else
                J := J + 1
            end
        end
    end,
    Counts = Cnts.

multinomial_dist_n(N, Ps, Num) = [multinomial_dist(N, Ps) : _ in 1..Num].

%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%
% multinomial_dist_pdf(N, Ps, Ks)
% ------------------------------------------------------------
% PDF (PMF) for Multinomial(N, Ps):
%   f(Ks; N, Ps) = N! / (∏ k_i!) * ∏ p_i^k_i
% Ps may be unnormalized; we normalize internally.
% Ks must sum to N, same length as Ps.
%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%
multinomial_dist_pdf(N, Ps, Ks) = P =>
    if not integer(N) ; N < 0 then
        throw('multinomial_dist_pdf: N must be a nonnegative integer')
    elseif Ps == [] ; length(Ps) != length(Ks) then
        throw('multinomial_dist_pdf: Ps and Ks must have same non-empty length')
    elseif sum(Ks) != N then
        throw('multinomial_dist_pdf: sum(Ks) must equal N')
    elseif not all_nonneg(Ps) then
        throw('multinomial_dist_pdf: all Ps must be >= 0')
    else
        S = sum(Ps),
        if S =:= 0.0 then
            throw('multinomial_dist_pdf: sum(Ps) must be > 0')
        else
            % Normalize Ps
            Qs = [P0 / S : P0 in Ps],
            K = length(Qs),

            % If any Qs[I]=0 with Ks[I]>0 => probability 0
            HasZeroProb = false,
            foreach(I in 1..K)
                if Qs[I] =:= 0.0, Ks[I] > 0 then
                    HasZeroProb := true
                end
            end,

            if HasZeroProb then
                P = 0.0
            else
                % log P = lgamma(N+1) - sum lgamma(k_i+1) + sum k_i*log(q_i)
                LogP1 = lgamma(N + 1.0) - sum([lgamma(Ks[I] + 1.0) : I in 1..K]),
                SumKLogQ = 0.0,
                foreach(I in 1..K)
                    if Ks[I] =:= 0 then
                        SumKLogQ := SumKLogQ + 0.0
                    else
                        SumKLogQ := SumKLogQ + Ks[I] * log(Qs[I])
                    end
                end,
                P = exp(LogP1 + SumKLogQ)
            end
        end
    end.

% The CDF is the same as the PDF
multinomial_dist_cdf(N, Ps, Ks) = multinomial_dist_pdf(N, Ps, Ks).

multinomial_dist_mean(N, Ps) = [N * P / sum(Ps) : P in Ps].


multinomial_dist_varcov(N, Ps) = (Vars, Covs) =>
    S = sum(Ps),
    Qs = [P / S : P in Ps],
    Vars = [N * Q * (1.0 - Q) : Q in Qs],
    Covs = new_array(length(Qs), length(Qs)),
    foreach(I in 1..length(Qs))
        foreach(J in 1..length(Qs))
            if I == J then
                Covs[I,J] := Vars[I]
            else
                Covs[I,J] := -N * Qs[I] * Qs[J]
            end
        end
    end.

/*
------------------------------------------------------------
Dirichlet Distribution
------------------------------------------------------------
Name:
    dirichlet_dist(Alpha)

Description:
    The Dirichlet distribution is a multivariate generalization
    of the Beta distribution. It is defined over the probability
    simplex:
        x_i >= 0,  sum_i x_i = 1
    with concentration parameters:
        Alpha = [A1, A2, ..., Ak], where each A_i > 0.

PDF:
    PDF(X; Alpha) =
        (1 / B(Alpha)) * product_i( X_i^(A_i - 1) )

    where
        B(Alpha) = product_i( Gamma(A_i) ) / Gamma( sum_i A_i )

Support:
    X_i >= 0 for all i, and sum_i X_i = 1.

Mean:
    mean_i = A_i / A0
    where A0 = sum_i Alpha_i

Variance (per component):
    var_i = (A_i * (A0 - A_i)) / (A0^2 * (A0 + 1))

Covariance (i != j):
    cov_ij = - (A_i * A_j) / (A0^2 * (A0 + 1))

Random generation:
    Sample G_i ~ Gamma(A_i, 1.0)
    Return X_i = G_i / sum(G)

CDF:
    No closed form for general k > 2.
    For k = 2, it reduces to the Beta distribution.

Quantile:
    No closed form for general k > 2.
    For k = 2, use the Beta quantile function.

Parameter constraints:
    * Alpha must be a list of positive numbers.
    * Length of X and Alpha must be equal.
    * Each X_i >= 0, sum(X) = 1.
    * Throws explicit domain_error exceptions on invalid input.

Function summary:
    dirichlet_dist(Alpha)
        -> returns a random sample vector X on the simplex

    dirichlet_dist_pdf(Alpha, Xs)
        -> returns the probability density at point Xs

    dirichlet_dist_cdf(Alpha, Xs)
        -> not implemented (except via Beta for k=2)

    dirichlet_dist_quantile(Alpha, P)
        -> not implemented (except via Beta for k=2)

    dirichlet_dist_mean(Alpha)
        -> returns list of component means

    dirichlet_dist_variance(Alpha)
        -> returns list of component variances

Notes:
    * Implemented using float-first arithmetic.
    * Numerically stable (log-space PDF computation).
    * Throws errors instead of returning NaN.
    * Follows Picat syntax conventions (uppercase vars,
      if/then/elseif/else/end, no break, no ternary).

------------------------------------------------------------
*/ 

/*
------------------------------------------------------------
Dirichlet distribution implementation for Picat PPL
Functions:
    dirichlet_dist(Alpha)             % random generation
    dirichlet_dist_pdf(Alpha,X)
    dirichlet_dist_cdf(Alpha,X)      % not implemented (except k=2 optional)
    dirichlet_dist_quantile(Alpha,P) % not implemented
    dirichlet_dist_mean(Alpha)
    dirichlet_dist_variance(Alpha)
------------------------------------------------------------
*/

/*------------------------------------------------------------
  Validation
------------------------------------------------------------*/
ensure_dirichlet_params(Alpha) =>
    if not list(Alpha) then
        throw($error('domain_error(list,Alpha),dirichlet_dist/1'))
    else
        foreach (A in Alpha)
            if A =< 0.0 then
                throw($error('domain_error(number,Alpha),dirichlet_dist/1'))
            end
        end
    end.

/*------------------------------------------------------------
  Random generation:
  X_i = Gamma(A_i,1.0) / sum(Gamma(A_i,1.0))
------------------------------------------------------------*/
dirichlet_dist(Alpha) = Xs =>
    ensure_dirichlet_params(Alpha),
    Gs = [gamma_dist(A,1.0) : A in Alpha],
    SumG = sum(Gs),
    if SumG =:= 0.0 then
        throw($error('domain_error(number,Alpha),dirichlet_dist/1'))
    else
        Xs = [G / SumG : G in Gs]
    end.
dirichlet_dist_n(Alpha,N) = [dirichlet_dist(Alpha) : _ in 1..N].

/*------------------------------------------------------------
  PDF
------------------------------------------------------------*/
dirichlet_dist_pdf(Alpha,Xs) = PDF =>
    ensure_dirichlet_params(Alpha),
    if not list(Xs) then
        throw($error('domain_error(list,Xs),dirichlet_dist_pdf/2)'))
    elseif length(Xs) != length(Alpha) then
        throw($error('domain_error(length,Xs),dirichlet_dist_pdf/2'))
    elseif sum(Xs) > 1.0000001 ; sum(Xs) < 0.9999999 then
        throw($error('domain_error(number,sum(Xs)),dirichlet_dist_pdf/2)'))
    else
        % If any x_i <= 0, the pdf = 0.0
        % if member(X,Xs), X =< 0.0 then
        if [Check : X in Xs, Check = cond(X =< 0.0,1,0) ].sum > 0 then
            PDF = 0.0
        else
            SumA = sum(Alpha),
            LogNumer = sum([(A - 1.0) * log(X) : {A, X} in zip(Alpha, Xs)]),
            LogDenom = sum([lgamma(A) : A in Alpha]) - lgamma(SumA),
            PDF = exp(LogNumer - LogDenom)
        end
    end.

/*------------------------------------------------------------
  CDF (no closed form for k > 2)
------------------------------------------------------------*/
dirichlet_dist_cdf(_Alpha,_Xs) =>
    throw($error(not_implemented, dirichlet_dist_cdf/2)).

/*------------------------------------------------------------
  Quantile (no closed form)
------------------------------------------------------------*/
dirichlet_dist_quantile(_Alpha,_P) =>
    throw($error('not_implemented, dirichlet_dist_quantile/2')).

/*------------------------------------------------------------
  Mean
------------------------------------------------------------*/
dirichlet_dist_mean(Alpha) = Mean =>
    ensure_dirichlet_params(Alpha),
    SumA = sum(Alpha),
    Mean = [A / SumA : A in Alpha].

/*------------------------------------------------------------
  Variance (per component)
------------------------------------------------------------*/
dirichlet_dist_variance(Alpha) = Vars =>
    ensure_dirichlet_params(Alpha),
    SumA = sum(Alpha),
    Vars = [(A * (SumA - A)) / (SumA * SumA * (SumA + 1.0)) : A in Alpha].


/*
  Laplace distribution
  --------------------
  The Laplace (double exponential) distribution is a continuous
  probability distribution characterized by a sharp peak at its
  location parameter Mu and heavier tails than the Normal distribution.
  
  Parameters:
    Mu : location parameter (float)
    B  : scale parameter (float, must be > 0.0)

  PDF (Probability Density Function):
    f(x; Mu, B) = (1 / (2*B)) * exp(-abs(x - Mu) / B)

  CDF (Cumulative Distribution Function):
    F(x; Mu, B) =
      0.5 * exp((x - Mu)/B)                , if x < Mu
      1 - 0.5 * exp(-(x - Mu)/B)           , if x >= Mu

  Quantile function (Inverse CDF):
    Q(p; Mu, B) =
      Mu + B * log(2*p)                    , if 0 < p < 0.5
      Mu - B * log(2*(1 - p))              , if 0.5 <= p < 1

  Mean:
    E[X] = Mu

  Variance:
    Var[X] = 2 * B^2

  Random generation:
    If U ~ Uniform(-0.5, 0.5), then
      X = Mu - B * sign(U) * log(1 - 2*abs(U))
    follows Laplace(Mu, B).

  Functions:
    laplace_dist(Mu, B)
        -> Generate a random value from Laplace(Mu, B).

    laplace_dist_pdf(Mu, B, X)
        -> Probability density at X.

    laplace_dist_cdf(Mu, B, X)
        -> Cumulative probability at X.

    laplace_dist_quantile(Mu, B, P)
        -> Inverse CDF (quantile) at probability P, 0 < P < 1.

    laplace_dist_mean(Mu, B)
        -> Return the mean (Mu).

    laplace_dist_variance(Mu, B)
        -> Return the variance (2*B^2).

  Error handling:
    All functions throw($invalid_parameter(...)) if B =< 0.0.
    The quantile function also throws($invalid_parameter_range(...))
    if P =< 0.0 or P >= 1.0.

  ----------------------------------------------------------------------
  Typical use cases
  ----------------------------------------------------------------------
  The Laplace distribution is used when data are centered around a mean
  but have more frequent extreme deviations than a Normal distribution.
  It is often a better fit when residuals or errors show "spiky" centers
  and heavy tails.

  Examples:
    * Robust modeling of residuals:
        When fitting models where outliers occur more often than
        Gaussian assumptions allow, Laplace-distributed errors yield
        robust estimators (e.g. L1 regression instead of L2).

    * Signal or noise modeling:
        In signal processing, Laplace noise models impulsive or bursty
        interference better than Normal noise.

    * Differential privacy:
        The Laplace mechanism adds Laplace-distributed noise to preserve
        privacy while maintaining expected mean accuracy.

    * Financial modeling:
        Asset returns or forecast errors with occasional large jumps
        often follow a Laplace-like distribution rather than a Normal one.

    * Bayesian modeling:
        Laplace priors (equivalent to L1 regularization) encourage sparse
        parameter estimates, similar to the Lasso method.

  Examples:
    Picat> println(laplace_dist(0.0, 1.0)).
      -0.05650873577630594

    Picat> println(laplace_dist_pdf(0.0, 1.0, 0.0)).
      0.5

    Picat> println(laplace_dist_cdf(0.0, 1.0, 1.0)).
      0.8160602794142788

    Picat> println(laplace_dist_quantile(0.75, 0.0, 1.0)).
      0.6931471805599453

    Picat> println(laplace_dist_mean(0.0, 1.0)).
      0.0

    Picat> println(laplace_dist_variance(0.0, 1.0)).
      2.0
*/


/*
  Laplace distribution random generator
  -------------------------------------
  laplace_dist(Mu, B) generates a random value from
  Laplace(Mu, B) where
     Mu = location parameter (float)
     B  = scale parameter > 0 (float)
*/

laplace_dist(Mu, B) = X =>
    (B =< 0.0 ->
        throw($invalid_parameter(laplace_dist, B))
    ;
        U = uniform(-0.5, 0.5),
        Sign = cond(U < 0.0, -1.0, 1.0),
        X = Mu - B * Sign * log(1.0 - 2.0 * abs(U))
    ).
laplace_dist_n(Mu,B,N) = [laplace_dist(Mu,B) : _ in 1..N].


/*
  Laplace distribution PDF
*/
laplace_dist_pdf(Mu, B, X) = P =>
    (B =< 0.0 ->
        throw($invalid_parameter(laplace_dist_pdf, B))
    ;
        P = (1.0 / (2.0 * B)) * exp(-abs(X - Mu) / B)
    ).

/*
  Laplace distribution CDF
*/
laplace_dist_cdf(Mu, B, X) = CDF =>
    (B =< 0.0 ->
        throw($invalid_parameter(laplace_dist_cdf, B))
    ;
        (X < Mu ->
            CDF = 0.5 * exp((X - Mu) / B)
        ;
            CDF = 1.0 - 0.5 * exp(-(X - Mu) / B)
        )
    ).

/*
  Laplace distribution quantile function (inverse CDF)
*/
laplace_dist_quantile(Mu, B, P) = X =>
    (B =< 0.0 ->
        throw($invalid_parameter(laplace_dist_quantile, B))
    ;
        (P <= 0.0 ; P >= 1.0) ->
            throw($invalid_parameter_range(laplace_dist_quantile, P))
        ;
        (P < 0.5 ->
            X = Mu + B * log(2.0 * P)
        ;
            X = Mu - B * log(2.0 * (1.0 - P))
        )
    ).

/*
  Laplace distribution mean
*/
laplace_dist_mean(Mu, B) = Mean =>
    (B =< 0.0 ->
        throw($invalid_parameter(laplace_dist_mean, B))
    ;
        Mean = Mu
    ).

/*
  Laplace distribution variance
*/
laplace_dist_variance(Mu, B) = Var =>
    (B =< 0.0 ->
        throw($invalid_parameter(laplace_dist_variance, B))
    ;
        Var = 2.0 * B * B
    ).


/*
  Zipf distributions. 

  Note: Here are two different versions: One "canonical" version, and
        one version that is compliant with Mathematica's 
        ZipfDistribution[N,S] and ZipfDistribution[rho] (as much as possible)..

  The canonical version is
    zipf_dist_*(...)
  And the Mathematica compliant version is
    zipf_mma_dist_*(...)

 

------------------------------------------------------------------------------
  Relationship between the two Zipf parameterizations
------------------------------------------------------------------------------
The one-parameter form used here:
    P(K) = K^(-S) / ζ(S)
is the *canonical mathematical* (zeta) form commonly found in statistics
and mathematical references.

Mathematica’s ZipfDistribution[ρ], however, defines:
    P(K) = K^(-(1+ρ)) / ζ(1+ρ)
which simply uses a *shifted exponent*:
    S = 1 + ρ.

This choice makes Mathematica’s Zipf consistent with its ParetoDistribution,
whose PDF ∝ x^(-(1+ρ)) uses the same tail index convention.

Both formulations describe the same family of distributions; the parameter
relation is linear, and all derived functions (PDF, CDF, mean, variance)
match exactly when substituting S = 1 + ρ.
------------------------------------------------------------------------------
*/

/*
==============================================================================
  Two-Parameter Zipf Distribution  (Zipf / Zeta with finite support)
==============================================================================

Overview
--------
The two-parameter Zipf distribution is the finite-support version of the
classic Zipf (or Zeta) distribution. It represents the probability of the
ranked integer K ∈ {1, 2, …, N} when the frequency of items decays as a
power law.

Two forms are supported:

1. zipf_dist(N, S)
   -----------------
   The *classic* finite Zipf distribution:
       P(X = K) = K^(-S) / H(N, S),    K = 1..N
   where H(N, S) = Σ_{i=1..N} i^(-S) is the generalized harmonic number.
   This form corresponds to ZipfDistribution[N, S] in most mathematical texts.

   Domain:
       N ≥ 1,  S > 0

   Special cases:
       As N → ∞, the distribution approaches the one-parameter Zipf(S)
       with normalization constant ζ(S).

2. zipf_mma_dist(N, Rho)
   ----------------------
   The *Mathematica-compatible* form:
       P(X = K) = K^(-(1 + Rho)) / H(N, 1 + Rho),    K = 1..N
   Corresponds exactly to Mathematica's ZipfDistribution[n, ρ],
   where ρ > 0.

------------------------------------------------------------------------------

Implemented Functions
---------------------

  Random generation
      zipf_dist(N, S)                 -> integer variate on 1..N
      zipf_mma_dist(N, Rho)           -> same, Mathematica-compatible

  Probability mass function (PDF)
      zipf_dist_pdf(N, S, K)
      zipf_mma_dist_pdf(N, Rho, K)

  Cumulative distribution function (CDF)
      zipf_dist_cdf(N, S, K)
      zipf_mma_dist_cdf(N, Rho, K)

  Quantile (inverse CDF)
      zipf_dist_quantile(N, S, P)
      zipf_mma_dist_quantile(N, Rho, P)
      Returns smallest integer K in 1..N such that CDF(K) >= P.

  Mean
      zipf_dist_mean(N, S)
          = H(N, S-1) / H(N, S)            (finite for any S>1)
      zipf_mma_dist_mean(N, Rho)
          = H(N, Rho) / H(N, 1+Rho)

  Variance
      zipf_dist_variance(N, S)
          = H(N, S-2)/H(N, S) - (H(N, S-1)/H(N, S))^2
      zipf_mma_dist_variance(N, Rho)
          = H(N, Rho-1)/H(N, 1+Rho) - (H(N, Rho)/H(N, 1+Rho))^2

  where H(N, S) is the generalized harmonic number:
      H(N, S) = Σ_{i=1..N} i^(-S)

------------------------------------------------------------------------------

Examples
---------

  % Mean and variance (Mathematica-compatible)
  println(zipf_mma_dist_mean(15, 0.3)).
    % → 3.577379582807594
  println(zipf_mma_dist_variance(15, 0.3)).
    % → 2.3381775...

  % Probability mass function
  println(zipf_mma_dist_pdf(15, 0.3, 1)).
    % → 0.302943452874...
  println(zipf_mma_dist_pdf(15, 0.3, 5)).
    % → 0.087539...

  % Cumulative distribution function
  println(zipf_mma_dist_cdf(15, 0.3, 5)).
    % → 0.660813 (matches Mathematica's CDF[ZipfDistribution[15,0.3],5])

  % Quantile (inverse CDF)
  println(zipf_mma_dist_quantile(15, 0.3, 0.5)).
    % → 4   (median rank)

  % Random sampling
  Xs = [zipf_mma_dist(15, 0.3) : _ in 1..10],
  println(Xs).

------------------------------------------------------------------------------

Notes
-----
• These finite-support variants normalize by H(N,S) rather than ζ(S).
• As N → ∞, the results converge to those of the one-parameter Zipf.
• Random generation uses inverse transform sampling over 1..N.
• All functions are implemented in pure Picat, with double precision.
• Variable names follow Picat’s conventions (uppercase identifiers).
• All loops are structured; no break/continue constructs are used.

==============================================================================
*/

/*
  Zipf distribution

  Note: There are two versions:
  - The "plain" version that ChatGPT insists is the "correct" version.
    These are the zipf_dist_* functions.
  - Mathematica's version which has the S parameter one off.
    These are the zipf_mma_dist_* functions.

*/

/*
  Our "full-exponent" Zipf: P(k) ∝ k^{-S}
  Parameters: N >= 1, S > 0
*/
zipf_dist(N, S) = K =>
    (N < 1  -> throw($error('zipf_dist_rand: N must be >= 1')) ; true),
    (S =< 0 -> throw($error('zipf_dist_rand: S must be > 0'))   ; true),
    U = u01(),
    H = harmonic_number(N, S),
    I = 1, 
    Acc = (1**(-S)) / H,
    while (I < N, Acc < U) do
        I := I + 1,
        Acc := Acc + (I**(-S)) / H
    end,
    K = I.
zipf_dist_n(N,S,Num) = [zipf_dist(N,S) : _ in 1..Num].

/* (optional) Means, to verify quickly */
zipf_dist_mean(N, S) = Mean =>
    Num = harmonic_number(N, S-1.0),
    Den = harmonic_number(N, S),
    Mean = Num / Den.

/*
  Mathematica-compatible ZipfDistribution[n, ρ]:
  P(k) ∝ k^{-(1+ρ)}  i.e., S = 1 + ρ
*/
zipf_mma_dist(N, Rho) = K =>
    (Rho =< 0 -> throw($error('zipf_mma_dist: Rho must be > 0')) ; true),
    K = zipf_dist(N, 1.0 + Rho).
zipf_mma_dist_n(N, Rho, Num) = [ zipf_mma_dist(N, Rho) : _ in 1..Num].

zipf_mma_dist_mean(N, Rho) = Mean =>
    Num = harmonic_number(N, Rho),
    Den = harmonic_number(N, 1.0 + Rho),
    Mean = Num / Den.

/*
  zipf_dist_pdf(N, S, K)
  ----------------------
  Probability mass function for Zipf distribution:
    P(K) = K^(-S) / H(N,S)
  Parameters:
    N >= 1 : range (maximum rank)
    S > 0  : exponent
  Domain:
    K in 1..N (integer); otherwise 0.0
*/
zipf_dist_pdf(N, S, K) = P =>
    if S =< 0 then
        throw($error('zipf_dist_pdf: S must be > 0'))
    elseif N < 1 then
        throw($error('zipf_dist_pdf: N must be >= 1'))
    elseif not integer(K) then
        P = 0.0
    elseif K < 1 ; K > N then
        P = 0.0
    else
        H = harmonic_number(N, S),
        P = (K**(-S)) / H
    end.


/*
  zipf_mma_dist_pdf(N, Rho, K)
  -----------------------------
  Mathematica-compatible ZipfDistribution[n, ρ]:
    P(K) = K^(-(1+ρ)) / H(N,1+ρ)
  Parameters:
    N >= 1 : range
    Rho > 0: exponent parameter ρ
  Domain:
    K in 1..N (integer); otherwise 0.0
*/
zipf_mma_dist_pdf(N, Rho, K) = P =>
    if Rho =< 0 then
        throw($error('zipf_mma_dist_pdf: Rho must be > 0'))
    else
        P = zipf_dist_pdf(N, 1.0 + Rho, K)
    end.



/*
  zipf_dist_cdf(N, S, K)
  ----------------------
  CDF for Zipf with full exponent S (P(k) ∝ k^{-S}).
  Parameters:
    N >= 1, S > 0
  Argument:
    K : value at which CDF is evaluated.
        For non-integer K, we follow the standard discrete convention: use floor(K).
  Returns:
    CDF = sum_{i=1..floor(K)} i^{-S} / H(N,S)
*/
zipf_dist_cdf(N, S, K) = C =>
    if S =< 0 then
        throw($error('zipf_dist_cdf: S must be > 0'))
    elseif N < 1 then
        throw($error('zipf_dist_cdf: N must be >= 1'))
    else
        K1 = floor(K),                            % discrete CDF uses floor
        if K1 < 1 then
            C = 0.0
        elseif K1 >= N then
            C = 1.0
        else
            H = harmonic_number(N, S),
            Num = sum([I**(-S) : I in 1..K1]),
            C = Num / H
        end
    end.


/*
  zipf_mma_dist_cdf(N, Rho, K)
  ----------------------------
  Mathematica-compatible ZipfDistribution[n, ρ] where P(k) ∝ k^{-(1+ρ)}.
  Parameters:
    N >= 1, Rho > 0
  Argument:
    K (uses floor(K) for discrete CDF)
*/
zipf_mma_dist_cdf(N, Rho, K) = C =>
    if Rho =< 0 then
        throw($error('zipf_mma_dist_cdf: Rho must be > 0'))
    else
        C = zipf_dist_cdf(N, 1.0 + Rho, K)
    end.


/*
  zipf_dist_quantile(N, S, P)
  ---------------------------
  Quantile function (inverse CDF) for Zipf with exponent S.

  Returns the smallest integer K such that
      CDF(K) >= P

  Parameters:
    N >= 1 : range (maximum rank)
    S > 0  : exponent
    P ∈ [0,1]

  Uses a simple linear search since N is usually modest.
  (For large N, a binary search over a precomputed cumulative table would be faster.)
*/
zipf_dist_quantile(N, S, P) = K =>
    if S =< 0 then
        throw($error('zipf_dist_quantile: S must be > 0'))
    elseif N < 1 then
        throw($error('zipf_dist_quantile: N must be >= 1'))
    elseif P < 0.0 ; P > 1.0 then
        throw($error('zipf_dist_quantile: P must be in [0,1]'))
    elseif P =< 0.0 then
        K = 1
    elseif P >= 1.0 then
        K = N
    else
        H = harmonic_number(N, S),
        Acc = 0.0,
        I = 1,
        K = _,
        while (I < N, Acc < P) do
            Acc := Acc + (I**(-S)) / H,
            if Acc >= P then
                K := I
            else
                I := I + 1
            end
        end,
        (var(K) -> K = N ; true)
    end.


/*
  zipf_mma_dist_quantile(N, Rho, P)
  ---------------------------------
  Quantile for Mathematica-style ZipfDistribution[n, ρ],
  which uses exponent (1+ρ).

  Parameters:
    N >= 1 : range
    Rho > 0: exponent parameter ρ
    P ∈ [0,1]
*/
zipf_mma_dist_quantile(N, Rho, P) = K =>
    if Rho =< 0 then
        throw($error('zipf_mma_dist_quantile: Rho must be > 0'))
    else
        K = zipf_dist_quantile(N, 1.0 + Rho, P)
    end.


/*
  zipf_dist_variance(N, S)
  ------------------------
  Variance of Zipf distribution with full exponent S.

  Var(X) = E[X^2] - (E[X])^2
          = H(N, S-2)/H(N, S) - (H(N, S-1)/H(N, S))^2

  Parameters:
    N >= 1 : range (maximum rank)
    S > 0  : exponent
*/
zipf_dist_variance(N, S) = Var =>
    if S =< 0 then
        throw($error('zipf_dist_variance: S must be > 0'))
    elseif N < 1 then
        throw($error('zipf_dist_variance: N must be >= 1'))
    else
        Hs   = harmonic_number(N, S),
        Hs1  = harmonic_number(N, S - 1.0),
        Hs2  = harmonic_number(N, S - 2.0),
        Mean = Hs1 / Hs,
        Var  = (Hs2 / Hs) - Mean**2
    end.


/*
  zipf_mma_dist_variance(N, Rho)
  ------------------------------
  Variance of Mathematica-compatible ZipfDistribution[n, ρ].

  Var(X) = H(N, ρ-1)/H(N,1+ρ) - (H(N,ρ)/H(N,1+ρ))^2
         = H(N, Rho - 1)/H(N, 1 + Rho) - Mean^2
  (uses S = 1 + Rho)
*/
zipf_mma_dist_variance(N, Rho) = Var =>
    if Rho =< 0 then
        throw($error('zipf_mma_dist_variance: Rho must be > 0'))
    else
        Var = zipf_dist_variance(N, 1.0 + Rho)
    end.

/*
==============================================================================
  Zipf Distribution (one-parameter and Mathematica-compatible variants)
==============================================================================

Overview
--------
The Zipf distribution is a discrete probability distribution that models
the frequency of events where the rank of an event is inversely proportional
to its probability. It is often used for modeling word frequencies, city
populations, and other rank–frequency phenomena.

Two forms are supported:

1. zipf_dist(S)
   -----------------
   The *classic* one-parameter Zipf distribution:
       P(X = k) = k^(-S) / ζ(S),     k = 1,2,3,...
   where ζ(S) is the Riemann zeta function.
   This form corresponds to ZipfDistribution[S] in many mathematical texts.

   Domain:
       S > 1         (distribution defined)
       S > 2         (finite mean)
       S > 3         (finite variance)

2. zipf_mma_dist(Rho)
   -----------------
   The *Mathematica-compatible* form:
       P(X = k) = k^(-(1 + Rho)) / ζ(1 + Rho),    k = 1,2,3,...
   This corresponds to Mathematica's ZipfDistribution[ρ],
   where ρ > 0.

   Domain:
       Rho > 0       (distribution defined)
       Rho > 1       (finite mean)
       Rho > 2       (finite variance)

------------------------------------------------------------------------------

Implemented Functions
---------------------

  Random generation
      zipf_dist(S)                -> integer variate for given exponent S
      zipf_mma_dist(Rho)          -> same, Mathematica-style

  Probability mass function (PDF)
      zipf_dist_pdf(S, K)
      zipf_mma_dist_pdf(Rho, K)

  Cumulative distribution function (CDF)
      zipf_dist_cdf(S, K)
      zipf_mma_dist_cdf(Rho, K)

  Quantile (inverse CDF)
      zipf_dist_quantile(S, P)
      zipf_mma_dist_quantile(Rho, P)
      Returns smallest integer K such that F(K) >= P.
      For infinite support, undefined at P = 1.0.

  Mean
      zipf_dist_mean(S)           -> finite if S > 2
          = ζ(S-1) / ζ(S)

      zipf_mma_dist_mean(Rho)     -> finite if Rho > 1
          = ζ(Rho) / ζ(1+Rho)

  Variance
      zipf_dist_variance(S)       -> finite if S > 3
          = ζ(S-2)/ζ(S) - (ζ(S-1)/ζ(S))^2

      zipf_mma_dist_variance(Rho) -> finite if Rho > 2
          = ζ(Rho-1)/ζ(1+Rho) - (ζ(Rho)/ζ(1+Rho))^2

------------------------------------------------------------------------------

Examples
---------

  % Mean and variance (Mathematica-compatible)
  println(zipf_mma_dist_mean(1.3)).
    % → 2.744973807275084
  println(zipf_mma_dist_variance(2.5)).
    % → 0.609892...

  % Probability mass function
  println(zipf_mma_dist_pdf(1.3, 1)).
    % → 0.698120339238389
  println(zipf_mma_dist_pdf(1.3, 2)).
    % → 0.141763073364556

  % Cumulative distribution function
  println(zipf_mma_dist_cdf(1.3, 2)).
    % → 0.839883 (matches Mathematica's CDF[ZipfDistribution[1.3],2])

  % Quantile (inverse CDF)
  println(zipf_mma_dist_quantile(2.0, 0.9999)).
    % → 65  (smallest K s.t. CDF(K) ≥ 0.9999)

  % Random sampling (pseudo-random example)
  Xs = [zipf_mma_dist(1.5) : _ in 1..10],
  println(Xs).

------------------------------------------------------------------------------

Notes
-----
• The functions are implemented in pure Picat, using IEEE double precision.
• For random generation, inverse transform sampling is used.
• For quantiles, a numerically stable tail-bound method is applied.
• All variable names follow Picat’s convention (uppercase identifiers).
• The implementations avoid break/continue and use structured loops only.

==============================================================================
*/



/*
  zipf_dist_rand(S)
  -----------------
  Random integer from Zipf(S) with S > 1.
  Uses Devroye’s rejection method, written without break/continue.
*/
zipf_dist(S) = K =>
    if S =< 1.0 then
        throw($error('zipf_dist_rand: S must be > 1 for convergence'))
    else
        Alpha = 1.0 / (S - 1.0),
        Accepted = false,
        K = 1,
        while (not Accepted)
            U = u01(),
            V = u01(),
            X = floor(U ** (-Alpha)),       % candidate ≥ 1
            % acceptance test:
            if X > 0 then
                % (1 + 1/X)^(S-1) - 1  gives acceptance ratio
                % add a small amout to avoid division by 0
                A =  0.0000001 + (1.0 + 1.0 / X) ** (S - 1.0) - 1.0,
                if V * X >= 1.0 / A then
                    % reject: do nothing (loop continues)
                    true
                else
                    K :=  X,               
                    Accepted := true
                end
            end
        end
    end.
zipf_dist_n(S,N) = [zipf_dist(S) : _ in 1..N].

zipf_mma_dist(Rho) = K =>
    if Rho =< 0.0 then
        throw($error('zipf_mma_dist_rand: Rho must be > 0'))
    else
        K = zipf_dist(1.0 + Rho)
    end.
zipf_mma_dist_n(S,N) = [zipf_mma_dist(S) : _ in 1..N].

/*
  zipf_dist_mean(S)
  -----------------
  Mean of one-parameter Zipf distribution (P(k) ∝ k^{-S}).
  Finite only for S > 2.
*/
zipf_dist_mean(S) = Mean =>
    if S =< 2.0 then
        throw($error('zipf_dist_mean: mean diverges for S <= 2'))
    else
        Mean = zeta(S - 1.0) / zeta(S)
    end.


/*
  zipf_mma_dist_mean(Rho)
  -----------------------
  Mean of Mathematica-style ZipfDistribution[ρ].
  Finite only for Rho > 1.
*/
zipf_mma_dist_mean(Rho) = Mean =>
    if Rho =< 1.0 then
        throw($error('zipf_mma_dist_mean: mean diverges for Rho <= 1'))
    else
        Mean = zeta(Rho) / zeta(1.0 + Rho)
    end.


/*
  zipf_dist_pdf(S, K)
  -------------------
  Probability mass function for the one-parameter Zipf distribution:
      P(K) = K^(-S) / ζ(S)
  Parameters:
      S > 1  : exponent
      K >= 1 : integer support
  Returns 0.0 for invalid K.
*/
zipf_dist_pdf(S, K) = P =>
    if S =< 1.0 then
        throw($error('zipf_dist_pdf: S must be > 1'))
    elseif not integer(K) ; K < 1 then
        P = 0.0
    else
        P = (K**(-S)) / zeta(S)
    end.


/*
  zipf_mma_dist_pdf(Rho, K)
  -------------------------
  PDF for Mathematica-style ZipfDistribution[ρ]:
      P(K) = K^(-(1+ρ)) / ζ(1+ρ)
  Parameters:
      Rho > 0 : exponent parameter
      K >= 1  : integer support
*/
zipf_mma_dist_pdf(Rho, K) = P =>
    if Rho =< 0.0 then
        throw($error('zipf_mma_dist_pdf: Rho must be > 0'))
    elseif not integer(K) ; K < 1 then
        P = 0.0
    else
        P = (K**(-(1.0 + Rho))) / zeta(1.0 + Rho)
    end.


/*
  hurwitz_zeta(S, A)
  ------------------
  Hurwitz zeta ζ(S,A) = Σ_{n=0..∞} (n+A)^(-S)
  Approximation using partial sum + integral remainder.
  Accurate to about 1e-10 for S>1, A>0.
*/
hurwitz_zeta(S, A) = Z =>
    if S =< 1.0 then
        throw($error('hurwitz_zeta: S must be > 1'))
    elseif A =< 0.0 then
        throw($error('hurwitz_zeta: A must be > 0'))
    else
        Nmax = 10000,
        Sum = sum([(A + I)**(-S) : I in 0..Nmax]),
        Tail = ((A + Nmax)**(1.0 - S))/(S - 1.0)
             + 0.5*((A + Nmax)**(-S)),
        Z = Sum + Tail
    end.



/*
  zipf_dist_cdf(S, K)
  -------------------
  CDF for Zipf(S):  F(K) = 1 - ζ(S, K+1)/ζ(S)
  Parameters:
      S > 1
      K >= 1 (integer)
*/
zipf_dist_cdf(S, K) = F =>
    if S =< 1.0 then
        throw($error('zipf_dist_cdf: S must be > 1'))
    elseif not integer(K) ; K < 1 then
        F = 0.0
    else
        F = 1.0 - hurwitz_zeta(S, K + 1.0) / zeta(S)
    end.



/*
  zipf_mma_dist_cdf(Rho, K)
  -------------------------
  Mathematica-style ZipfDistribution[ρ]:
      F(K) = 1 - ζ(1+ρ, K+1)/ζ(1+ρ)
  Parameters:
      Rho > 0
      K >= 1 (integer)
*/
zipf_mma_dist_cdf(Rho, K) = F =>
    if Rho =< 0.0 then
        throw($error('zipf_mma_dist_cdf: Rho must be > 0'))
    elseif not integer(K) ; K < 1 then
        F = 0.0
    else
        F = 1.0 - hurwitz_zeta(1.0 + Rho, K + 1.0) / zeta(1.0 + Rho)
    end.


/*
  zipf_dist_quantile(S, P)
  ------------------------
  One-parameter Zipf, S>1. Smallest K with F(K) >= P.
  No zeta used: partial sums + integral tail upper bound.
*/
zipf_dist_quantile(S, P) = K =>
    if S =< 1.0 then
        throw($error('zipf_dist_quantile: S must be > 1'))
    elseif P < 0.0 ; P > 1.0 then
        throw($error('zipf_dist_quantile: P must be in [0,1]'))
    elseif P =< 0.0 then
        K = 1
    elseif P >= 1.0 then
        throw($error('zipf_dist_quantile: undefined at P=1 for infinite support'))
    else
        % Optional: start near the tail-only guess
        % Start = max(1, ceiling(((1.0)/((S-1.0)*(1.0-P))) ** (1.0/(S-1.0))))
        I   = 1,         % or Start
        ACC = 0.0,
        K   = 1,
        Done = false,
        while (not Done)
            ACC := ACC + (I ** (-S)),
            TailUpper = (I ** (1.0 - S)) / (S - 1.0),   % ∫_{I}^{∞} x^{-S} dx
            Flow      = ACC / (ACC + TailUpper),        % lower bound on CDF
            if Flow >= P then
                K := I,
                Done := true
            else
                I := I + 1
            end
        end
    end.

/*
  Mathematica-compatible: ZipfDistribution[ρ] with S = 1 + ρ
*/
zipf_mma_dist_quantile(Rho, P) = K =>
    if Rho =< 0.0 then
        throw($error('zipf_mma_dist_quantile: Rho must be > 0'))
    else
        K = zipf_dist_quantile(1.0 + Rho, P)
    end.


/*
  zipf_dist_variance(S)
  ---------------------
  Variance of one-parameter Zipf (infinite support), P(k) ∝ k^{-S}.
  Finite iff S > 3.

  Var(X) = ζ(S-2)/ζ(S) - (ζ(S-1)/ζ(S))^2
*/
zipf_dist_variance(S) = Var =>
    if S =< 3.0 then
        throw($error('zipf_dist_variance: variance diverges for S <= 3'))
    else
        ZS   = zeta(S),
        ZSm1 = zeta(S - 1.0),
        ZSm2 = zeta(S - 2.0),
        Var  = (ZSm2 / ZS) - (ZSm1 / ZS)**2
    end.


/*
  zipf_mma_dist_variance(Rho)
  ---------------------------
  Mathematica-compatible ZipfDistribution[ρ], where S = 1 + Rho.
  Finite iff Rho > 2.

  Var(X) = ζ(ρ-1)/ζ(1+ρ) - (ζ(ρ)/ζ(1+ρ))^2
*/
zipf_mma_dist_variance(Rho) = Var =>
    if Rho =< 2.0 then
        throw($error('zipf_mma_dist_variance: variance diverges for Rho <= 2'))
    else
        S    = 1.0 + Rho,
        ZS   = zeta(S),
        ZSm1 = zeta(S - 1.0),  % = ζ(ρ)
        ZSm2 = zeta(S - 2.0),  % = ζ(ρ-1)
        Var  = (ZSm2 / ZS) - (ZSm1 / ZS)**2
    end.




/*
  Matching distribution

  Probability of correct guessing drawn numbers without replacement.

  https://ora.ox.ac.uk/objects/uuid:478fa6d8-bc7f-458d-92f5-7e5d63d9824a/files/mf8a85f79596ca2150078192a755dc020
  https://stats.libretexts.org/Bookshelves/Probability_Theory/Probability_Mathematical_Statistics_and_Stochastic_Processes_(Siegrist)/12:_Finite_Sampling_Models/12.05:_The_Matching_Problem

  See ppl_matching_distribution.rkt for more on this.

*/


% Generate random variates
matching_dist(N) = matching_dist_quantile(N,U) => U = uniform(0,1).
matching_dist_n(N,Num) = [matching_dist(N) : _ in 1..Num].

matching_dist_pdf(N,R) = (1 / factorial(R)) * sum([ (-1)**T / factorial(T)  : T in 0..N-R]).

matching_dist_cdf(N,R) = sum([matching_dist_pdf(N,I) : I in 0..R]).

matching_dist_quantile(N,Q) = Res =>
  Res = _,
  OK = true,  
  foreach(I in 0..N, break(OK == false))
    if matching_dist_cdf(N,I) >= Q then
      Res = I,
      OK := false
    end
  end.

matching_dist_mean(N) = 1.
matching_dist_variance(N) = 1.



/*
  Pareto distributions.

  These are ported from Mathematica ParetoDistribution

  Type 1: pareto1_dist(K,Alpha):
          represents a Pareto distribution with minimum value parameter k 
          and shape parameter Alpha. 
  
  Type 2: pareto2_dist(K,Alpha,Mu):
          represents a Pareto type II distribution with location parameter Mu
   
  Type 3: pareto3_dist(K,1,Gamma,Mu)
          
  Type 4: pareto4_dist(K,Alpha,Gamma,Mu):
          represents a Pareto type IV distribution with shape parameter Gamma.

  pareto2_dist[K,Alpha,0] is also known as Lomax distribution.

*/

/*
  Overview
  --------
  The Pareto Type I distribution is a continuous power-law distribution that
  describes phenomena where large values occur with small probability.
  It is often used to model income, wealth, file sizes, and other quantities
  with heavy-tailed behavior.

*/

% pareto1_dist(K, Alpha) = pareto1_dist_quantile(K, Alpha, U) =>
%    U = uniform_dist(0,1).
pareto1_dist(K, Alpha) = X =>
    if K =< 0.0 then
        throw($error('pareto1_dist: K must be > 0'))
    elseif Alpha =< 0.0 then
        throw($error('pareto1_dist: Alpha must be > 0'))
    else
        U = u01(),                              % uniform in (0,1)
        % X = K * ((1.0 - U) ** (-1.0 / Alpha))   % inverse transform
        X = K * safe_pow(1.0 - U,-1.0 / Alpha)   % inverse transform
    end.

pareto1_dist_n(K, Alpha, N) = [pareto1_dist(K,Alpha) : _ in 1..N].

pareto1_dist_pdf(K, Alpha, X) = cond(X >= K, K**Alpha * X**(-1-Alpha) * Alpha,0).

pareto1_dist_cdf(K, Alpha, X) = cond(X >= K, 1 - (K/X)**Alpha,0).

pareto1_dist_quantile(K, Alpha, X) = Res =>
  if X >= 0, X <= 1 then
     Res = K*(1-X)**(-1/Alpha)
  else
     throw($error('pareto1_dist_quantile: X must be >= 0 and <= 1'))
  end.
  
pareto1_dist_mean(K, Alpha) = Res =>
  if Alpha >= 1 then
    Res = (K*Alpha)/(Alpha-1)
  else
    % throw($error('pareto1_dist_mean: Alpha must be >= 2'))
    throw($error('pareto1_dist_mean: Alpha must be >= 1'))
  end.

pareto1_dist_variance(K, Alpha) = Res =>
  if Alpha >= 1 then
    Res = (K**2*Alpha)/((Alpha-2)*(Alpha-1)**2)
  else
    throw($error('pareto1_dist_mean: Alpha must be >= 1'))
  end.

/*
  Pareto distribution type II

  pareto2_dist(K,Alpha,Mu)

  This is the *shifted* (or *translated*) Pareto Type I distribution.
  It generalizes ParetoDistribution[K, Alpha] by introducing a location
  parameter Mu that shifts the distribution along the x-axis.

  Parameters:
    K     > 0   : scale (minimum positive offset)
    Alpha > 0   : shape (tail index)
    Mu    ∈ ℝ   : location (shift)

  Support:
    x ∈ [Mu + K, ∞)

  PDF:
    f(x) = (Alpha * K^Alpha) / (x - Mu)^(Alpha + 1),
           for x ≥ Mu + K.

  CDF:
    F(x) = 1 - (K / (x - Mu))^Alpha,   for x ≥ Mu + K.

  Quantile:
    Q(p) = Mu + K * (1 - p)^(-1/Alpha),   0 ≤ p < 1.

  Moments:
    Mean     = Mu + (Alpha * K) / (Alpha - 1),   (finite for Alpha > 1)
    Variance = (Alpha * K^2) / ((Alpha - 1)^2 * (Alpha - 2)),  (finite for Alpha > 2)

  Notes:
  * When Mu = 0, this reduces to ParetoDistribution[K, Alpha].
  * The location parameter shifts the entire distribution to start at Mu + K.
  * Random generation:  X = Mu + K * (1 - U)^(-1/Alpha),  U ∼ Uniform(0,1).

*/

% pareto2_dist(K, Alpha, Mu) = pareto2_dist_quantile(K, Alpha, Mu, U) =>
%     U = uniform_dist(0,1).
pareto2_dist(K, Alpha, Mu) = Mu + K * (1-U)**(-1/Alpha) =>
  U = uniform(0,1).

pareto2_dist_n(K, Alpha, Mu, N) = [pareto2_dist(K,Alpha,Mu) : _ in 1..N].

pareto2_dist_pdf(K, Alpha, Mu, X) = cond(X >= Mu, (((K-Mu+X)/K)**(-1-Alpha)*Alpha)/K ,0).

pareto2_dist_cdf(K, Alpha, Mu, X) = cond(X >= Mu, 1 - (1+ (X-Mu)/K)**(-Alpha),0).

pareto2_dist_quantile(K, Alpha, Mu, X) = Res =>
  if X >= 0, X <= 1 then
     if X > 0, X < 1 then
       Res = Mu + K*((1-X)**(-1/Alpha)-1)
     elseif X <= 0 then
       Res = Mu
     else
       throw($error('pareto2_dist_quantile: X must be >= 0 and <= 1'))
     end
  else
     throw($error('pareto2_dist_quantile: X must be >= 0 and <= 1'))
  end.
  
pareto2_dist_mean(K, Alpha, Mu) = Res =>
  if Alpha > 1 then
    Res = Mu + K/(Alpha-1)
  else
    throw($error('pareto1_dist_mean: Alpha must be >= 1'))
  end.

pareto2_dist_variance(K, Alpha, Mu) = Res =>
  if Alpha > 2 then
    Res = (K**2*Alpha)/((Alpha-2)*(Alpha-1)**2)
  else
    throw($error('pareto2_dist_mean: Alpha must be >= 2'))
  end.


/*
  Pareto dist type IV

  pareto3_dist(K,Gamma,Mu)
 
  -> pareto4_dist(K,1,Gamma,Mu) 

*/
pareto3_dist(K, Gamma, Mu) = pareto4_dist(K, 1, Gamma, Mu).
pareto3_dist_n(K, Alpha, Mu, N) = [pareto3_dist(K,Alpha,Mu) : _ in 1..N].
pareto3_dist_pdf(K, Gamma, Mu,X) = pareto4_dist_pdf(K, 1, Gamma, Mu,X).
pareto3_dist_cdf(K, Gamma, Mu,X) = pareto4_dist_cdf(K, 1, Gamma, Mu,X).
pareto3_dist_quantile(K, Gamma, Mu,X) = pareto4_dist_quantile(K, 1, Gamma, Mu,X).
pareto3_dist_mean(K, Gamma, Mu) = pareto4_dist_mean(K, 1, Gamma, Mu).
pareto3_dist_variance(K, Gamma, Mu) = pareto4_dist_variance(K, 1, Gamma, Mu).


/*
  Pareto dist type IV

  pareto4_dist(K,Alpha,Gamma,Mu)

*/
pareto4_dist(K, Alpha, Gamma, Mu) = pareto4_dist_quantile(K, Alpha, Mu, Gamma, U) =>
   U = uniform_dist(0,1).

pareto4_dist_n(K, Alpha, Gamma, Mu, N) = [pareto4_dist(K,Alpha,Gamma, Mu) : _ in 1..N].

pareto4_dist_pdf(K, Alpha, Gamma, Mu, X) = Res =>
  if X >= Mu then
   Res = (K**(-1/Gamma) * (X-Mu)**(1/Gamma-1) * (1+(K/(X-Mu))**(-1/Gamma))**(-Alpha-1) * Alpha)/Gamma
  else
    Res = 0
  end.

pareto4_dist_cdf(K, Alpha, Gamma, Mu, X) = Res =>
  if X >= Mu then
    Res = 1 - (1 + ((X-Mu)/K)**(1/Gamma))**(-Alpha)
  else
    Res = 0
  end.

pareto4_dist_quantile(K, Alpha, Gamma, Mu, X) = Res =>
  if X >= 0, X <= 1 then
     if X > 0, X < 1 then
       Res = Mu + K*((1-X)**(-1/Alpha)-1)**Gamma
     elseif X <= 0 then
       Res = Mu
     else
       throw($error('pareto4_dist_quantile: X must be >= 0 and <= 1'))
     end
  else
     throw($error('pareto4_dist_quantile: X must be >= 0 and <= 1'))
  end.
  
pareto4_dist_mean(K, Alpha, Gamma, Mu) = Res =>
  if Alpha > Gamma then
    Res = Mu + (K * gamma_func(1+Gamma)*gamma_func(Alpha-Gamma))/gamma_func(Alpha)
  else
    throw($error('pareto1_dist_mean: Alpha must be > Gamma'))
  end.

pareto4_dist_variance(K, Alpha, Gamma, Mu) = Res =>
  if Alpha > 2*Gamma then
    Res = 1/gamma_func(Alpha)**2 *
         K**2 *(
         gamma_func(1+2*Gamma)*gamma_func(Alpha)*gamma_func(Alpha-2*Gamma)
         - 
         gamma_func(1+Gamma)**2 * gamma_func(Alpha-Gamma)**2)
  else
    throw($error('pareto4_dist_mean: Alpha must be >= 2*Gamma'))
  end.


/*
  Extreme Value distribution

  From Mathematica ExtremeValueDistribution:
  The extreme value distribution gives the asymptotic distribution of the maximum 
  value in a sample from a distribution such as the normal distribution.

  This is compliant with Mathematica's 
  * ExtremeValueDistribution[alpha,beta]
    represents an extreme value distribution with location parameter Alpha 
    and scale parameter Beta..

  * ExtremeValueDistribution[] -> ExtremeValueDistribution[0,1]


*/
extreme_value_dist(Alpha,Beta) = extreme_value_dist_quantile(Alpha, Beta, U) =>
  U = u01().
extreme_value_dist_n(Alpha,Beta,N) = [extreme_value_dist(Alpha,Beta) : _ in 1..N].

extreme_value_dist_pdf(Alpha, Beta, X) =
  exp( -exp(Alpha-X)/Beta + (Alpha-X)/Beta)/Beta.

extreme_value_dist_cdf(Alpha, Beta, X) = exp(-exp( (Alpha-X)/Beta)).

extreme_value_dist_quantile(Alpha, Beta, X) = Res =>
  if X > 0, X < 1 then
    Res = Alpha-Beta * log(-log(X))
  else
    throw($error('extreme_value_dist_quantile: X > 0, X < 1'))
  end.

extreme_value_dist_mean(Alpha, Beta) = Alpha + euler_gamma() * Beta.

extreme_value_dist_variance(Alpha, Beta) = pi**2 * Beta**2 / 6.


/* 
  extreme_value_dist() -> extreme_value_dist(0,1) 
*/
extreme_value_dist() = extreme_value_dist(0,1).
extreme_value_dist_n(N) = [extreme_value_dist() : _ in 1..N].

extreme_value_dist_pdf(X) = extreme_value_dist_pdf(0,1,X).
extreme_value_dist_cdf(X) = extreme_value_dist_cdf(0,1,X).
extreme_value_dist_quantile(X) = extreme_value_dist_quantile(0,1,X).
extreme_value_dist_mean() = extreme_value_dist_mean(0,1).
extreme_value_dist_variance() = extreme_value_dist_variance(0,1).


/*
  Gumbel distribution

  gumbel_distribution(Alpha,Beta)

  From Mathematica:
  -  GumbelDistribution[Alpha,Beta]
     Represents a Gumbel distribution with location parameter Alpha and scale parameter Beta.
  - GumbelDistribution[] -> GumbelDistribution[0,1]
  """
  The Gumbel distribution gives the asymptotic distribution of the *minimum* value in a sample 
  from a distribution such as the normal distribution.

  The asymptotic distribution of the *maximum* value, also sometimes called a Gumbel
  distribution, is implemented in the Wolfram Language as ExtremeValueDistribution.
  """
*/

gumbel_dist(Alpha,Beta) = gumbel_dist_quantile(Alpha, Beta, U) =>
  U = u01().
gumbel_dist_n(Alpha,Beta,N) = [gumbel_dist(Alpha,Beta) : _ in 1..N].

gumbel_dist_pdf(Alpha, Beta, X) =
  exp( -exp(X-Alpha)/Beta + (X-Alpha)/Beta)/Beta.

gumbel_dist_cdf(Alpha, Beta, X) = 1-exp(-exp( (X-Alpha)/Beta)).

gumbel_dist_quantile(Alpha, Beta, X) = Res =>
  if X > 0, X < 1 then
    Res = Alpha + Beta * log(-log(1-X))
  else
    throw($error('gumbel_dist_quantile: X > 0, X < 1'))
  end.

gumbel_dist_mean(Alpha, Beta) = Alpha - euler_gamma() * Beta.

gumbel_dist_variance(Alpha, Beta) = pi**2 * Beta**2 / 6.


/*
  gumbel_dist() -> gumbel_dist(0,1)
  
*/
gumbel_dist() = gumbel_dist(0,1).
gumbel_dist_n(N) = [gumbel_dist() : _ in 1..N].

gumbel_dist_pdf(X) = gumbel_dist_pdf(0,1,X).
gumbel_dist_cdf(X) = gumbel_dist_cdf(0,1,X).
gumbel_dist_quantile(X) = gumbel_dist_quantile(0,1,X).
gumbel_dist_mean() = gumbel_dist_mean(0,1).
gumbel_dist_variance() = gumbel_dist_variance(0,1).



/*
  Weilbull distribution

  weibull_distribution(Alpha,Beta)

  This is compliant with Mathematica
  * WeibullDistribution[Alpha,Beta]
    represents a Weibull distribution with shape parameter Alpha and scale parameter Beta
  * WeibullDistribution[Alpha,Beta,Mu]
    Represents a Weibull distribution with shape parameter Alpha, scale parameter Beta, 
    and location parameter Mu.

  Note (and beware): My Gamble variant is (weibull lambda k):
  lambda: scale
  k     : shape

  I.e. the inverse parameter order than Picat PPL's Mathematica compliant version.

*/
weibull_dist(Alpha, Beta) = weibull_dist_quantile(Alpha,Beta,U) =>
  U = u01().
weibull_dist_n(Alpha, Beta, N) = [weibull_dist(Alpha,Beta) : _ in 1..N].

weibull_dist_pdf(Alpha, Beta, X) = Res =>
   if X > 0 then
     Res = exp(-(X/Beta)**Alpha) * Alpha * (X/Beta)**(Alpha-1) / Beta
   else
     Res = 0
   end.

weibull_dist_cdf(Alpha, Beta, X) = Res =>
  if X > 0 then
    Res = 1 - exp(-(X/Beta)**Alpha)
  else
    Res = 0
  end.

weibull_dist_quantile(Alpha, Beta, X) = Res =>
  if X >= 0, X <= 1 then
    if X > 0, X < 1 then
      % Res = Beta * (-log(1-X))**(1/Alpha)
      Res = Beta * safe_pow((-log(1-X)),(1/Alpha))
    else
      Res = 0
    end
  else
    throw($error('weibull_dist_quantile: X >= 0, X < 1'))
  end.

weibull_dist_mean(Alpha, Beta) = Beta * gamma_func(1+ 1/Alpha).

weibull_dist_variance(Alpha, Beta) =
  Beta**2 * (-gamma_func(1+ 1/Alpha)**2 + gamma_func(1+2/Alpha)).

/*
  weibull_dist(Alpha,Beta,Mu)

*/

weibull_dist(Alpha, Beta, Mu) = weibull_dist_quantile(Alpha,Beta,U) =>
  U = u01().
weibull_dist_n(Alpha, Beta, Mu,N) = [weibull_dist(Alpha,Beta, Mu) : _ in 1..N].

weibull_dist_pdf(Alpha, Beta, Mu, X) = Res =>
   if X > Mu then
     Res = exp(-((X-Mu)/Beta)**Alpha * ((X-Mu)/Beta)**(Alpha-1)) / Beta
   else
     Res = 0
   end.

weibull_dist_cdf(Alpha, Beta, Mu, X) = Res =>
  if X > Mu then
    Res = 1 - exp(-( ((X-Mu)/Beta)**Alpha))
  else
    Res = 0
  end.

weibull_dist_quantile(Alpha, Beta, Mu, X) = Res =>
  if X >= 0, X <= 1 then
    if X > 0, X < 1 then
      Res = Mu + Beta * (-log(1-X)**1/Alpha)
    else
      Res = Mu
    end
  else
    throw($error('weibull_dist_quantile: X >= 0, X < 1'))
  end.

weibull_dist_mean(Alpha, Beta, Mu) = Mu + Beta * gamma_func(1 + 1/Alpha).

weibull_dist_variance(Alpha, Beta, Mu) =
  Beta**2 * (-gamma_func(1+ 1/Alpha)**2 + gamma_func(1+2/Alpha)).



/*
   Frechet distribution

   frechet_dist(Alpha,Beta)
   frechet_dist(Alpha,Beta,Mu)

   This is compliant with Mathematica
   * FrechetDistribution[Alpha,Beta]
     Represents the Fréchet distribution with shape parameter Alpha and 
     scale parameter Beta.
   * FrechetDistribution[Alpha,Beta,Mu]
     Represents the Fréchet distribution with shape parameter Alpha, 
     scale parameter Beta, and location parameter Mu.

   The Fréchet distribution gives the asymptotic distribution of the maximum 
   value in a sample from a distribution such as the Cauchy distribution.
   FrechetDistribution is also known as type II extreme value distribution.

*/
frechet_dist(Alpha, Beta) = frechet_dist_quantile(Alpha,Beta,U) =>
  U = u01().
frechet_dist_n(Alpha, Beta, N) = [frechet_dist(Alpha,Beta) : _ in 1..N].

frechet_dist_pdf(Alpha, Beta, X) = Res =>
  if X > 0 then
    Res = exp( -(X/Beta)**(-Alpha)) * Alpha * (X/Beta)**(-1-Alpha) / Beta
  else
    Res = 0
  end.

frechet_dist_cdf(Alpha, Beta, X) = Res =>
  if X > 0 then
    Res = exp(-(X/Beta)**(-Alpha))
  else
    Res = 0
  end.

frechet_dist_quantile(Alpha, Beta, X) = Res =>
  if X >= 0, X <= 1 then
    if X > 0, X < 1 then
      % Res = Beta * (-log(X))**(-1/Alpha)
      Res = Beta * safe_pow(-log(X),(-1/Alpha))
    else
      Res = 0
    end
  else
    throw($error('frechet_dist_quantile: X >= 0, X < 1'))    
  end.
  
frechet_dist_mean(Alpha, Beta) = Res =>
  if 1 < Alpha then
    Res = Beta * gamma_func(1-1/Alpha)
  else
    throw($error('frechet_dist_mean: 1 < Alpha'))      
  end.

frechet_dist_variance(Alpha, Beta) = Res =>
  if Alpha > 2 then
    Res = Beta**2 * (gamma_func(1-2/Alpha)-gamma_func(1-1/Alpha)**2)
  else
    throw($error('frechet_dist_mean: Alpha > 2 '))      
  end.


/*
  frechet_distribution(Alpha,Beta,Mu)
*/

frechet_dist(Alpha, Beta, Mu) = frechet_dist_quantile(Alpha,Beta,Mu,U) =>
  U = u01().
frechet_dist_n(Alpha, Beta, Mu, N) = [frechet_dist(Alpha,Beta,Mu) : _ in 1..N].


frechet_dist_pdf(Alpha, Beta, Mu, X) = Res =>
  if X > Mu then
    Res = exp(-((X-Mu)/Beta)**(-Alpha)) * Alpha * ((X-Mu)/Beta)**(-1-Alpha) / Beta
  else
    Res = 0
  end.

frechet_dist_cdf(Alpha, Beta, Mu, X) = Res =>
  if X > Mu then
    Res = exp(-((X-Mu)/Beta)**(-Alpha))
  else
    Res = 0
  end.

frechet_dist_quantile(Alpha, Beta, Mu, X) = Res =>
  if X >= 0, X <= 1 then
    if X > 0, X < 1 then
      Res = Mu + Beta * (-log(X))**(-1/Alpha)
    else
      Res = Mu
    end
  else
    throw($error('frechet_dist_quantile: X >= 0, X < 1 '))      
  end.

frechet_dist_mean(Alpha, Beta, Mu) = Res =>
  if 1 < Alpha then
    Res = Mu + Beta*gamma_func(1-1/Alpha)
  else
    throw($error('frechet_dist_mean: 1 < Alpha'))        
  end.

frechet_dist_variance(Alpha, Beta, Mu) = Res =>
  if Alpha > 2 then
    Res = Beta**2 * (gamma_func(1-2/Alpha) - gamma_func(1-1/Alpha)**2)
  else
    throw($error('frechet_dist_variance: Alpha > 2'))        
  end.


/*
  Max stable distribution

  max_stable_dist(Mu,Sigma,Xi)
  
  This is compliant with Mathematica:
  * MaxStableDistribution[Mu,Sigma,Xi] 
    Represents a generalized maximum extreme value distribution with location 
    parameter Mu, scale parameter Sigma, and shape parameter Xi.

  MaxStableDistribution is also known as Fisher\[Dash]Tippett distribution.

  The generalized maximum extreme value distribution gives the asymptotic 
  distribution of the maximum value in a sample from a distribution such as 
  the normal, Cauchy, or beta distribution.


*/

max_stable_dist(Mu, Sigma, Xi) = max_stable_dist_quantile(Mu, Sigma, Xi, U) =>
  U = u01().
max_stable_dist_n(Mu, Sigma, Xi, N) = [max_stable_dist(Mu,Sigma,Xi) : _ in 1..N].

max_stable_dist_pdf(Mu, Sigma, Xi, X) = Res =>
  if Xi == 0 then
    Res = exp( -(exp( (Mu-X)/Sigma)) - ((X-Mu)/Sigma)) / Sigma
  elseif Xi != 0, 1 + (X-Mu)*Xi/Sigma > 0 then
    Res = (exp(-(1+ ((X-Mu)*Xi/Sigma))**(-1/Xi)) * (1+(X-Mu)*Xi/Sigma)**(-1-1/Xi)) / Sigma
  else
    Res = 0
  end.


max_stable_dist_cdf(Mu, Sigma, Xi, X) = Res =>
  if Xi == 0 then
    Res = exp(-exp( (Mu-X)/Sigma))
  elseif Xi != 0, 1 + ((X-Mu)*Xi/Sigma) > 0 then
    Res = exp(-(1+ ((X-Mu)*Xi/Sigma))**(-1/Xi))
  elseif Xi != 0, 1 + ((X-Mu)*Xi/Sigma) <= 0 then
    Res = 0
  else
    Res = 1
  end.

max_stable_dist_quantile(Mu, Sigma, Xi, X) = Res =>
  if X >= 0, X <= 1 then
    if Xi == 0, X > 0, X < 1 then
      Res = Mu - Sigma*log(-log(X)) / Xi
    elseif Xi != 0, X > 0, X < 1 then
      % Res = Mu - (Sigma*(1-(-log(X)**(-Xi)))) / Xi
      Res = Mu - (Sigma*(1-safe_pow(-log(X),-Xi))) / Xi
    elseif ((X <= 0, Xi > 0) ; (X >= 1, Xi < 0)) then
      Res = Mu - Sigma/Xi
    else
      throw($error('max_stable_quantile: parameter error'))        
    end
  else
     throw($error('max_stable_quantile: X >= 0, X <= 1'))        
  end.
  
max_stable_dist_mean(Mu, Sigma, Xi) = Res =>
  if Xi == 0 then
    Res = Mu + euler_gamma() * Sigma
  elseif Xi != 0, Xi < 1 then
    Res = (Mu*Xi -Sigma + Sigma*gamma_func(1-Xi)) / Xi
  else
    throw($error('max_stable_mean: parameter error'))        
  end.


max_stable_dist_variance(Mu, Sigma, Xi) = Res =>
  if Xi == 0 then
    Res = pi**2 * Sigma**2/6
  elseif Xi != 0, 2*Xi < 1 then
    Res = (Sigma**2 * (gamma_func(1-2*Xi) - gamma_func(1-Xi)**2)) / Xi**2
  else
    throw($error('max_stable_variance: parameter error'))        
  end.



/*

  Min stable distribution

  min_stable_dist(Mu,Sigma,Xi)
  
  This is compliant with Mathematica:
  * MinStableDistribution[Mu,Sigma,Xi] 
    Represents a generalized maximum extreme value distribution with location 
    parameter Mu, scale parameter Sigma, and shape parameter Xi.

  MinStableDistribution is also known as Fisher\[Dash]Tippett distribution.

  The generalized minimum extreme value distribution gives the asymptotic 
  distribution of the minimum value in a sample from a distribution such as the normal, 
  Cauchy, or beta distribution.


*/


min_stable_dist(Mu, Sigma, Xi) = min_stable_dist_quantile(Mu, Sigma, Xi, U) =>
  U = u01().
min_stable_dist_n(Mu, Sigma, Xi, N) = [min_stable_dist(Mu,Sigma,Xi) : _ in 1..N].

min_stable_dist_pdf(Mu, Sigma, Xi, X) = Res =>
  if Xi == 0 then
    Res = exp( -(exp( (X-Mu)/Sigma)) - ((Mu-X)/Sigma)) / Sigma
  elseif Xi != 0, 1 + (Mu-X)*Xi/Sigma > 0 then
    Res = (exp(-(1+ ((Mu-X)*Xi/Sigma))**(-1/Xi)) * (1+(Mu-X)*Xi/Sigma)**(-1-1/Xi)) / Sigma
  else
    Res = 0
  end.


min_stable_dist_cdf(Mu, Sigma, Xi, X) = Res =>
  if Xi == 0 then
    Res = 1 - exp(-exp( (X-Mu)/Sigma))
  elseif Xi != 0, 1 + ((Mu-X)*Xi/Sigma) > 0 then
    Res = 1-exp(-(1+ ((Mu-X)*Xi/Sigma))**(-1/Xi))
  elseif Xi > 0, 1 + ((Mu-X)*Xi/Sigma) <= 0 then
    Res = 0
  else
    Res = 1
  end.

min_stable_dist_quantile(Mu, Sigma, Xi, X) = Res =>
  if X >= 0, X <= 1 then
    if Xi == 0, X > 0, X < 1 then
      Res = Mu + Sigma*log(-log(1-X)) / Xi
    elseif Xi != 0, X > 0, X < 1 then
      % Res = Mu - (Sigma*(1-(-log(X)**(-Xi)))) / Xi
      Res = Mu + (Sigma*(1-safe_pow(-log(1-X),-Xi))) / Xi
    elseif ((X <= 0, Xi > 0) ; (X >= 1, Xi < 0)) then
      Res = Mu + Sigma/Xi
    else
      throw($error('min_stable_quantile: parameter error'))        
    end
  else
     throw($error('min_stable_quantile: X >= 0, X <= 1'))        
  end.
  
min_stable_dist_mean(Mu, Sigma, Xi) = Res =>
  if Xi == 0 then
    Res = Mu - euler_gamma() * Sigma
  elseif Xi != 0, Xi < 1 then
    Res = (Mu*Xi + Sigma - Sigma*gamma_func(1-Xi)) / Xi
  else
    throw($error('min_stable_mean: parameter error'))        
  end.


min_stable_dist_variance(Mu, Sigma, Xi) = Res =>
  if Xi == 0 then
    Res = pi**2 * Sigma**2/6
  elseif Xi != 0, 2*Xi < 1 then
    Res = (Sigma**2 * (gamma_func(1-2*Xi) - gamma_func(1-Xi)**2)) / Xi**2
  else
    throw($error('min_stable_variance: parameter error'))        
  end.



/* 
  Order statistics: Estimator of M

  From Siegrist "Probability Mathematical Statisics and Stochastic Processes"  
  """
  Estimators of m Based on Order Statistics

  Suppose that the population size m is unknown. In this subsection we consider estimators of m 
  constructed from the various order statistics.

  For i ∈ {1, 2, … , n} , the following statistic is an unbiased estimator of m :
         (n + 1)
   Ui =  ------* X(i) - 1
         i

   E(Ui)=m 
  """


  See ppl_order_statistics_estimator_of_m.pi

*/

% 
% Return the i'th smallest element in list xs
% Use Sorted = true if the list is already sorted (increasing)
%
ith_smallest(Xs,I) = ith_smallest(Xs,I,false).
ith_smallest(Xs,I,Sorted) = Res =>
  Xs1 = Xs,
  if Sorted == false then
     Xs1 := Xs.sort
  end,
  Res = Xs1[I].

%  U(i) estimator
order_statistics_m_estimator_u(Xs,I) = Res =>
  N = Xs.len,
  Res = (N+1)/I * ith_smallest(Xs,I) - 1.

order_statistics_m_estimator_u_all(Xs) = Res =>
  Res = [ order_statistics_m_estimator_u(Xs,I) : I in 1..Xs.len].

%  V(i) estimator
%  2*M-1
%  where M = average xs
order_statistics_m_estimator_v(Xs) = Res =>
  Res = 2 * Xs.avg - 1.



/*
  Order statistics for continuous distributions

  From Siegrist "Probability Mathematical Statisics and Stochastic Processes",
  section "Distribution of the k th order statistic"

  From Siegrist "Probability Mathematical Statisics and Stochastic Processes",
  Chapter "6.6: Order Statistics"
  """
  Suppose that x is a real-valued variable for a population and that x = (x1,x2, … , xn ) 
  are the observed values of a sample of size n corresponding to this variable. 
  The order statistic of rank k is the k th smallest value in the data set, and is usually 
  denoted xn:k . To emphasize the dependence on the sample size, another common notation 
  is x(k) .
  """

  Wikipedia https://en.wikipedia.org/wiki/Order_statistic

*/

% PDF:
% table
order_statistics_continuous_pdf(PDF, CDF, N, R, X) = Res =>
  PDFX = apply(PDF,X),
  CDFX = apply(CDF,X),
  Res = (factorialf(N)  / (factorialf(R-1) * factorialf(N-R))) *
        PDFX *
        safe_pow(CDFX,R-1) *
        safe_pow(1-CDFX,N-R).

% CDF:
% table
order_statistics_continuous_cdf(F, N, K, X) = Res =>
  if K > N then
    throw($error('order_statistics_continuous_cdf  K <= N)'))
  end,
  S = 0,
  foreach(J in K..N)
    T = apply(F,X),
    S := S + binomialf(N,J) * safe_pow(T,J) * safe_pow(1-T,N-J)
  end,
  Res = S.

% Quantile
% table
order_statistics_continuous_quantile(CDF, N, R, Q) = order_statistics_continuous_quantile(CDF, N, R, Q, 1).
order_statistics_continuous_quantile(CDF, N, R, Q, Prec) = Res =>
  T = 0,
  OK = false,
  while(OK == false)
    V = order_statistics_continuous_cdf(CDF, N, R, T),
    % println(T=V),    
    if V >= Q then
      OK := T
    else
     T := T + Prec
    end
  end,
  Res = OK.

% The median (using quantiles)
order_statistics_continuous_median(CDF, N, R) =
   order_statistics_continuous_median(CDF, N, R, 1).
order_statistics_continuous_median(CDF, N, R, Prec) = order_statistics_continuous_quantile(CDF, N, R, 0.5, Prec).



/*
  Order statistics discrete without replacement

  From Siegrist "Probability Mathematical Statisics and Stochastic Processes",
  Chapter "12.4: Order Statistics"
  """
  Suppose that the objects in our population are numbered from 1 to m , so that 
  D = {1, 2, … , m} . 
  For example, the population might consist of manufactured items, and the labels 
  might correspond to serial numbers. As in the basic sampling model we select
  n objects at random, without replacement from D .

  ...

  The probability density function of X(i) is

               binom(x-1,i-1) * binom(m-x,n-i)
  P(X(i)=x) =  ...............................
                   binom(m n)
  x in (i,i+1,...m-m+i)
  """

  M: Number of (ordered) objects
  N: Selecting N objects
  I: Values of The I'th object

*/

% Random generation
order_statistics_without_replacement(M, N, I) = order_statistics_without_replacement_quantile(M, N, I, U) =>
  U = uniform(0,1).
order_statistics_without_replacement_n(M, N, I, Num) = [order_statistics_without_replacement(M, N, I) : _ in 1..Num].

order_statistics_without_replacement_pdf(M, N, I, X) = Res =>
  Res = binomialf_float(X-1,I-1) * binomialf_float(M-X,N-I) / binomialf_float(M,N).

order_statistics_without_replacement_cdf(M, N, I, X) = Res =>
  Res = sum([order_statistics_without_replacement_pdf(M, N, I, K) : K in 1..X]).

order_statistics_without_replacement_quantile(M, N, I, Q) = Res =>
  Val = _,
  foreach(X in 1..(M-N)+I,break(nonvar(Val)))
    V = order_statistics_without_replacement_cdf(M, N, I, X),
    if V >= Q then
      Val := X
    end
  end,
  Res = Val.

order_statistics_without_replacement_mean(M, N, I) = Res =>
  Res =  I * ((M+1) / (N+1)).

order_statistics_without_replacement_variance(M, N, I) = Res =>
  Res = I *
        (1 + (N-I)) *
        (  (M+1) * (M-N) / ((N+1)**2 * (N+2))).




/*
  Order statistics with replacement (discrete)

  This is for order statistics with replacement for discrete distributions. 
  For continous distribution, see order_statistics_continuous_dist

  From Siegrist "Probability Mathematical Statisics and Stochastic Processes",
  Chapter "6.6: Order Statistics"
  """
  Suppose that x is a real-valued variable for a population and that x = (x1,x2, … , xn ) 
  are the observed values of a sample of size n corresponding to this variable. 
  The order statistic of rank k is the k th smallest value in the data set, and is usually 
  denoted xn:k . To emphasize the dependence on the sample size, another common notation 
  is x(k) .
  """

  From https://en.wikipedia.org/wiki/Order_statistic
  "Dealing with discrete variables"
  Three values are needed
   p1 = P(X < x) = F(x) - f(x)
   p2 = P(X = x) = f(x)
   p3 = P(X > x) = 1- F(x)
  The CDF:
   P(X(k) <= x) = sum(j=0..n-k,p3^j*(p1+p2)^(n-j))

  For discrete_uniform_dist(1,6), e.g. throwing a die:
  PDF: discrete_uniform_dist_pdf(1,6)
  CDF: discrete_uniform_dist_cdf(1,6)  

  We see the probability (PDF and CDF) that 1 is the Nth smallest 
  value:
   
  PDF:
    order_statistics_with_replacement_discrete_pdf(PDF,CDF,4,1,2): 
    0.51774691358024694

  CDF: 
   order_statistics_with_replacement_discrete_leq_cdf(PDF,CDF,4,1,2): 
   0.80246913580246926

  This corresponds to Mathematica's OrderDistribution which 
  is for order statistics with replacement (for <=), e.g.
     dist = OrderDistribution[{DiscreteUniformDistribution[{1, 6}], 4}, 1]
     PDF[dist, 2]
     -> 0.284722
     CDF[dist, 2]
     -> 0.802469

  See ppl_order_statistics_with_replacement_discrete_dist.pi
*/

/*
  "Note that the probability mass function of X(k) is just the difference of these values,"
  Note that this is for taking n (all) samples. 
*/
order_statistics_with_replacement_discrete_pdf(PDF,CDF, N,K,X) = Res =>
  PDFX = apply(PDF,X),
  CDFX = apply(CDF,X),
  P1 = CDFX - PDFX,
  P2 = PDFX,
  P3 = 1 - CDFX,
  Res = sum([ binomialf(N,J) *
              (
                (P3**J * (P1+P2)**(N-J)) -
                ((P2+P3)**J * P1**(N-J))
              )              
             : J in 0..N-K]).

/*
  CDF: for P(X(k) <= x). 
  Note that it's <=, i.e. less than or equal
*/
order_statistics_with_replacement_discrete_leq_cdf(PDF,CDF,N,K,X) = Res =>
  PDFX = apply(PDF,X),
  CDFX = apply(CDF,X),
  P1 = CDFX - PDFX,
  P2 = PDFX,
  P3 = 1 - CDFX,
  Res = sum([ binomialf(N,J) *
              P3**J *
              (P1+P2)**(N-J)
             : J in 0..N-K]).

/*
  For P(X(k) <= x)
  Note that it's <=, i.e. less than

  Hmm, this does give 0 for x = 1, and for n it's not fully 1!
  But it does make some sense, since there can be no values < 1 and there are
  not a possibility of getting 6, but for n+1 (e.g. 7) the CDF will become 1.
  I'm not sure how useful this really is, perhaps for other distributions
  such as poisson?

*/
order_statistics_with_replacement_discrete_lt_cdf(PDF,CDF,N,K,X) = Res =>
  PDFX = apply(PDF,X),
  CDFX = apply(CDF,X),
  P1 = CDFX - PDFX,
  P2 = PDFX,
  P3 = 1 - CDFX,
  Res = sum([ binomialf(N,J) *
              (P2 + P3)**J *
              P1**(N-J)
             : J in 0..N-K]).


order_statistics_with_replacement_discrete_leq_quantile(PDF,CDF,N,K,X) = Res =>
  Val = _,
  I = 0,
  while(var(Val))
    V = order_statistics_with_replacement_discrete_leq_cdf(PDF,CDF,N,K,I),
    if V >= X then
      Val := I
    else
      I := I + 1
    end
  end,
  Res = Val.


/*
  Cauchy distribution

  https://en.wikipedia.org/wiki/Cauchy_distribution
  and from Mathematica's CauchyDistribution

  cauchy_dist(A,B)
  Represents a Cauchy distribution with location parameter A and scale parameter B.
  
  cauchy_dist()
  Represents a Cauchy distribution with location parameter 0 and scale parameter 1.

*/ 
cauchy_dist(A,B) = Res =>
  U = uniform(0,1),
  Res = A + B * tan(pi * U - 1/2).
cauchy_dist_n(A,B,N) = [cauchy_dist(A,B) : _ in 1..N].


cauchy_dist_pdf(A,B,X) = 1 / (pi * B * (1 + ((X-A)**2) / B**2)).

cauchy_dist_cdf(A,B,X) = 1/2 + atan( (X-A)/B ) / pi.

cauchy_dist_quantile(A,B,X) = Res =>
  if X > 0, X < 1 then
    Res = A + B * tan(pi * (X-1/2))
  else
    throw($error('cauchy_dist_cdf: X > 0, X < 1'))
  end.

cauchy_dist_mean(A,B) = indeterminate.
cauchy_dist_variance(A,B) = indeterminate.


/*
  Benford distribution

  From Mathematica BenfordDistribution[b]
  represents a Benford distribution with base parameter b.

*/

benford_dist(B) = benford_dist_quantile(B,U) =>
  U = uniform(0,1).
% benford_dist(B) = ceiling(B**U) =>
%  U = uniform(0,1).
benford_dist_n(B,N) = [benford_dist(B) : _ in 1..N].

benford_dist_pdf(B, X) = Res =>
  if 1 <= X, X < B then
    Res = log(1+1/X) / log(B)
  else
    Res = 0
  end.

benford_dist_cdf(B, X) = Res =>
  if 1 <= X, X < B then
    Res = log(1+floor(X))/log(B)
  elseif X < 1 then
    Res = 0
  else
    Res = 1
  end.


benford_dist_quantile(B, X) = Res =>
  if 0 <= X, X <= 1 then
     if 0 < X, X < 1 then
       Res = ceiling(B**X) -1
     elseif X <= 0 then
       Res = 1
     else
       Res = B - 1
     end
  end.

benford_dist_mean(B) = B - lgamma(1+B) / log(B).

% benford_dist_variance(B) % not implemented


/*

  Pólya Distribution

  https://www.randomservices.org/random/urn/Polya.html
  """
  An urn initially contains a red and b green balls, where a and b  are positive integers. 
  At each discrete time (trial), a ball is selected from the urn and then returned to the urn 
  along with c new balls of the same color. The random process is known as Pólya's urn process, 
  named for George Pólya.

  ...

  In terms of the colors of the selected balls, Pólya's urn scheme generalizes the standard 
  models of sampling with and without replacement.

  c = 0: corresponds to sampling with replacement.
  c = -1: corresponds to sampling without replacement.
  """

 See ppl_beta_binomial_urn_model.pi for an example

*/
polya_dist(N, A, B, C) = beta_binomial_dist(N, A/C, B/C).
polya_dist_n(N, A, B, C, Num) = [polya_dist(N, A, B, C) : _ in 1..Num].


/*
  From Mathematica BetaBinomialDistribution:
  """
  The distribution models an urn scheme. An urn contains w white balls and b black balls. 
  When a ball is drawn it is returned to the urn together with c additional balls of 
  the same color. The distribution gives the probability of drawing k white balls in n draws.
  """

  See ppl_beta_binomial_urn_model.pi for an example
*/
polya_eggenberg_dist(N,W,B,C) = beta_binomial_dist(N,W/C,B/C).
polya_eggenberg_dist_n(N,W,B,C,Num) = [polya_eggenberg_dist(N,W,B,C) : _ in 1..Num].



/*
  Chi distribution

  From Handbook on probability distributions
  page 78
  """
  Take the square root of a chi-squared random variable.
  """

  From Mathematica:
  ChiDistribution[Nu] 
  Represents a Chi distribution with Nu degrees of freedom.

*/

/* ---------------------------------------------------------------------------
Chi distribution
Chi(Nu)

Description:
    The Chi distribution is a continuous probability distribution that
    describes the square root of a Chi-square distributed variable.

        If Y ~ ChiSquare(Nu),
        then X = sqrt(Y)
        follows a Chi(Nu) distribution.

    It has one parameter:
        Nu : degrees of freedom (> 0)

    The Chi distribution is a special case of the Gamma family
    (through the Chi-square) and is closely related to the
    Maxwell–Boltzmann distribution.

    It often arises as the distribution of the Euclidean norm (length)
    of a vector of Nu independent standard normal random variables.

        If Z₁, Z₂, …, Z_Nu ~ Normal(0,1),
        then X = sqrt(Z₁² + Z₂² + … + Z_Nu²) ~ Chi(Nu)

-------------------------------------------------------------------------------
Mathematical definitions

    PDF:
        f(x; Nu) =
            [ 2^(1 - Nu/2) / Γ(Nu/2) ] * x^(Nu - 1) * e^(-x² / 2)
        for x > 0, and f(x) = 0 for x <= 0.

    CDF:
        F(x; Nu) = Γ_reg(Nu/2, x²/2)
        where Γ_reg(a, z) is the regularized lower incomplete gamma function.

    Quantile (inverse CDF):
        Q(p; Nu) = sqrt( 2 * Γ_reg⁻¹(Nu/2, p) )

    Mean:
        E[X] = sqrt(2) * Γ((Nu+1)/2) / Γ(Nu/2)

    Variance:
        Var[X] = Nu - E[X]²

-------------------------------------------------------------------------------
Example usage:

    % Random sampling
    println(chi_dist(5))
    println(chi_dist(10))

    % PDF
    println(chi_dist_pdf(5, 2.0))        % ≈ 0.2076
    println(chi_dist_pdf(10, 3.0))       % ≈ 0.1760

    % CDF
    println(chi_dist_cdf(5, 2.0))        % ≈ 0.286
    println(chi_dist_cdf(10, 3.0))       % ≈ 0.400

    % Quantile
    println(chi_dist_quantile(5, 0.95))  % ≈ 3.33
    println(chi_dist_quantile(10, 0.95)) % ≈ 4.30

    % Mean and variance
    println(chi_dist_mean(5))            % ≈ 2.13
    println(chi_dist_variance(5))        % ≈ 0.69

-------------------------------------------------------------------------------
Applications:

  • Euclidean norm of Gaussian vectors:
        Models the length of random vectors in high-dimensional normal space.

  • Signal and noise analysis:
        Used in statistics of RMS amplitude, magnitude, or noise power.

  • Maxwell–Boltzmann distribution:
        The Chi(3) distribution describes the speed of particles in a
        3D ideal gas.

  • Goodness-of-fit and confidence intervals:
        Appears in the denominator of t- and F-distributions through
        square-rooted Chi-square terms.

  • Simulation and probabilistic geometry:
        Useful when generating random radii in multidimensional normal spaces.

-------------------------------------------------------------------------------
Implementation notes:

  - Random generation via sqrt(ChiSquare(Nu))
  - PDF computed via gamma_func and exponential term
  - CDF and quantile use gamma_regularized/2 and inverse_gamma_regularized/2
  - Mean and variance follow from Gamma function ratios
  - Numerical precision typically ~1e-15 for moderate degrees of freedom

--------------------------------------------------------------------------- */


chi_dist(Nu) = sqrt(sum([normal_dist(0,1)**2  : _ in 1..Nu])).
chi_dist_n(Nu, N) = [chi_dist(Nu) : _ in 1..N].

% chi_dist_pdf(Nu,X) = Res =>
%   if X > 0 then
%     Res = (2**(1-(Nu/2)) * exp(-(X^2)/2) * X**(Nu-1)) / gamma_func(Nu/2)
%   else
%     Res = 0
%   end.
%
% Chi distribution PDF
% chi_dist_pdf(Nu, X)
%   Nu : degrees of freedom (> 0)
%   X  : point (>= 0)
%
chi_dist_pdf(Nu, X) = Pdf =>
    if Nu =< 0.0 then
        throw($error('chi_dist_pdf: Nu must be > 0'))
    elseif X < 0.0 then
        Pdf = 0.0
    elseif X =:= 0.0 then
        if Nu < 1.0 then
            % Diverges to +inf for Nu < 1
            Pdf = 1.0/0.0
        elseif Nu =:= 1.0 then
            Pdf = sqrt(2.0/pi)      % half-normal at 0
        else
            Pdf = 0.0               % Nu > 1
        end
    else
        % f(x;nu) = 1/(2^{nu/2 - 1} * Gamma(nu/2)) * x^{nu-1} * exp(-x^2/2)
        LogC = -((Nu/2.0 - 1.0) * log(2.0) + lgamma(Nu/2.0)),
        LogF = LogC + (Nu - 1.0) * log(X) - (X*X)/2.0,
        Pdf = exp(LogF)
    end.

%
% Chi distribution CDF
% chi_dist_cdf(Nu, X)
%   Nu : degrees of freedom (> 0)
%   X  : value to evaluate (>= 0)
%
chi_dist_cdf(Nu, X) = Res =>
    if Nu =< 0.0 then
        throw($error('chi_dist_cdf: Nu must be > 0'))
    elseif X =< 0.0 then
        Res = 0.0
    elseif X >= 1.0e10 then
        Res = 1.0
    else
        Res = gamma_regularized(Nu / 2.0, (X * X) / 2.0)
    end.


%
% Chi distribution quantile (inverse CDF)
% chi_dist_quantile(Nu, P)
%   Nu : degrees of freedom (> 0)
%   P  : probability in [0,1]
%
chi_dist_quantile(Nu, P) = Q =>
    if Nu =< 0.0 then
        throw($error('chi_dist_quantile: Nu must be > 0'))
    elseif P =< 0.0 then
        Q = 0.0
    elseif P >= 1.0 then
        Q = 1.0e10
    else
        X = inverse_gamma_regularized(Nu / 2.0, P),
        Q = sqrt(2.0 * X)
    end.



%
% Chi distribution mean
% chi_dist_mean(Nu)
%   Nu : degrees of freedom (> 0)
%
chi_dist_mean(Nu) = M =>
    if Nu =< 0.0 then
        throw($error('chi_dist_mean: Nu must be > 0'))
    else
        % μ = sqrt(2) * Γ((ν+1)/2) / Γ(ν/2)
        % Use lgamma for numerical stability.
        M = sqrt(2.0) * exp(lgamma((Nu + 1.0) / 2.0) - lgamma(Nu / 2.0))
    end.

%
% Chi distribution variance
% chi_dist_variance(Nu)
%   Nu : degrees of freedom (> 0)
%
chi_dist_variance(Nu) = V =>
    if Nu =< 0.0 then
        throw($error('chi_dist_variance: Nu must be > 0'))
    else
        % Var = ν - μ^2
        Mu = chi_dist_mean(Nu),
        V = Nu - Mu * Mu
    end.



/* ---------------------------------------------------------------------------
Chi-square distribution
ChiSquare(Nu)

Description:
    The Chi-square (χ²) distribution is a continuous probability distribution
    that arises as the sum of squares of independent standard normal random
    variables.

        If Z₁, Z₂, …, Z_Nu ~ Normal(0,1), then
            X = Σ (Zᵢ²)
        follows a ChiSquare(Nu) distribution.

    It is parameterized by:
        Nu : degrees of freedom (> 0)

    The Chi-square distribution is a special case of the Gamma distribution:
        ChiSquare(Nu) = Gamma( shape = Nu/2, scale = 2 )

    and it forms the basis of several other important distributions:
        - Student’s t  :  t(Nu) = Z / sqrt(ChiSquare(Nu) / Nu)
        - F-distribution:  F(Nu₁, Nu₂) = (ChiSquare(Nu₁)/Nu₁) / (ChiSquare(Nu₂)/Nu₂)

-------------------------------------------------------------------------------
Mathematical definitions

    PDF:
        f(x; Nu) =
            [ 1 / (2^(Nu/2) * Γ(Nu/2)) ] * x^(Nu/2 - 1) * e^(-x/2)
        for x > 0, and f(x) = 0 for x <= 0.

    CDF:
        F(x; Nu) = Γ_reg(Nu/2, x/2)
        where Γ_reg(a, z) is the regularized lower incomplete gamma function.

    Quantile (inverse CDF):
        Q(p; Nu) = 2 * Γ_reg⁻¹(Nu/2, p)

    Mean:
        E[X] = Nu

    Variance:
        Var[X] = 2 * Nu

-------------------------------------------------------------------------------
Example usage:

    % Random sampling
    println(chi_square_dist(5))
    println(chi_square_dist(10))

    % PDF
    println(chi_square_dist_pdf(5, 2.0))        % ≈ 0.2076
    println(chi_square_dist_pdf(10, 10.0))      % ≈ 0.1251

    % CDF
    println(chi_square_dist_cdf(5, 2.0))        % ≈ 0.050
    println(chi_square_dist_cdf(10, 10.0))      % ≈ 0.5595

    % Quantile
    println(chi_square_dist_quantile(5, 0.95))  % ≈ 11.07
    println(chi_square_dist_quantile(10, 0.99)) % ≈ 23.21

    % Mean and variance
    println(chi_square_dist_mean(5))            % 5.0
    println(chi_square_dist_variance(5))        % 10.0

-------------------------------------------------------------------------------
Applications:

  • Hypothesis testing for variance:
        Used in tests where the population variance σ² is compared to a
        hypothesized value based on a sample variance s²:
            χ² = (Nu * s²) / σ₀²

  • Goodness-of-fit tests:
        Measures how well observed frequencies match expected frequencies.

  • Independence tests:
        Used in contingency tables (χ² test of independence).

  • Building blocks for other distributions:
        Underlies t- and F-distributions and many likelihood-ratio tests.

  • Monte Carlo and Bayesian modeling:
        Used as a prior or likelihood for precision/variance parameters.

-------------------------------------------------------------------------------
Implementation notes:

  - Random generation via Gamma(Nu/2, 2)
  - PDF uses log-domain evaluation for numerical stability
  - CDF and quantile computed via gamma_regularized/2 and inverse_gamma_regularized/2
  - Mean = Nu; Variance = 2 * Nu
  - Fully consistent with standard numerical precision (~1e-15)

--------------------------------------------------------------------------- */


%
% Chi-square distribution random generator
% chi_square_dist(Nu)
%   Nu : degrees of freedom (> 0)
%
chi_square_dist(Nu) = Y =>
    if Nu =< 0.0 then
        throw($error('chi_square_dist: Nu must be > 0'))
    else
        % Chi-square(nu) = Gamma(shape=nu/2, scale=2)
        Y = gamma_dist(Nu / 2.0, 2.0)
    end.

chi_square_dist_n(Nu, N) = [chi_square_dist(Nu) : _ in 1..N].

%
% Chi-square distribution PDF
% chi_square_dist_pdf(Nu, X)
%   Nu : degrees of freedom (> 0)
%   X  : value (>= 0)
%
chi_square_dist_pdf(Nu, X) = Pdf =>
    if Nu =< 0.0 then
        throw($error('chi_square_dist_pdf: Nu must be > 0'))
    elseif X < 0.0 then
        Pdf = 0.0
    elseif X =:= 0.0 then
        if Nu < 2.0 then
            % Diverges to infinity when Nu < 2
            Pdf = 1.0e300 % Pdf = 1.0/0.0
        elseif Nu =:= 2.0 then
            % f(0;2) = 0.5
            Pdf = 0.5
        else
            Pdf = 0.0
        end
    else
        % log-form for stability:
        % log(f) = (Nu/2 - 1)*log(X) - X/2 - (Nu/2)*log(2) - lgamma(Nu/2)
        LogF = (Nu/2.0 - 1.0) * log(X)
               - X / 2.0
               - (Nu/2.0) * log(2.0)
               - lgamma(Nu / 2.0),
        Pdf = exp(LogF)
    end.


%
% Chi-square distribution CDF
% chi_square_dist_cdf(Nu, X)
%   Nu : degrees of freedom (> 0)
%   X  : value (>= 0)
%
chi_square_dist_cdf(Nu, X) = Res =>
    if Nu =< 0.0 then
        throw($error('chi_square_dist_cdf: Nu must be > 0'))
    elseif X =< 0.0 then
        Res = 0.0
    elseif X >= 1.0e10 then
        Res = 1.0
    else
        Res = gamma_regularized(Nu / 2.0, X / 2.0)
    end.


%
% Chi-square distribution quantile (inverse CDF)
% chi_square_dist_quantile(Nu, P)
%   Nu : degrees of freedom (> 0)
%   P  : probability in [0,1]
%
chi_square_dist_quantile(Nu, P) = Q =>
    if Nu =< 0.0 then
        throw($error('chi_square_dist_quantile: Nu must be > 0'))
    elseif P =< 0.0 then
        Q = 0.0
    elseif P >= 1.0 then
        Q = 1.0e10
    else
        X = inverse_gamma_regularized(Nu / 2.0, P),
        Q = 2.0 * X
    end.


%
% Chi-square distribution mean
% chi_square_dist_mean(Nu)
%   Nu : degrees of freedom (> 0)
%
chi_square_dist_mean(Nu) = M =>
    if Nu =< 0.0 then
        throw($error('chi_square_dist_mean: Nu must be > 0'))
    else
        M = Nu
    end.


%
% Chi-square distribution variance
% chi_square_dist_variance(Nu)
%   Nu : degrees of freedom (> 0)
%
chi_square_dist_variance(Nu) = V =>
    if Nu =< 0.0 then
        throw($error('chi_square_dist_variance: Nu must be > 0'))
    else
        V = 2.0 * Nu
    end.

/*
  Chi square inverse distribution

  From Handbook on probability distributions
  page 82
  """
  Simply inverse a chi-squared random variable
  """

  Only random generation

*/
chi_square_inverse_dist(K) = 1 / chi_square_dist(K).
chi_square_inverse_dist_n(K,N) = [chi_square_inverse_dist(K) : _ in 1..N].


/*
  Erlang distribution

  From Mathematica

  ErlangDistribution[K,Lambda] 
  represents the Erlang distribution with shape parameter K and rate Lambda.

*/

/*
  From Handbook on probability distributions
  page 64
  """
  The algorithm is very easy simulate independently d random variables
  exponentially E(lambda_j) distributed and sum them.
  """
*/
erlang_dist(K,Lambda) = [exponential_dist(Lambda) : _ in 1..K].sum.

% The quantile is too slow for this...
% erlang_dist(K,Lambda) = Res =>
%   U = u01(),
%   Res = erlang_dist_quantile(K,Lambda,U),
%   if Res > 100_000 then
%     println(large_value=[u=U,res=Res])
%   end.

erlang_dist_n(K,Lambda,N) = [erlang_dist(K,Lambda) : _ in 1..N].

erlang_dist_pdf(K,Lambda,X) = Res =>
  if X > 0 then
    Res = exp(-X*Lambda) * X**(K-1) * Lambda**K / gamma_func(K)
  else
    Res = 0
  end.

erlang_dist_cdf(K,Lambda,X) = Res =>
  if X > 0 then
    Res = gamma_regularized(K,X*Lambda)
  else
    Res = 0
  end.

/*
% This does not work:
% Picat> erlang_dist_quantile(10,3,0.99)=X 
% X = 24265957.713672984391451
% Picat> erlang_dist_quantile(10,3,0.999)=X
% X = 1132.01258265243473
%
erlang_dist_quantile(K,Lambda,X) = Res =>
  if 0 <= X, X < 1 then
    if X > 0 then
      Res = inverse_gamma_regularized(K,X) / Lambda
    else
     Res = 0
    end
  else
    throw($error('erlang_dist_quantile: 0 <= X, X < 1'))
  end.
*/
% Use binary search instead
erlang_dist_quantile(K,Lambda,X) = Res =>
  Res = binary_search_quantile($erlang_dist_cdf(K,Lambda),X,1.0e-15,0.1,1.0e100).
 
erlang_dist_mean(K,Lambda) = K / Lambda.
erlang_dist_variance(K,Lambda) = K / Lambda**2.


/*
  Inverse exponential

  From Handbook on probability distributions
  page 60
  """
  The algorithm is simply to inverse an exponential variate of parameter
  1/lambda, i.e. (−lambda log(U))−1 for an uniform variable U.
  """

  Only the random generator.

*/
inverse_exponential_dist(Lambda) = - Lambda * log(U) =>
  U = uniform(0,1).
inverse_exponential_dist_n(Lambda,N) = [inverse_exponential_dist(Lambda) : _ in 1..N].


/*

  Shifted_exponential distribution

  From Handbook on probability distributions
  page 60
  """
  The random generation is simple: just add τ to the
  algorithm of exponential distribution
  """

  Only random generator.

*/
shifted_exponential_dist(Lambda,T) = Res =>
  U = uniform(0,1),
  Res = (-1/log(U))/Lambda + T.
shifted_exponential_dist_n(Lambda,T,N) = [shifted_exponential_dist(Lambda,T) : _ in 1..N].


/*
  Kumaraswamy distribution

  From Mathematica:
  KumaraswamyDistribution[Alpha,Beta] 
  represents a Kumaraswamy distribution with shape parameters Alpha and Beta.
  
*/
% kumaraswamy_dist(Alpha,Beta) =  (1-(1-U)**(1/Beta))**(1/Alpha) =>
kumaraswamy_dist(Alpha,Beta) =  safe_pow((1-safe_pow((1-U),(1/Beta))),(1/Alpha)) =>
  U = uniform(0,1).
kumaraswamy_dist_n(Alpha,Beta,N) = [kumaraswamy_dist(Alpha,Beta) : _ in 1..N].


kumaraswamy_dist_pdf(Alpha,Beta,X) = Res =>
  if 0 < X, X < 1 then
    Res = X**(Alpha-1)*(1-X)**(Beta-1)*Alpha*Beta
  else
    Res = 0
  end.

kumaraswamy_dist_cdf(Alpha,Beta,X) = Res =>
  if 0 < X, X < 1 then
    Res = 1 - (1-X**Alpha)**Beta
  elseif X >= 1 then
    Res = 1
  else
    Res = 0    
  end.

kumaraswamy_dist_quantile(Alpha,Beta,X) = Res =>
  if 0 < X, X < 1 then
    Res = (1-(1-X)**(1/Beta))**(1/Alpha)
  elseif X <= 0 then
    Res = 0
  else
    Res = 1
  end.

kumaraswamy_dist_mean(Alpha,Beta) = Beta*beta_func(Beta,1+1/Alpha).
kumaraswamy_dist_variance(Alpha,Beta) = -Beta**2 * beta_func(Beta,1+1/Alpha)**2 +
                                        Beta * beta_func(Beta,1+2/Alpha).



/*
  Triangular distributions
  
  https://en.wikipedia.org/wiki/Triangular_distribution

  From Mathematica:
  TriangularDistribution[{Min,Max}] 
  represents a symmetric triangular statistical distribution giving values between Min and Max. 

  TriangularDistribution[]
  represents a symmetric triangular statistical distribution giving values between 0 and 1.

  TriangularDistribution[{Min,Max},C]
  represents a triangular distribution with mode at C.
*/

/*
  TriangularDistribution[] 
*/

triangular_dist() = triangular_dist(0,1).
triangular_dist_n(N) = [triangular_dist() : _ in 1..N].

triangular_dist_pdf(X) = triangular_dist_pdf(0,1,X).
triangular_dist_cdf(X) = triangular_dist_cdf(0,1,X).
triangular_dist_quantile(X) = triangular_dist_quantile(0,1,X).

triangular_dist_mean() = triangular_dist_mean(0,1).
triangular_dist_variance() = triangular_dist_variance(0,1).

/*
  TriangularDistribution[{Min,Max}] 

  Called triangular_dist
*/
triangular_dist(Min,Max) = triangular_dist_quantile(Min,Max,U) =>
  U = uniform(0,1).
triangular_dist_n(Min,Max,N) = [triangular_dist(Min,Max) : _ in 1..N].


triangular_dist_pdf(Min,Max,X) = Res =>
  if Min <= X, X <= (Max+Min)/2 then
    Res = 4*(X-Min) / (Max-Min)**2
  elseif ((Max+Min)/2 < X, X <= Max) then
    Res = 4*(Max-X) / (Max-Min)**2 
  else
    Res = 0
  end.
  
triangular_dist_cdf(Min,Max,X) = Res =>
  if Min <= X, X <= (Max-Min) / 2 then
    Res = (X-Min)**2 / ( (Max-Min) * ( (Max+Min)/2 -Min ))
  elseif ((Max+Min)/2 < X, X <= Max) then
    Res = 1- ( (Max-X)**2 / ( (Max + (1/2)*(-Max-Min)) * (Max-Min)))
  elseif X > Max then
    Res = 1
  else
    Res = 0
  end.

triangular_dist_quantile(Min,Max,X) = Res =>
  T = (((Max+Min)/ 2) - Min) / (Max-Min),
  if 0 <= X, X <= T then
    Res = Min + sqrt( (Max-Min) * (((Max+Min)/2) - Min) * X )
  elseif 1 >= X, X > T then
    Res = Max - sqrt( (Max + (1/2)*(-Max-Min)) * (Max-Min) * (1-X))
  elseif X < 0 then
    Res = Min
  else
    Res = Max
  end.

triangular_dist_mean(Min,Max) = (Min+Max) / 2.

triangular_dist_variance(Min,Max) = (1/24)*(Max-Min)**2.

/*
  TriangularDistribution[{Min,Max},C]
  represents a triangular distribution with mode at C.

  Called triangular2_dist
*/
triangular2_dist(Min,Max,C) = triangular2_dist_quantile(Min,Max,C,U) =>
  U = uniform(0,1).
triangular2_dist_n(Min,Max,C,N) = [triangular2_dist(Min,Max,C) : _ in 1..N].

triangular2_dist_pdf(Min,Max,C,X) = Res =>
  if Min <= X, X <= C then
    Res = 2*(X-Min) / ((C-Min)*(Max-Min))
  elseif C < X, C <= Max then
    Res = 2*(Max-X) / ((Max-C)*(Max-Min))
  else
    Res = 0
  end.

triangular2_dist_cdf(Min,Max,C,X) = Res =>
  if Min <= X, X <= C then
    Res = (X-Min)**2 / ((C-Min)*(Max-Min))
  elseif C < X, C <= Max then
    Res = 1 - (Max-X)**2 / ((Max-C)*(Max-Min))
  elseif X > Max then
    Res = 1
  else
    Res = 0
  end.

triangular2_dist_quantile(Min,Max,C,X) = Res =>
  if 0 <= X, X <= (C-Min)/(Max-Min) then
    Res = Min + sqrt((C-Min) * (Max-Min) * X)
  elseif 1 >= X, X > (C-Min) / (Max-Min) then
    Res = Max - sqrt( (Max-C) * (Max-Min) * (1-X))
  elseif X < 0 then
    Res = Min
  else
    Res = Max
  end.

triangular2_dist_mean(Min,Max,C) = 1/3 * (C + Max + Min).

triangular2_dist_variance(Min,Max,C) = 1/18 * (C**2 - C*Max + Max**2 - C*Min - Max*Min + Min**2).



/*
  Log gamma distribution

  From Handbook on probability distributions
  page 69
  """
  Simply simulate a gamma G(k, 1) distributed variable and returns a + b log(X).
  """

  From Mathematica:
  LogGammaDistribution[Alpha,Beta,Mu] 
  represents a log-gamma distribution with shape parameters Alpha and Beta and location parameter Mu.
*/
% log_gamma_dist(K,A,B) = A + B * log(gamma_dist(K,1)).
% log_gamma_dist(Alpha,Beta,Mu) = Alpha + Beta * log(gamma_dist(Mu,1)).

log_gamma_dist(Alpha,Beta,Mu) = log_gamma_dist_quantile(Alpha,Beta,Mu,U) =>
  U = u01().
  
log_gamma_dist_n(Alpha,Beta,Mu,N) = [log_gamma_dist(Alpha,Beta,Mu) : _ in 1..N].

log_gamma_dist_pdf(Alpha,Beta,Mu,X) = Res =>
  if X >= Mu then
    Res = Beta**(-Alpha) * safe_pow((1+X-Mu),(- (Beta+1)/Beta)) * safe_pow(log(1+X-Mu),(Alpha-1)) / gamma_func(Alpha)
  else
    Res = 0
  end.

log_gamma_dist_cdf(Alpha,Beta,Mu,X) = Res =>
  if X >= Mu then
    Res = gamma_regularized(Alpha,log(1+X-Mu)/Beta)
  else
    Res = 0
  end.

log_gamma_dist_quantile(Alpha,Beta,Mu,X) = Res =>
  if 0 <= X, X <= 1 then
    if 0 < X, X < 1 then
      Res = -1 + exp(Beta * inverse_gamma_regularized(Alpha,X)) + Mu,
    elseif X <= 0 then
      Res = Mu
    else
      throw($error('log_gamma_dist_quantile: 0 <= X < 1'))
    end
  else
    throw($error('log_gamma_dist_quantile: 0 <= X < 1'))
  end.

log_gamma_dist_mean(Alpha,Beta,Mu) = Res =>
  if Beta < 1 then
    Res = -1 + (1-Beta)**(-Alpha) + Mu
  else
   throw($error('log_gamma_dist_mean: Beta < 1'))
  end.

log_gamma_dist_variance(Alpha,Beta,Mu) = Res =>
  if Beta < 1/2 then
    Res = (1-2*Beta)**(-Alpha) - (1-Beta)**(-2*Alpha)
  else
    throw($error('log_gamma_dist_variance: Beta < 1/12'))
  end.

/* ---------------------------------------------------------------------------
Log-normal distribution
LogNormal(Mu, Sigma)

Description:
    The Log-normal distribution models a random variable whose natural
    logarithm follows a Normal distribution.

        If Y ~ Normal(Mu, Sigma),
        then X = exp(Y)
        follows a LogNormal(Mu, Sigma) distribution.

    The parameters are:
        Mu     : mean of log(X)
        Sigma  : standard deviation of log(X) (> 0)

    The Log-normal distribution is right-skewed and strictly positive.
    It is commonly used to model multiplicative effects, growth processes,
    and data spanning several orders of magnitude.

-------------------------------------------------------------------------------
Mathematical definitions

    PDF:
        f(x; Mu, Sigma) =
            [ 1 / (x * Sigma * sqrt(2π)) ] *
            exp( - (ln(x) - Mu)² / (2 * Sigma²) )

        for x > 0, and f(x) = 0 for x <= 0.

    CDF:
        F(x; Mu, Sigma) =
            Φ( (ln(x) - Mu) / Sigma )
        where Φ is the CDF of the standard Normal distribution.

    Quantile (inverse CDF):
        Q(p; Mu, Sigma) =
            exp( Mu + Sigma * Φ⁻¹(p) )

    Mean:
        E[X] = exp(Mu + Sigma² / 2)

    Variance:
        Var[X] = (exp(Sigma²) - 1) * exp(2*Mu + Sigma²)

-------------------------------------------------------------------------------
Example usage:

    % Random sampling
    println(lognormal_dist_random(0.0, 1.0))
    println(lognormal_dist_random(1.0, 0.5))

    % PDF
    println(lognormal_dist_pdf(0.0, 1.0, 1.0))        % ≈ 0.3989
    println(lognormal_dist_pdf(0.0, 1.0, exp(1.0)))   % ≈ 0.2420

    % CDF
    println(lognormal_dist_cdf(0.0, 1.0, 1.0))        % 0.5
    println(lognormal_dist_cdf(0.0, 1.0, exp(1.0)))   % ≈ 0.8413

    % Quantile
    println(lognormal_dist_quantile(0.0, 1.0, 0.5))   % 1.0
    println(lognormal_dist_quantile(0.0, 1.0, 0.8413))% ≈ 2.718

    % Mean and variance
    println(lognormal_dist_mean(0.0, 1.0))            % ≈ 1.6487
    println(lognormal_dist_variance(0.0, 1.0))        % ≈ 4.6708

-------------------------------------------------------------------------------
Applications:

  • Modeling positive-valued random variables:
        - incomes, house prices, reaction times, biological measurements
        - quantities with multiplicative error structure

  • Finance and economics:
        - models of asset prices, since prices cannot go negative
        - geometric Brownian motion models use Log-normal returns

  • Reliability and lifetime modeling:
        - used for failure times when the logarithm of time-to-failure
          is approximately Normal

  • Environmental and natural sciences:
        - sizes of particles, raindrops, or other positive-valued phenomena

  • Monte Carlo simulations:
        - used when multiplicative randomness or long right tails are required

-------------------------------------------------------------------------------
Implementation notes:

  - Random generation: exp(Normal(Mu, Sigma))
  - PDF evaluated in log-domain for numerical stability
  - CDF and quantile reuse Normal distribution functions
  - Mean and variance computed exactly from analytic formulas
  - Numerical precision typically better than 1e-15

--------------------------------------------------------------------------- */

%
% Log-normal distribution random generator
% lognormal_dist_random(Mu, Sigma)
%   Mu     : mean of log(X)
%   Sigma  : standard deviation of log(X) (> 0)
%
lognormal_dist(Mu, Sigma) = X =>
    if Sigma =< 0.0 then
        throw($error('lognormal_dist_random: Sigma must be > 0'))
    else
        Y = normal_dist(Mu, Sigma),
        X = exp(Y)
    end.
lognormal_dist_n(Mu,Sigma,N) = [lognormal_dist(Mu,Sigma) : _ in 1..N].

%
% Log-normal distribution PDF
% lognormal_dist_pdf(Mu, Sigma, X)
%   Mu     : mean of log(X)
%   Sigma  : standard deviation of log(X) (> 0)
%   X      : value (>= 0)
%
lognormal_dist_pdf(Mu, Sigma, X) = Pdf =>
    if Sigma =< 0.0 then
        throw($error('lognormal_dist_pdf: Sigma must be > 0'))
    elseif X =< 0.0 then
        Pdf = 0.0
    else
        Z = (log(X) - Mu) / Sigma,
        LogF = -log(X) - log(Sigma) - 0.5*log(2.0*pi) - 0.5*(Z*Z),
        Pdf = exp(LogF)
    end.

%
% Log-normal distribution CDF
% lognormal_dist_cdf(Mu, Sigma, X)
%   Mu     : mean of log(X)
%   Sigma  : standard deviation of log(X) (> 0)
%   X      : value (>= 0)
%
lognormal_dist_cdf(Mu, Sigma, X) = Res =>
    if Sigma =< 0.0 then
        throw($error('lognormal_dist_cdf: Sigma must be > 0'))
    elseif X =< 0.0 then
        Res = 0.0
    elseif X >= 1.0e10 then
        Res = 1.0
    else
        Z = (log(X) - Mu) / Sigma,
        Res = normal_dist_cdf(0.0, 1.0, Z)
    end.

%
% Log-normal distribution quantile (inverse CDF)
% lognormal_dist_quantile(Mu, Sigma, P)
%   Mu     : mean of log(X)
%   Sigma  : standard deviation of log(X) (> 0)
%   P      : probability in [0,1]
%
lognormal_dist_quantile(Mu, Sigma, P) = Q =>
    if Sigma =< 0.0 then
        throw($error('lognormal_dist_quantile: Sigma must be > 0'))
    elseif P =< 0.0 then
        Q = 0.0
    elseif P >= 1.0 then
        Q = 1.0e10
    else
        Z = normal_dist_quantile(0.0, 1.0, P),
        Q = exp(Mu + Sigma * Z)
    end.

%
% Log-normal distribution mean
% lognormal_dist_mean(Mu, Sigma)
%   Mu     : mean of log(X)
%   Sigma  : standard deviation of log(X) (> 0)
%
lognormal_dist_mean(Mu, Sigma) = M =>
    if Sigma =< 0.0 then
        throw($error('lognormal_dist_mean: Sigma must be > 0'))
    else
        M = exp(Mu + 0.5 * Sigma * Sigma)
    end.


%
% Log-normal distribution variance
% lognormal_dist_variance(Mu, Sigma)
%   Mu     : mean of log(X)
%   Sigma  : standard deviation of log(X) (> 0)
%
lognormal_dist_variance(Mu, Sigma) = V =>
    if Sigma =< 0.0 then
        throw($error('lognormal_dist_variance: Sigma must be > 0'))
    else
        V = (exp(Sigma * Sigma) - 1.0) * exp(2.0 * Mu + Sigma * Sigma)
    end.


% Alias as log_normal_dist*
log_normal_dist(Mu, Sigma) = lognormal_dist(Mu, Sigma).
log_normal_dist_n(Mu, Sigma,N) = [lognormal_dist(Mu, Sigma) : _ in 1..N].

log_normal_dist_pdf(Mu, Sigma, X) = lognormal_dist_pdf(Mu, Sigma,X).
log_normal_dist_cdf(Mu, Sigma, X) = lognormal_dist_cdf(Mu, Sigma,X).
log_normal_dist_quantile(Mu, Sigma, X) = lognormal_dist_quantile(Mu, Sigma,X).
log_normal_dist_mean(Mu, Sigma) = lognormal_dist_mean(Mu, Sigma).
log_normal_dist_variance(Mu, Sigma) = lognormal_dist_variance(Mu, Sigma).

/*

  Student's t-distribution

  Description:
    The Student's t-distribution (t-distribution) is a continuous, symmetric,
    bell-shaped probability distribution that arises when estimating the mean
    of a normally distributed population with unknown variance using a small
    sample size. It is parameterized by one parameter:

        Nu : degrees of freedom (> 0)

    The t-distribution is heavier-tailed than the standard normal distribution,
    but approaches Normal(0,1) as Nu → ∞.

    Relation to other distributions:
        - t(1)  = Cauchy(0,1)
        - As Nu → ∞, t(Nu) → Normal(0,1)
        - Can be expressed as:
              T = Z / sqrt(V / Nu)
          where Z ~ Normal(0,1) and V ~ ChiSquare(Nu)

    PDF:
        f(x; Nu) =
          Γ((Nu+1)/2)
          ---------------------------- * (1 + x²/Nu)^(-(Nu+1)/2)
          √(Nuπ) Γ(Nu/2)

    CDF:
        F(x; Nu) = 0.5 * I_{Nu / (Nu + x²)}(Nu/2, 1/2)
                   if x < 0 then F(x) = 0.5 * I_x(...)
                   else           F(x) = 1 - 0.5 * I_x(...)

    Quantile (inverse CDF):
        For p ∈ (0,1),
            if p < 0.5 then
                t = -sqrt(Nu * (1/x - 1)),  x = I⁻¹(2p; Nu/2, 1/2)
            else
                t =  sqrt(Nu * (1/x - 1)),  x = I⁻¹(2(1-p); Nu/2, 1/2)

    Mean:
        0                  (if Nu > 1)
        undefined          (if Nu <= 1)

    Variance:
        Nu / (Nu - 2)      (if Nu > 2)
        +∞                 (if 1 < Nu <= 2)
        undefined          (if Nu <= 1)

  Example usage:
      Random sample
    println(student_t_dist_random(5))

      PDF and CDF values
    println(student_t_dist_pdf(5, 0.0))        ≈ 0.379606
    println(student_t_dist_cdf(5, 2.015))      ≈ 0.95

      Quantiles
    println(student_t_dist_quantile(5, 0.95))    ≈ 2.015
    println(student_t_dist_quantile(30, 0.975))   ≈ 2.042

  Use cases:
    - Hypothesis testing (e.g., t-tests for mean differences)
    - Bayesian modeling when variance is uncertain
    - Robust regression and inference under heavy-tailed noise
    - As a scale-mixture of normals for flexible modeling of outliers

  Implementation notes:
    - Random generation uses Z / sqrt(V / Nu)
    - PDF and CDF use lgamma and regularized incomplete beta
    - Quantile computed via inverse regularized incomplete beta
    - Numerically stable for wide range of Nu values (1e-10 precision typical)


  Examples: Hypothesis testing with the Student's t-distribution

  The t-distribution is commonly used when testing hypotheses about
  the mean of a normally distributed population when the population
  variance is unknown and the sample size is small.

  ----------------------------------------------------------------------------
  1. One-sample t-test (two-tailed)

    Given a sample of size N with sample mean M, sample standard deviation S,
    and hypothesized population mean Mu0.

    Null hypothesis  H0 : mean = Mu0
    Alternative      H1 : mean ≠ Mu0

    Test statistic:
        T = (M - Mu0) / (S / sqrt(N))

    Degrees of freedom:
        Nu = N - 1

    p-value (two-tailed):
        p = 2 * (1 - student_t_dist_cdf(Nu, abs(T)))

  Example:
    N   = 10,
    M   = 5.3,
    S   = 1.2,
    Mu0 = 5.0,
    Nu  = N - 1,
    T   = (M - Mu0) / (S / sqrt(N)),
    P   = 2 * (1.0 - student_t_dist_cdf(Nu, abs(T))),
    println([t_stat=T, p_value=P]).

  ----------------------------------------------------------------------------
  2. One-sample t-test (one-tailed, mean > Mu0)

    p-value (right-tailed):
        p = 1 - student_t_dist_cdf(Nu, T)

  Example:
    P = 1.0 - student_t_dist_cdf(Nu, T),
    if P < 0.05 then println('Reject H0') else println('Fail to reject H0') end.

  ----------------------------------------------------------------------------
  3. Two-sample (independent) t-test (equal variances)

    Suppose two independent samples X and Y:
        n1 = length(X), n2 = length(Y)
        m1 = mean(X),   m2 = mean(Y)
        s1 = stdev(X),  s2 = stdev(Y)

    Pooled variance:
        Sp2 = ((n1-1)*s1^2 + (n2-1)*s2^2) / (n1 + n2 - 2)

    Test statistic:
        T = (m1 - m2) / sqrt(Sp2 * (1/n1 + 1/n2))

    Degrees of freedom:
        Nu = n1 + n2 - 2

    Two-tailed p-value:
        p = 2 * (1.0 - student_t_dist_cdf(Nu, abs(T)))

  Example:
    X = [4.8,5.1,5.0,4.9,5.3],
    Y = [4.6,4.5,4.7,4.4,4.6],
    n1 = length(X), n2 = length(Y),
    m1 = sum(X)/n1, m2 = sum(Y)/n2,
    s1 = sqrt(sum([(Xi - m1)^2 : Xi in X])/(n1-1)),
    s2 = sqrt(sum([(Yi - m2)^2 : Yi in Y])/(n2-1)),
    Sp2 = ((n1-1)*s1^2 + (n2-1)*s2^2) / (n1 + n2 - 2),
    T = (m1 - m2) / sqrt(Sp2 * (1.0/n1 + 1.0/n2)),
    Nu = n1 + n2 - 2,
    P = 2.0 * (1.0 - student_t_dist_cdf(Nu, abs(T))),
    println([t_stat=T, p_value=P]).

  ----------------------------------------------------------------------------
  4. Confidence interval for mean (two-sided, 95%)

    CI = M ± t_{0.975,Nu} * S / sqrt(N)

  Example:
    Alpha = 0.05,
    Tcrit = student_t_dist_quantile(Nu, 1.0 - Alpha/2.0),
    Lower = M - Tcrit * S / sqrt(N),
    Upper = M + Tcrit * S / sqrt(N),
    println([conf_int=[Lower,Upper]]).

  ----------------------------------------------------------------------------
  Interpretation:
    - Small p-values (< 0.05) indicate evidence against H0.
    - The confidence interval gives a range of plausible true means.

*/

%
% Student's t-distribution random generator
% student_t_dist_random(Nu)
%   Nu : degrees of freedom (> 0)
%
student_t_dist(Nu) = T =>
    if Nu =< 0.0 then
        throw($error('student_t_dist_random: Nu must be > 0'))
    else
        Z = normal_dist(0.0, 1.0),
        V = chi_square_dist(Nu),
        T = Z / sqrt(V / Nu)
    end.
student_t_dist_n(Nu,N) = [student_t_dist(Nu) : _ in 1..N].

%
% Student's t-distribution PDF
% student_t_dist_pdf(Nu, X)
%   Nu : degrees of freedom (> 0)
%   X  : value
%
student_t_dist_pdf(Nu, X) = Pdf =>
    if Nu =< 0.0 then
        throw($error('student_t_dist_pdf: Nu must be > 0'))
    else
        % log form for stability:
        % log(f) = lgamma((Nu+1)/2)
        %           - lgamma(Nu/2)
        %           - 0.5*(log(Nu) + log(pi))
        %           - ((Nu+1)/2)*log(1 + X^2/Nu)
        LogF = lgamma((Nu + 1.0) / 2.0)
               - lgamma(Nu / 2.0)
               - 0.5 * (log(Nu) + log(pi))
               - ((Nu + 1.0) / 2.0) * log(1.0 + (X * X) / Nu),
        Pdf = exp(LogF)
    end.


%
% Student's t-distribution CDF
% student_t_dist_cdf(Nu, T)
%   Nu : degrees of freedom (> 0)
%   T  : value
%
student_t_dist_cdf(Nu, T) = Res =>
    if Nu =< 0.0 then
        throw($error('student_t_dist_cdf: Nu must be > 0'))
    elseif T =< -1.0e10 then
        Res = 0.0
    elseif T >=  1.0e10 then
        Res = 1.0
    elseif T =:= 0.0 then
        Res = 0.5
    else
        X = Nu / (Nu + T*T),
        A = Nu / 2.0,
        B = 0.5,
        Ix = reg_incomplete_beta(A, B, X),
        if T > 0.0 then
            Res = 1.0 - 0.5 * Ix
        else
            Res = 0.5 * Ix
        end
    end.




%
% Student's t-distribution quantile (inverse CDF)
% student_t_dist_quantile(Nu, P)
%   Nu : degrees of freedom (> 0)
%   P  : probability in [0,1]
%
student_t_dist_quantile(Nu, P) = T =>
    if Nu =< 0.0 then
        throw($error('student_t_dist_quantile: Nu must be > 0'))
    elseif P =< 0.0 then
        T = -1.0e10
    elseif P >= 1.0 then
        T = 1.0e10
    elseif P =:= 0.5 then
        T = 0.0
    else
        A = Nu / 2.0,
        B = 0.5,
        % tail symmetry
        if P < 0.5 then
            X = inverse_reg_incomplete_beta(A, B, 2.0 * P),
            T = -sqrt(Nu * (1.0 / X - 1.0))
        else
            X = inverse_reg_incomplete_beta(A, B, 2.0 * (1.0 - P)),
            T = sqrt(Nu * (1.0 / X - 1.0))
        end
    end.




%
% Student's t-distribution mean
% student_t_dist_mean(Nu)
%   Nu : degrees of freedom (> 0)
%   Mean exists only for Nu > 1
%
student_t_dist_mean(Nu) = M =>
    if Nu =< 0.0 then
        throw($error('student_t_dist_mean: Nu must be > 0'))
    elseif Nu =< 1.0 then
        throw($error('student_t_dist_mean: mean undefined for Nu <= 1'))
    else
        M = 0.0
    end.


%
% Student's t-distribution variance
% student_t_dist_variance(Nu)
%   Nu : degrees of freedom (> 0)
%   Variance:
%     - undefined for Nu <= 1
%     - infinite for 1 < Nu <= 2
%     - Nu / (Nu - 2) for Nu > 2
%
student_t_dist_variance(Nu) = V =>
    if Nu =< 0.0 then
        throw($error('student_t_dist_variance: Nu must be > 0'))
    elseif Nu =< 1.0 then
        throw($error('student_t_dist_variance: variance undefined for Nu <= 1'))
    elseif Nu =< 2.0 then
        V = 1.0/0.0    % +Infinity
    else
        V = Nu / (Nu - 2.0)
    end.



/* ---------------------------------------------------------------------------
Three-parameter Student's t-distribution
t(Mu, Sigma, Nu)

Description:
    The three-parameter Student's t-distribution is a generalization of the
    standard t-distribution that introduces:
        Mu     : location (center of distribution)
        Sigma  : scale   (> 0, controls spread)
        Nu     : degrees of freedom (> 0, controls tail heaviness)

    The standard t-distribution is recovered when Mu = 0 and Sigma = 1.

    If T ~ StudentT(0, 1, Nu), then:
        X = Mu + Sigma * T
    follows StudentT(Mu, Sigma, Nu).

    The distribution is symmetric around Mu, with heavier tails than the
    normal distribution. As Nu → ∞, the distribution approaches
    Normal(Mu, Sigma).

-------------------------------------------------------------------------------
Mathematical definitions

    PDF:
        f(x; Mu, Sigma, Nu) =
            [ Γ((Nu+1)/2) ]
            -----------------------------------
            [ Γ(Nu/2) * sqrt(Nuπ) * Sigma ]
            * ( 1 + ((x - Mu)²) / (Nu * Sigma²) )^(-(Nu+1)/2)

    CDF:
        F(x; Mu, Sigma, Nu) = 
            StudentT_CDF( (x - Mu)/Sigma ; Nu )

        where StudentT_CDF(z; Nu) is the CDF of the standard t-distribution.

    Quantile (inverse CDF):
        Q(p; Mu, Sigma, Nu) =
            Mu + Sigma * StudentT_Quantile(p; Nu)

    Mean:
        E[X] = Mu                     (for Nu > 1)
                undefined              (for Nu <= 1)

    Variance:
        Var[X] = Sigma² * Nu / (Nu - 2)   (for Nu > 2)
                  ∞                        (for 1 < Nu <= 2)
                  undefined                (for Nu <= 1)

-------------------------------------------------------------------------------
Example usage:

    % Random sampling
    println(student_t_dist_random(0.0, 1.0, 5.0))
    println(student_t_dist_random(5.0, 2.0, 10.0))

    % PDF
    println(student_t_dist_pdf(0.0, 1.0, 5.0, 0.0))      % ≈ 0.3796
    println(student_t_dist_pdf(5.0, 2.0, 10.0, 5.0))     % ≈ 0.1946

    % CDF
    println(student_t_dist_cdf(5.0, 2.0, 10.0, 7.0))     % ≈ 0.789
    println(student_t_dist_cdf(5.0, 2.0, 10.0, 3.0))     % ≈ 0.211

    % Quantile
    println(student_t_dist_quantile(5.0, 2.0, 10.0, 0.975))  % ≈ 9.456
    println(student_t_dist_quantile(5.0, 2.0, 10.0, 0.025))  % ≈ 0.544

    % Mean and variance
    println(student_t_dist_mean(5.0, 2.0, 10.0))         % 5.0
    println(student_t_dist_variance(5.0, 2.0, 10.0))     % 5.0

-------------------------------------------------------------------------------
Applications:

  • Bayesian and robust regression
        Used when residuals have heavy tails or outliers.
        Acts as a scale-mixture of normals with unknown precision.

  • Financial modeling
        Commonly used to model returns with high kurtosis.

  • Monte Carlo simulation
        Generates data with adjustable tail weight for stress testing.

  • Statistical inference
        General form used in t-tests with estimated mean and variance.

-------------------------------------------------------------------------------
Implementation notes:

  - Random generation uses: Mu + Sigma * student_t_dist_random(Nu)
  - PDF derived directly from standard Student-t scaled by 1/Sigma
  - CDF and quantile reuse the standard Student-t functions
  - Mean and variance follow standard formulas scaled by Sigma²
  - Handles edge cases with proper exceptions and infinite variance
  - Numerical precision ~1e-15 for moderate Nu values
--------------------------------------------------------------------------- */



%
% 3-parameter Student's t random generator
% student_t_dist_random(Mu, Sigma, Nu)
%   Mu     : location
%   Sigma  : scale (> 0)
%   Nu     : degrees of freedom (> 0)
%
student_t_dist(Mu, Sigma, Nu) = X =>
    if Sigma =< 0.0 then
        throw($error('student_t_dist_random: Sigma must be > 0'))
    elseif Nu =< 0.0 then
        throw($error('student_t_dist_random: Nu must be > 0'))
    else
        T = student_t_dist(Nu),
        X = Mu + Sigma * T
    end.
student_t_dist_n(Mu, Sigma, Nu, N) = [student_t_dist(Mu, Sigma, Nu) : _ in 1..N].



%
% 3-parameter Student's t-distribution PDF
% student_t_dist_pdf(Mu, Sigma, Nu, X)
%   Mu     : location
%   Sigma  : scale (> 0)
%   Nu     : degrees of freedom (> 0)
%   X      : value
%
student_t_dist_pdf(Mu, Sigma, Nu, X) = Pdf =>
    if Sigma =< 0.0 then
        throw($error('student_t_dist_pdf: Sigma must be > 0'))
    elseif Nu =< 0.0 then
        throw($error('student_t_dist_pdf: Nu must be > 0'))
    else
        Z = (X - Mu) / Sigma,
        LogF = lgamma((Nu + 1.0) / 2.0)
               - lgamma(Nu / 2.0)
               - 0.5 * (log(Nu) + log(pi))
               - log(Sigma)
               - ((Nu + 1.0) / 2.0) * log(1.0 + (Z * Z) / Nu),
        Pdf = exp(LogF)
    end.

%
% 3-parameter Student's t-distribution CDF
% student_t_dist_cdf(Mu, Sigma, Nu, X)
%   Mu     : location
%   Sigma  : scale (> 0)
%   Nu     : degrees of freedom (> 0)
%   X      : value
%
student_t_dist_cdf(Mu, Sigma, Nu, X) = Res =>
    if Sigma =< 0.0 then
        throw($error('student_t_dist_cdf: Sigma must be > 0'))
    elseif Nu =< 0.0 then
        throw($error('student_t_dist_cdf: Nu must be > 0'))
    else
        Z = (X - Mu) / Sigma,
        Res = student_t_dist_cdf(Nu, Z)
    end.


%
% 3-parameter Student's t-distribution quantile (inverse CDF)
% student_t_dist_quantile(Mu, Sigma, Nu, P)
%   Mu     : location
%   Sigma  : scale (> 0)
%   Nu     : degrees of freedom (> 0)
%   P      : probability in [0,1]
%
student_t_dist_quantile(Mu, Sigma, Nu, P) = Q =>
    if Sigma =< 0.0 then
        throw($error('student_t_dist_quantile: Sigma must be > 0'))
    elseif Nu =< 0.0 then
        throw($error('student_t_dist_quantile: Nu must be > 0'))
    elseif P =< 0.0 then
        Q = Mu - 1.0e10 * Sigma
    elseif P >= 1.0 then
        Q = Mu + 1.0e10 * Sigma
    else
        T = student_t_dist_quantile(Nu, P),
        Q = Mu + Sigma * T
    end.


%
% 3-parameter Student's t-distribution mean
% student_t_dist_mean(Mu, Sigma, Nu)
%   Mu     : location
%   Sigma  : scale (> 0)
%   Nu     : degrees of freedom (> 0)
%
student_t_dist_mean(Mu, Sigma, Nu) = M =>
    if Sigma =< 0.0 then
        throw($error('student_t_dist_mean: Sigma must be > 0'))
    elseif Nu =< 0.0 then
        throw($error('student_t_dist_mean: Nu must be > 0'))
    elseif Nu =< 1.0 then
        throw($error('student_t_dist_mean: mean undefined for Nu <= 1'))
    else
        M = Mu
    end.


%
% 3-parameter Student's t-distribution variance
% student_t_dist_variance(Mu, Sigma, Nu)
%   Mu     : location (not used in calculation)
%   Sigma  : scale (> 0)
%   Nu     : degrees of freedom (> 0)
%
student_t_dist_variance(Mu, Sigma, Nu) = V =>
    if Sigma =< 0.0 then
        throw($error('student_t_dist_variance: Sigma must be > 0'))
    elseif Nu =< 0.0 then
        throw($error('student_t_dist_variance: Nu must be > 0'))
    elseif Nu =< 1.0 then
        throw($error('student_t_dist_variance: variance undefined for Nu <= 1'))
    elseif Nu =< 2.0 then
        V = 10**100 % 1.0/0.0   % +Infinity
    else
        V = Sigma * Sigma * Nu / (Nu - 2.0)
    end.



/* ---------------------------------------------------------------------------
F-distribution

Description:
    The F-distribution (also known as the Fisher–Snedecor distribution)
    is a continuous probability distribution that arises frequently in
    hypothesis testing, especially in tests that compare variances or
    fit models such as ANOVA and regression.

    It is defined by two positive parameters:
        Nu1 : numerator degrees of freedom
        Nu2 : denominator degrees of freedom

    If
        X1 ~ ChiSquare(Nu1)
        X2 ~ ChiSquare(Nu2)
    are independent random variables, then
        F = (X1 / Nu1) / (X2 / Nu2)
    follows the F(Nu1, Nu2) distribution.

    The distribution is right-skewed but approaches a normal shape
    as both Nu1 and Nu2 increase.

-------------------------------------------------------------------------------
Mathematical definitions

    PDF:
        f(x; Nu1, Nu2) =
            [ (Nu1/Nu2)^(Nu1/2) * x^(Nu1/2 - 1) ]
            ----------------------------------------------
            [ B(Nu1/2, Nu2/2) * (1 + (Nu1*x/Nu2))^((Nu1+Nu2)/2) ]

        for x > 0, and f(x) = 0 for x <= 0.

    CDF:
        F(x; Nu1, Nu2) =
            I_{ Nu1*x / (Nu1*x + Nu2) }( Nu1/2, Nu2/2 )

        where I_z(a,b) is the regularized incomplete beta function.

    Quantile (inverse CDF):
        Given p ∈ (0,1), let
            x = I⁻¹_p( Nu1/2, Nu2/2 )
        then
            Q(p; Nu1, Nu2) = (Nu2 * x) / (Nu1 * (1 - x))

    Mean:
        E[F] = Nu2 / (Nu2 - 2)         (for Nu2 > 2)
        undefined                      (for Nu2 <= 2)

    Variance:
        Var[F] = 2 * Nu2² * (Nu1 + Nu2 - 2)
                 ----------------------------------------
                 Nu1 * (Nu2 - 2)² * (Nu2 - 4)

                 (for Nu2 > 4)
        infinite (for 2 < Nu2 <= 4)
        undefined (for Nu2 <= 2)

-------------------------------------------------------------------------------
Example usage:

    % Random sampling
    println(f_dist_random(5, 2))
    println(f_dist_random(10, 10))

    % PDF and CDF
    println(f_dist_pdf(5, 2, 1.0))        % ≈ 0.3193
    println(f_dist_cdf(5, 2, 1.0))        % ≈ 0.5205

    % Quantiles
    println(f_dist_quantile(5, 2, 0.95))  % ≈ 19.296
    println(f_dist_quantile(10, 10, 0.5)) % ≈ 1.0

    % Mean and variance
    println(f_dist_mean(5, 10))           % 1.25
    println(f_dist_variance(5, 10))       % ≈ 0.7031

-------------------------------------------------------------------------------
Applications:

  • Analysis of Variance (ANOVA)
        Tests whether several groups have equal means by comparing
        between-group variance to within-group variance.

  • Regression Analysis
        Tests overall model significance (F-test for R²).

  • Variance-ratio tests
        Tests whether two populations have equal variances.

  • Model comparison
        Used when comparing nested models via ratio of mean squares.

  • Monte Carlo simulations
        Can be used to generate random F-values for statistical modeling
        or bootstrapped hypothesis tests.

-------------------------------------------------------------------------------
Implementation notes:

  - Random generation uses two independent chi-square draws.
  - PDF uses beta_func(A, B) for normalization.
  - CDF and quantile use reg_incomplete_beta/3 and its inverse.
  - Handles parameter and edge cases according to statistical definitions.
  - Numerical precision typically 1e-15 for moderate degrees of freedom.

--------------------------------------------------------------------------- */

%
% F-distribution random generator
% f_dist_random(Nu1, Nu2)
%   Nu1 : numerator degrees of freedom (> 0)
%   Nu2 : denominator degrees of freedom (> 0)
%
f_dist(Nu1, Nu2) = F =>
    if Nu1 =< 0.0 then
        throw($error('f_dist_random: Nu1 must be > 0'))
    elseif Nu2 =< 0.0 then
        throw($error('f_dist_random: Nu2 must be > 0'))
    else
        X1 = chi_square_dist(Nu1),
        X2 = chi_square_dist(Nu2),
        F = (X1 / Nu1) / (X2 / Nu2)
    end.

f_dist_n(Nu1, Nu2, N) = [f_dist(Nu1, Nu2) : _ in 1..N].


%
% F-distribution PDF
% f_dist_pdf(Nu1, Nu2, X)
%   Nu1 : numerator degrees of freedom (> 0)
%   Nu2 : denominator degrees of freedom (> 0)
%   X   : value (>= 0)
%
f_dist_pdf(Nu1, Nu2, X) = Pdf =>
    if Nu1 =< 0.0 then
        throw($error('f_dist_pdf: Nu1 must be > 0'))
    elseif Nu2 =< 0.0 then
        throw($error('f_dist_pdf: Nu2 must be > 0'))
    elseif X =< 0.0 then
        Pdf = 0.0
    else
        A = Nu1 / 2.0,
        B = Nu2 / 2.0,
        % log form for numerical stability
        LogF = A * log(Nu1 / Nu2)
               + (A - 1.0) * log(X)
               - (A + B) * log(1.0 + (Nu1 * X / Nu2))
               - log(beta_func(A, B)),
        Pdf = exp(LogF)
    end.


%
% F-distribution CDF
% f_dist_cdf(Nu1, Nu2, X)
%   Nu1 : numerator degrees of freedom (> 0)
%   Nu2 : denominator degrees of freedom (> 0)
%   X   : value (>= 0)
%
f_dist_cdf(Nu1, Nu2, X) = Res =>
    if Nu1 =< 0.0 then
        throw($error('f_dist_cdf: Nu1 must be > 0'))
    elseif Nu2 =< 0.0 then
        throw($error('f_dist_cdf: Nu2 must be > 0'))
    elseif X =< 0.0 then
        Res = 0.0
    elseif X >= 1.0e10 then
        Res = 1.0
    else
        A = Nu1 / 2.0,
        B = Nu2 / 2.0,
        Z = (Nu1 * X) / (Nu1 * X + Nu2),
        Res = reg_incomplete_beta(A, B, Z)
    end.


%
% F-distribution quantile (inverse CDF)
% f_dist_quantile(Nu1, Nu2, P)
%   Nu1 : numerator degrees of freedom (> 0)
%   Nu2 : denominator degrees of freedom (> 0)
%   P   : probability in [0,1]
%
f_dist_quantile(Nu1, Nu2, P) = Q =>
    if Nu1 =< 0.0 then
        throw($error('f_dist_quantile: Nu1 must be > 0'))
    elseif Nu2 =< 0.0 then
        throw($error('f_dist_quantile: Nu2 must be > 0'))
    elseif P =< 0.0 then
        Q = 0.0
    elseif P >= 1.0 then
        Q = 1.0e10
    else
        A = Nu1 / 2.0,
        B = Nu2 / 2.0,
        X = inverse_reg_incomplete_beta(A, B, P),
        Q = (Nu2 * X) / (Nu1 * (1.0 - X))
    end.


%
% F-distribution mean
% f_dist_mean(Nu1, Nu2)
%   Nu1 : numerator degrees of freedom (> 0)
%   Nu2 : denominator degrees of freedom (> 0)
%   Mean defined only for Nu2 > 2
%
f_dist_mean(Nu1, Nu2) = M =>
    if Nu1 =< 0.0 then
        throw($error('f_dist_mean: Nu1 must be > 0'))
    elseif Nu2 =< 0.0 then
        throw($error('f_dist_mean: Nu2 must be > 0'))
    elseif Nu2 =< 2.0 then
        throw($error('f_dist_mean: mean undefined for Nu2 <= 2'))
    else
        M = Nu2 / (Nu2 - 2.0)
    end.


%
% F-distribution variance
% f_dist_variance(Nu1, Nu2)
%   Nu1 : numerator degrees of freedom (> 0)
%   Nu2 : denominator degrees of freedom (> 0)
%   Variance:
%     - undefined for Nu2 <= 2
%     - infinite for 2 < Nu2 <= 4
%     - 2*Nu2^2*(Nu1+Nu2-2) / (Nu1*(Nu2-2)^2*(Nu2-4)) for Nu2 > 4
%
f_dist_variance(Nu1, Nu2) = V =>
    if Nu1 =< 0.0 then
        throw($error('f_dist_variance: Nu1 must be > 0'))
    elseif Nu2 =< 0.0 then
        throw($error('f_dist_variance: Nu2 must be > 0'))
    elseif Nu2 =< 2.0 then
        throw($error('f_dist_variance: variance undefined for Nu2 <= 2'))
    elseif Nu2 =< 4.0 then
        V = -1.0e300 % 1.0/0.0    % +Infinity
        
    else
        V = 2.0 * Nu2 * Nu2 * (Nu1 + Nu2 - 2.0) /
            (Nu1 * (Nu2 - 2.0)**2 * (Nu2 - 4.0))
    end.


/* ---------------------------------------------------------------------------
Logistic distribution
Logistic(Mu, S)

Description:
    The Logistic distribution is a continuous probability distribution that
    resembles the Normal distribution but has heavier tails. It is symmetric
    about the location parameter Mu and characterized by a scale parameter S.

    The Logistic distribution is commonly used as a simple alternative to
    the Normal distribution in situations where heavier tails are observed,
    or when analytical simplicity is desired (e.g. logistic regression).

    Parameters:
        Mu : location parameter (mean)
        S  : scale parameter (> 0)

    The standard Logistic distribution is obtained when Mu = 0 and S = 1.

-------------------------------------------------------------------------------
Mathematical definitions

    PDF:
        f(x; Mu, S) =
            exp(-(x - Mu)/S) /
            [ S * (1 + exp(-(x - Mu)/S))² ]

    CDF:
        F(x; Mu, S) =
            1 / (1 + exp(-(x - Mu)/S))

    Quantile (inverse CDF):
        Q(p; Mu, S) =
            Mu + S * ln( p / (1 - p) )

    Mean:
        E[X] = Mu

    Variance:
        Var[X] = (π² * S²) / 3

-------------------------------------------------------------------------------
Example usage:

    % Random sampling
    println(logistic_dist_random(0.0, 1.0))
    println(logistic_dist_random(5.0, 2.0))

    % PDF
    println(logistic_dist_pdf(0.0, 1.0, 0.0))   % 0.25
    println(logistic_dist_pdf(0.0, 1.0, 2.0))   % ≈ 0.105

    % CDF
    println(logistic_dist_cdf(0.0, 1.0, 0.0))   % 0.5
    println(logistic_dist_cdf(0.0, 1.0, 2.0))   % ≈ 0.8808
    println(logistic_dist_cdf(0.0, 1.0, -2.0))  % ≈ 0.1192

    % Quantile
    println(logistic_dist_quantile(0.0, 1.0, 0.5))     % 0.0
    println(logistic_dist_quantile(0.0, 1.0, 0.73106)) % ≈ 1.0
    println(logistic_dist_quantile(0.0, 1.0, 0.26894)) % ≈ -1.0

    % Mean and variance
    println(logistic_dist_mean(0.0, 1.0))              % 0.0
    println(logistic_dist_variance(0.0, 1.0))          % 3.2899

-------------------------------------------------------------------------------
Applications:

  • Logistic regression:
        The error term in logistic regression follows a Logistic(0,1)
        distribution rather than a Normal(0,1).

  • Robust modeling:
        Used as a heavy-tailed alternative to the Normal distribution
        in statistics and Bayesian modeling.

  • Growth curves:
        The derivative of the logistic growth function follows the
        Logistic distribution.

  • Reliability and lifetime analysis:
        Used for symmetric lifetime data where the failure rate
        increases then decreases.

  • Monte Carlo simulation:
        Convenient due to the closed-form quantile function
        (requires no iterative solver).

-------------------------------------------------------------------------------
Implementation notes:

  - Random generation: inverse-transform using U(0,1)
        X = Mu + S * ln(U / (1 - U))
  - PDF and CDF computed via stable exponential forms
  - Quantile function is exact analytical inverse
  - Mean = Mu; Variance = π² * S² / 3
  - Numerical precision typically better than 1e-15

--------------------------------------------------------------------------- */


%
% Logistic distribution random generator
% logistic_dist_random(Mu, S)
%   Mu : location
%   S  : scale (> 0)
%
logistic_dist(Mu, S) = X =>
    if S =< 0.0 then
        throw($error('logistic_dist_random: S must be > 0'))
    else
        U = u01(),  % safe uniform(0,1)
        X = Mu + S * log(U / (1.0 - U))
    end.
logistic_dist_n(Mu, S, N) = [logistic_dist(Mu, S) : _ in 1..N].


%
% Logistic distribution PDF
% logistic_dist_pdf(Mu, S, X)
%   Mu : location
%   S  : scale (> 0)
%   X  : value
%
logistic_dist_pdf(Mu, S, X) = Pdf =>
    if S =< 0.0 then
        throw($error('logistic_dist_pdf: S must be > 0'))
    else
        Z = (X - Mu) / S,
        Ez = exp(-Z),
        Pdf = Ez / (S * (1.0 + Ez) * (1.0 + Ez))
    end.

%
% Logistic distribution CDF
% logistic_dist_cdf(Mu, S, X)
%   Mu : location
%   S  : scale (> 0)
%   X  : value
%
logistic_dist_cdf(Mu, S, X) = Res =>
    if S =< 0.0 then
        throw($error('logistic_dist_cdf: S must be > 0'))
    else
        Z = (X - Mu) / S,
        if Z >= 40.0 then        % avoid overflow in exp(-Z)
            Res = 1.0
        elseif Z =< -40.0 then   % avoid underflow
            Res = 0.0
        else
            Res = 1.0 / (1.0 + exp(-Z))
        end
    end.


%
% Logistic distribution quantile (inverse CDF)
% logistic_dist_quantile(Mu, S, P)
%   Mu : location
%   S  : scale (> 0)
%   P  : probability in [0,1]
%
logistic_dist_quantile(Mu, S, P) = Q =>
    if S =< 0.0 then
        throw($error('logistic_dist_quantile: S must be > 0'))
    elseif P =< 0.0 then
        Q = Mu - 1.0e10 * S
    elseif P >= 1.0 then
        Q = Mu + 1.0e10 * S
    else
        Q = Mu + S * log(P / (1.0 - P))
    end.


%
% Logistic distribution mean
% logistic_dist_mean(Mu, S)
%   Mu : location
%   S  : scale (> 0)
%
logistic_dist_mean(Mu, S) = M =>
    if S =< 0.0 then
        throw($error('logistic_dist_mean: S must be > 0'))
    else
        M = Mu
    end.


%
% Logistic distribution variance
% logistic_dist_variance(Mu, S)
%   Mu : location
%   S  : scale (> 0)
%
logistic_dist_variance(Mu, S) = V =>
    if S =< 0.0 then
        throw($error('logistic_dist_variance: S must be > 0'))
    else
        V = (pi * pi * S * S) / 3.0
    end.


/*
  Logistic / 0 distribution

  Same as logistic_dist(0,1)

*/

logistic_dist() = logistic_dist(0,1).
logistic_dist_n(N) = [logistic_dist() : _ in 1..N].

logistic_dist_pdf(X) = logistic_dist_pdf(0,1,X).
logistic_dist_cdf(X) = logistic_dist_cdf(0,1,X).
logistic_dist_quantilef(X) = logistic_dist_quantile(0,1,X).

logistic_dist_mean() = logistic_dist_mean(0,1).
logistic_dist_variance() = logistic_dist_mean(0,1).



/* ---------------------------------------------------------------------------
Discrete Laplace distribution
DiscreteLaplace(Mu, P)

Description:
    The Discrete Laplace distribution (also known as the two-sided geometric)
    is the discrete analogue of the continuous Laplace (double exponential)
    distribution. It models symmetric integer-valued data centered at Mu
    with exponentially decaying probabilities in both directions.

    Parameters:
        Mu : location (integer or real)
        P  : dispersion parameter (0 < P < 1)
             Smaller P -> sharper peak, lighter tails
             Larger P -> heavier tails

    The PMF is symmetric around Mu. When P is small, the distribution is
    tightly concentrated near Mu; when P approaches 1, the tails grow heavy.

-------------------------------------------------------------------------------
Mathematical definitions

    PMF:
        P(X = k) = (1 - P) / (1 + P) * P^{|k - Mu|}
        for integer k

    CDF:
        F(x; Mu, P) =
            p^{Mu - floor(x)} / (1 + P),           for x < Mu
            1 - p^{floor(x) - Mu + 1} / (1 + P),   for x >= Mu

    Quantile (inverse CDF):
        Q(q; Mu, P) =
            Mu - ceil( ln(q*(1+P)) / ln(P) ),               for q < 1/(1+P)
            Mu + ceil( ln((1-q)*(1+P)) / ln(P) ) - 1,       otherwise

    Mean:
        E[X] = Mu

    Variance:
        Var[X] = 2P / (1 - P)²

-------------------------------------------------------------------------------
Example usage:

    % Random sampling
    println(discrete_laplace_dist_random(0, 0.5))
    println(discrete_laplace_dist_random(5, 0.9))

    % PMF
    println(discrete_laplace_dist_pdf(0, 0.5, 0))    % ≈ 0.333
    println(discrete_laplace_dist_pdf(0, 0.5, 1))    % ≈ 0.167

    % CDF
    println(discrete_laplace_dist_cdf(0, 0.5, 0))    % ≈ 0.667
    println(discrete_laplace_dist_cdf(0, 0.5, 1))    % ≈ 0.833

    % Quantile
    println(discrete_laplace_dist_quantile(0, 0.5, 0.5))  % 0
    println(discrete_laplace_dist_quantile(0, 0.5, 0.8))  % 1
    println(discrete_laplace_dist_quantile(0, 0.5, 0.2))  % -1

    % Mean and variance
    println(discrete_laplace_dist_mean(0, 0.5))           % 0
    println(discrete_laplace_dist_variance(0, 0.5))       % 4.0

-------------------------------------------------------------------------------
Applications:

  • Noise modeling and privacy:
        Used in differential privacy mechanisms (adding discrete noise).

  • Integer-valued analog of Laplace:
        Models absolute deviations of integer data from a center point.

  • Communication and error modeling:
        Describes integer-valued noise (e.g. quantized or count noise).

  • Statistical modeling:
        Useful for symmetric discrete heavy-tailed data.

  • Bayesian statistics:
        Common as a prior for sparse integer parameters.

-------------------------------------------------------------------------------
Implementation notes:

  - Random generation via inverse-transform using uniform(0,1)
  - PMF, CDF, and quantile all have exact closed forms
  - Mean = Mu; Variance = 2P / (1-P)²
  - Supports both integer and real Mu for generality
  - Numerical precision better than 1e-15 for typical P values

--------------------------------------------------------------------------- */


%
% Discrete Laplace distribution random generator
% discrete_laplace_dist_random(Mu, P)
%   Mu : location (integer or float)
%   P  : dispersion parameter (0 < P < 1)
%
discrete_laplace_dist(Mu, P) = X =>
    if P =< 0.0 ; P >= 1.0 then
        throw($error('discrete_laplace_dist_random: P must be in (0,1)'))
    else
        U = u01(),
        if U < 0.5 then
            K = floor(log(2.0 * U) / log(P)),
            X = Mu + abs(K)
        else
            K = floor(log(2.0 * (1.0 - U)) / log(P)),
            X = Mu - abs(K)
        end
    end.

discrete_laplace_dist_n(Mu, P, N) = [discrete_laplace_dist(Mu, P) : _ in 1..N].


%
% Discrete Laplace distribution PMF
% discrete_laplace_dist_pdf(Mu, P, K)
%   Mu : location
%   P  : dispersion parameter (0 < P < 1)
%   K  : integer value
%
discrete_laplace_dist_pdf(Mu, P, K) = Prob =>
    if P =< 0.0 ; P >= 1.0 then
        throw($error('discrete_laplace_dist_pmf: P must be in (0,1)'))
    else
        Prob = ((1.0 - P) / (1.0 + P)) * (P ** abs(K - Mu))
    end.


%
% Discrete Laplace distribution CDF
% discrete_laplace_dist_cdf(Mu, P, X)
%   Mu : location
%   P  : dispersion parameter (0 < P < 1)
%   X  : integer or real value
%
discrete_laplace_dist_cdf(Mu, P, X) = Res =>
    if P =< 0.0 ; P >= 1.0 then
        throw($error('discrete_laplace_dist_cdf: P must be in (0,1)'))
    else
        Xfloor = floor(X),
        if Xfloor < Mu then
            Res = P ** (Mu - Xfloor) / (1.0 + P)
        else
            Res = 1.0 - (P ** (Xfloor - Mu + 1)) / (1.0 + P)
        end
    end.

%
% Discrete Laplace distribution quantile (inverse CDF)
% discrete_laplace_dist_quantile(Mu, P, Q)
%   Mu : location
%   P  : dispersion parameter (0 < P < 1)
%   Q  : probability in [0,1]
%
discrete_laplace_dist_quantile(Mu, P, Q) = K =>
    if P =< 0.0 ; P >= 1.0 then
        throw($error('discrete_laplace_dist_quantile: P must be in (0,1)'))
    elseif Q =< 0.0 then
        K = Mu - 1.0e10
    elseif Q >= 1.0 then
        K = Mu + 1.0e10
    else
        if Q < 1.0 / (1.0 + P) then
            K = Mu - ceiling(log(Q * (1.0 + P)) / log(P))
        else
            K = Mu + ceiling(log((1.0 - Q) * (1.0 + P)) / log(P)) - 1
        end
    end.




%
% Discrete Laplace distribution mean
% discrete_laplace_dist_mean(Mu, P)
%   Mu : location (integer or real)
%   P  : dispersion parameter (0 < P < 1)
%
discrete_laplace_dist_mean(Mu, P) = M =>
    if P =< 0.0 ; P >= 1.0 then
        throw($error('discrete_laplace_dist_mean: P must be in (0,1)'))
    else
        % Symmetry about Mu ⇒ mean = Mu
        M = Mu
    end.


%
% Discrete Laplace distribution variance
% discrete_laplace_dist_variance(Mu, P)
%   Mu : location (not used)
%   P  : dispersion parameter (0 < P < 1)
%
discrete_laplace_dist_variance(Mu, P) = V =>
    if P =< 0.0 ; P >= 1.0 then
        throw($error('discrete_laplace_dist_variance: P must be in (0,1)'))
    else
        % Var = 2P / (1-P)^2
        Den = 1.0 - P,
        V = 2.0 * P / (Den * Den)
    end.



/* ---------------------------------------------------------------------------
Logarithmic (Log-series) distribution
LogSeries(Theta)

Description:
    The Logarithmic (or Log-series) distribution is a discrete probability
    distribution with support on the positive integers {1,2,3,...}.
    It was introduced by R.A. Fisher (1943) to model species abundances
    in ecology and remains a standard heavy-tailed count distribution.

    Parameters:
        Theta : parameter (0 < Theta < 1)
                 Controls the tail weight:
                   small Theta -> short tail
                   large Theta -> heavy tail

    The distribution is heavily right-skewed, with many small counts and
    a few large ones, making it ideal for overdispersed count data.

-------------------------------------------------------------------------------
Mathematical definitions

    PMF:
        P(X = k) =
            - (1 / ln(1 - Theta)) * (Theta^k / k),
        for integer k >= 1

    CDF:
        F(x; Theta) = Σ_{k=1}^{floor(x)} P(X = k)

    Quantile (inverse CDF):
        Q(p; Theta) = min { k : F(k; Theta) >= p }

    Mean:
        E[X] = -Theta / ((1 - Theta) * ln(1 - Theta))

    Variance:
        Var[X] = -Theta * (Theta + ln(1 - Theta)) /
                  ((1 - Theta)^2 * (ln(1 - Theta))^2)

-------------------------------------------------------------------------------
Example usage:

    % Random sampling
    println(logseries_dist_random(0.3))
    println(logseries_dist_random(0.7))

    % PDF (PMF)
    println(logseries_dist_pdf(0.3, 1))   % ≈ 0.386
    println(logseries_dist_pdf(0.3, 2))   % ≈ 0.0579
    println(logseries_dist_pdf(0.3, 3))   % ≈ 0.0174

    % CDF
    println(logseries_dist_cdf(0.3, 1))   % ≈ 0.386
    println(logseries_dist_cdf(0.3, 2))   % ≈ 0.444
    println(logseries_dist_cdf(0.3, 3))   % ≈ 0.461

    % Quantile
    println(logseries_dist_quantile(0.3, 0.5))  % → 1 or 2
    println(logseries_dist_quantile(0.7, 0.9))  % → 2

    % Mean and variance
    println(logseries_dist_mean(0.3))           % ≈ 1.20157
    println(logseries_dist_variance(0.3))       % ≈ 0.31061
    println(logseries_dist_mean(0.7))           % ≈ 2.77001
    println(logseries_dist_variance(0.7))       % ≈ 1.89846

-------------------------------------------------------------------------------
Applications:

  • Ecology:
        Fisher's log-series model for species abundance distributions
        (many rare species, few common ones).

  • Information theory and linguistics:
        Modeling symbol or word frequencies in heavy-tailed vocabularies.

  • Bayesian nonparametrics:
        Arises in compound Poisson and Dirichlet–process models.

  • Network and clustering models:
        Models count distributions with many small clusters and few large ones.

-------------------------------------------------------------------------------
Implementation notes:

  - Random generation via exact inverse transform.
  - PMF and CDF computed using recurrence:
        P(k+1) = P(k) * (Theta * k / (k+1))
  - Mean and variance formulas match Mathematica and SciPy.
  - Stable for Theta up to ~0.99; for Theta→1, use arbitrary precision if needed.
  - Numerical accuracy typically within 1e-15 (double precision).

--------------------------------------------------------------------------- */

%
% Logarithmic (log-series) distribution random generator (exact)
% logseries_dist_random(Theta)
%   Theta : parameter (0 < Theta < 1)
%
logseries_dist(Theta) = K =>
    if Theta =< 0.0 ; Theta >= 1.0 then
        throw($error('logseries_dist_random: Theta must be in (0,1)'))
    else
        U = u01(),
        Norm = -1.0 / log(1.0 - Theta),
        Sum = 0.0,
        K = 1,
        Prob = Norm * (Theta ** K) / K,
        Sum := Prob,
        while (U > Sum) do
            K := K + 1,
            Prob := Norm * (Theta ** K) / K,
            Sum := Sum + Prob
        end
    end.

logseries_dist_n(Theta, N) = [logseries_dist(Theta) : _ in 1..N].

%
% Logarithmic (log-series) distribution PMF
% logseries_dist_pdf(Theta, K)
%   Theta : parameter (0 < Theta < 1)
%   K     : integer value (>= 1)
%
logseries_dist_pdf(Theta, K) = P =>
    if Theta =< 0.0 ; Theta >= 1.0 then
        throw($error('logseries_dist_pdf: Theta must be in (0,1)'))
    elseif K < 1 then
        P = 0.0
    else
        P = - (1.0 / log(1.0 - Theta)) * (Theta ** K) / K
    end.


%
% Logarithmic (log-series) distribution CDF
% logseries_dist_cdf(Theta, X)
%   Theta : parameter (0 < Theta < 1)
%   X     : real/integer value (CDF up to floor(X))
%
logseries_dist_cdf(Theta, X) = Res =>
    if Theta =< 0.0 ; Theta >= 1.0 then
        throw($error('logseries_dist_cdf: Theta must be in (0,1)'))
    elseif X < 1.0 then
        Res = 0.0
    elseif X >= 1.0e10 then
        Res = 1.0
    else
        N = floor(X),
        C = -1.0 / log(1.0 - Theta),   % normalization constant
        % pmf(1) = C * Theta^1 / 1
        P = C * Theta,
        Sum = P,
        K = 1,
        while (K < N) do
            % pmf(k+1) = pmf(k) * (Theta * k / (k+1))
            P := P * (Theta * K / (K + 1.0)),
            Sum := Sum + P,
            K := K + 1
        end,
        Res = Sum
    end.


%
% Logarithmic (log-series) distribution quantile
% logseries_dist_quantile(Theta, P)
%   Theta : parameter (0 < Theta < 1)
%   P     : probability in [0,1]
% Returns the smallest integer k >= 1 such that CDF(k) >= P.
%
logseries_dist_quantile(Theta, P) = K =>
    if Theta =< 0.0 ; Theta >= 1.0 then
        throw($error('logseries_dist_quantile: Theta must be in (0,1)'))
    elseif P =< 0.0 then
        K = 1
    elseif P >= 1.0 then
        K = 10000000000     % very large sentinel (support is unbounded)
    else
        C = -1.0 / log(1.0 - Theta),   % normalization
        % pmf(1) = C * Theta^1 / 1
        Pk = C * Theta,
        Sum = Pk,
        K = 1,
        while (Sum < P) do
            K := K + 1,
            % pmf(k) via recurrence: pmf(k) = pmf(k-1) * (Theta * (k-1) / k)
            Pk := Pk * (Theta * (K - 1.0) / K),
            Sum := Sum + Pk
        end
    end.


%
% Log-series distribution mean
% logseries_dist_mean(Theta)
%   Theta : parameter (0 < Theta < 1)
%
logseries_dist_mean(Theta) = M =>
    if Theta =< 0.0 ; Theta >= 1.0 then
        throw($error('logseries_dist_mean: Theta must be in (0,1)'))
    else
        L = log(1.0 - Theta),
        % E[X] = -Theta / ((1-Theta) * ln(1-Theta))
        M = -Theta / ((1.0 - Theta) * L)
    end.


%
% Log-series distribution variance
% logseries_dist_variance(Theta) = V
%   Theta : parameter (0 < Theta < 1)
%
logseries_dist_variance(Theta) = V =>
    if Theta =< 0.0 ; Theta >= 1.0 then
        throw($error('logseries_dist_variance: Theta must be in (0,1)'))
    else
        L = log(1.0 - Theta),
        Den = (1.0 - Theta)**2 * (L * L),
        % Var[X] = -Theta * (Theta + ln(1-Theta)) / ((1-Theta)^2 * ln(1-Theta)^2)
        V = -(Theta * (Theta + L)) / Den
    end.


/*
----------------------------------------------------------------------
beta_negative_binomial distribution  (also called Beta–Pascal or Beta–Geometric generalization)

Definition:
------------
This discrete distribution models the number of failures (K)
before N successes, when the probability of success (P)
follows a Beta(Alpha, Beta) prior.

Mathematically:
    P(X = K) =
        C(N + K - 1, K) * B(Alpha + N, Beta + K) / B(Alpha, Beta)
    for K = 0, 1, 2, ...

where B(a,b) = Γ(a)Γ(b)/Γ(a+b), and C(n,k) = n! / (k!(n−k)!).

Parameters:
------------
    Alpha > 0.0
    Beta  > 0.0
    N     > 0    (number of successes)

Support:
------------
    K ∈ {0, 1, 2, ...}

Functions implemented:
----------------------
    beta_negative_binomial_dist_pdf(Alpha, Beta, N, K)
        - Probability mass function.

    beta_negative_binomial_dist_cdf(Alpha, Beta, N, K)
        - Cumulative probability P(X ≤ K).

    beta_negative_binomial_dist_quantile(Alpha, Beta, N, P)
        - Smallest integer K such that CDF(K) ≥ P.

    beta_negative_binomial_dist_mean(Alpha, Beta, N)
        - Mean = N * Beta / (Alpha - 1), requires Alpha > 1.

    beta_negative_binomial_dist_variance(Alpha, Beta, N)
        - Variance = N*(A2 - A1) + N^2*(A2 - A1^2),
          where A1 = (Alpha+Beta-1)/(Alpha-1),
                A2 = (Alpha+Beta-1)*(Alpha+Beta-2)/((Alpha-1)*(Alpha-2)),
          requires Alpha > 2.

    beta_negative_binomial_dist_rand(Alpha, Beta, N)
        - Random variate via hierarchical sampling:
              P ~ Beta(Alpha, Beta)
              X ~ NegBinomial(N, P)
              return X

Examples:
----------
    % Mean and variance
    println(beta_negative_binomial_dist_mean(2.0, 3.0, 5))
    println(beta_negative_binomial_dist_variance(3.0, 3.0, 4))

    % Probability mass for k = 0..5
    foreach(K in 0..5)
        println([K, beta_negative_binomial_dist_pdf(2.0, 3.0, 5, K)])
    end.

    % Random generation
    X = beta_negative_binomial_dist_rand(2.0, 3.0, 5),
    println(X).

    % Quantile example
    K = beta_negative_binomial_dist_quantile(2.0, 3.0, 5, 0.9),
    println([quantile_0.9=K]).

Notes:
-------
  * Requires lgamma/1 and exp/1 (built into Picat math library).
  * When Alpha ≤ 1 or Alpha ≤ 2, mean or variance respectively are undefined.
  * As Alpha → ∞, Beta → finite, the distribution approaches the standard
    Negative Binomial with success probability p = Alpha / (Alpha + Beta).
----------------------------------------------------------------------
*/


% Beta Negative Binomial Distribution
% Parameters: Alpha>0, Beta>0, N>0
% Returns number of failures before N successes.
beta_negative_binomial_dist(Alpha, Beta, N) = X =>
    if Alpha =< 0.0 ; Beta =< 0.0 ; N =< 0 then
        throw(invalid_parameters_for_betanegbinomial)
    end,
    P = beta_dist(Alpha, Beta),
    X = negative_binomial_dist(N, P).
beta_negative_binomial_dist_n(Alpha, Beta, N, Num) = [beta_negative_binomial_dist(Alpha, Beta, N) : _ in 1..Num].


% Beta Negative Binomial (Beta–Pascal) distribution
% PDF: Probability of exactly K failures before N successes,
% when success prob ~ Beta(Alpha,Beta)
% Parameters: Alpha>0, Beta>0, N>0, K>=0 (integer)

beta_negative_binomial_dist_pdf(Alpha, Beta, N, K) = P =>
    if Alpha =< 0.0 ; Beta =< 0.0 ; N =< 0 ; K < 0 ; not integer(K) then
        P = 0.0
    else
        % log(Beta(a,b)) = lgamma(a)+lgamma(b)-lgamma(a+b)
        LBeta1 = lgamma(Alpha + N) + lgamma(Beta + K) - lgamma(Alpha + Beta + N + K),
        LBeta0 = lgamma(Alpha) + lgamma(Beta) - lgamma(Alpha + Beta),
        LComb  = lgamma(N + K) - lgamma(N) - lgamma(K + 1),
        LogP   = LComb + (LBeta1 - LBeta0),
        P = exp(LogP)
    end.


% Quantile for Beta Negative Binomial (Beta–Pascal)
% Returns the smallest K >= 0 such that CDF(K) >= P.
% Parameters: Alpha>0, Beta>0, N>0, P in [0,1]

beta_negative_binomial_dist_quantile(Alpha, Beta, N, P) = K =>
    % --- Parameter checks and edge cases -----------------------------------
    if Alpha =< 0.0 ; Beta =< 0.0 ; N =< 0 then
        throw(invalid_parameters_for_beta_negative_binomial_quantile)
    end,
    if P =< 0.0 then
        K = 0
    elseif P >= 1.0 then
        K = 1.0e300 % inf  % or a large sentinel value per your policy
    else
        K = _,
        % --- Initialize at k = 0 ------------------------------------------
        LBeta0 = lgamma(Alpha) + lgamma(Beta) - lgamma(Alpha + Beta),
        LBeta1_k0 = lgamma(Alpha + N) + lgamma(Beta) - lgamma(Alpha + Beta + N),
        LogP0 = LBeta1_k0 - LBeta0,
        Term = exp(LogP0),   % PMF at k=0
        S = Term,            % cumulative sum so far
        J = 0,
        OK = true,

        % --- Iterate using the PMF recurrence ------------------------------
        while (OK == true)
            if S >= P then
                OK := false
            else
                Num = (N + J) * (Beta + J),
                Den = (J + 1) * (Alpha + Beta + N + J),
                Term := Term * (Num / Den),
                S := S + Term,
                J := J + 1,
                if S >= P then
                    OK := false
                end
            end
        end,
        K = J
    end.


% Mean: E[X] = N * Beta / (Alpha - 1), requires Alpha > 1
beta_negative_binomial_dist_mean(Alpha, Beta, N) = M =>
    if Alpha =< 1.0 ; Beta =< 0.0 ; N =< 0 then
        throw(invalid_parameters_for_beta_negative_binomial_mean)
    end,
    M = N * (Beta / (Alpha - 1.0)).

% Variance:
% Var[X] = N*(A2 - A1) + N^2*(A2 - A1^2)
% where A1 = (Alpha+Beta-1)/(Alpha-1), A2 = (Alpha+Beta-1)*(Alpha+Beta-2)/((Alpha-1)*(Alpha-2))
% requires Alpha > 2
beta_negative_binomial_dist_variance(Alpha, Beta, N) = V =>
    if Alpha =< 2.0 ; Beta =< 0.0 ; N =< 0 then
        throw(invalid_parameters_for_beta_negative_binomial_variance)
    end,
    A1 = (Alpha + Beta - 1.0) / (Alpha - 1.0),
    A2 = ((Alpha + Beta - 1.0) * (Alpha + Beta - 2.0)) / ((Alpha - 1.0) * (Alpha - 2.0)),
    V = N * (A2 - A1) + (N * N) * (A2 - A1 * A1).


/*
----------------------------------------------------------------------
beta_prime_dist  (Mathematica-compatible: BetaPrimeDistribution[p,q,α,β])

Definition:
------------
Generalized Beta of the Second Kind (GB2) distribution.

If
    Y ~ Beta(p, q)
then
    X = β * (Y / (1 - Y))^(1/α)
follows a BetaPrime/GB2 distribution.

PDF:
    f(x) =
        (α / (β * B(p, q))) *
        (x/β)^(α*p - 1) *
        (1 + (x/β)^α)^(-(p + q)),     x > 0.

CDF:
    F(x) = I_z(p, q),
        where z = (x/β)^α / (1 + (x/β)^α)
        and I_z(p, q) is the regularized incomplete beta.

Quantile:
    X = β * (Z / (1 - Z))^(1/α),
        where Z = Quantile(Beta(p, q), U).

Mean:
    E[X] = β * B(p + 1/α, q - 1/α) / B(p, q),
        defined for q > 1/α.

Variance:
    Var[X] = β² * [B(p + 2/α, q - 2/α)/B(p, q)
                    - (B(p + 1/α, q - 1/α)/B(p, q))²],
        defined for q > 2/α.

Parameters:
------------
    p > 0        (shape 1)
    q > 0        (shape 2)
    α > 0        (shape/“alpha”)
    β > 0        (scale/“beta”)

Support:
------------
    x > 0.

Implemented functions:
----------------------
    beta_prime_dist(P, Q, A, B)          % Random variate
    beta_prime_dist_pdf(P, Q, A, B, X)   % PDF
    beta_prime_dist_cdf(P, Q, A, B, X)   % CDF
    beta_prime_dist_quantile(P, Q, A, B, U) % Quantile
    beta_prime_dist_mean(P, Q, A, B)     % Mean
    beta_prime_dist_variance(P, Q, A, B) % Variance

Aliases:
---------
    beta_prime_dist(P, Q, B)             % α=1
    beta_prime_dist(P, Q)                % α=1, β=1
    beta_prime_dist_pdf(P, Q, B, X)
    beta_prime_dist_pdf(P, Q, X)
    beta_prime_dist_cdf(P, Q, B, X)
    beta_prime_dist_cdf(P, Q, X)
    beta_prime_dist_quantile(P, Q, B, U)
    beta_prime_dist_quantile(P, Q, U)
    beta_prime_dist_mean(P, Q, B)
    beta_prime_dist_variance(P, Q, B)
    beta_prime_dist_mean(P, Q)
    beta_prime_dist_variance(P, Q)

Examples:
----------
    % Standard version
    println(beta_prime_dist(10, 4)).

    % 4-parameter version (matches Mathematica)
    P=10.0, Q=4.0, A=4.0, B=3.0,

    % PDF
    println(beta_prime_dist_pdf(P, Q, A, B, 5.0)),   % ≈ 0.117201

    % CDF
    println(beta_prime_dist_cdf(P, Q, A, B, 5.0)),   % ≈ 0.8823

    % Quantile
    println(beta_prime_dist_quantile(P, Q, A, B, 0.9)), % ≈ 6.77

    % Mean and variance
    println(beta_prime_dist_mean(P, Q, A, B)),      % ≈ 3.18929
    println(beta_prime_dist_variance(P, Q, A, B)).  % ≈ 0.417617

Notes:
-------
  * Uses natural logarithms and lgamma/1 for stable computation.
  * Mean exists only for q > 1/α, variance for q > 2/α.
  * Matches Mathematica’s BetaPrimeDistribution exactly.
----------------------------------------------------------------------
*/


% BetaPrime / GB2 random variate (Mathematica-compatible)
% Parameters: P>0, Q>0 (shapes), A>0 (shape alpha), B>0 (scale beta)
% Support: X > 0
beta_prime_dist(P, Q, A, B) = X =>
    if P =< 0.0 ; Q =< 0.0 ; A =< 0.0 ; B =< 0.0 then
        throw(invalid_parameters_for_beta_prime)
    end,
    Y = beta_dist(P, Q),     % uses your existing Beta RNG
    % Y in (0,1). Transform to GB2:
    X = B * ( (Y / (1.0 - Y)) ** (1.0 / A) ).

% Aliases (consistent with Mathematica):
% 3-parameter: BetaPrimeDistribution[p,q,β]  (alpha=1)
beta_prime_dist(P, Q, B) = beta_prime_dist(P, Q, 1.0, B).

% 2-parameter: BetaPrimeDistribution[p,q] (alpha=1, beta=1)
beta_prime_dist(P, Q) = beta_prime_dist(P, Q, 1.0, 1.0).


% BetaPrime / GB2 PDF (Mathematica-compatible)
% Parameters: P>0, Q>0, A>0 (shape alpha), B>0 (scale beta)
% Support: X > 0
beta_prime_dist_pdf(P, Q, A, B, X) = F =>
    if P =< 0.0 ; Q =< 0.0 ; A =< 0.0 ; B =< 0.0 then
        throw(invalid_parameters_for_beta_prime_pdf)
    end,
    if X =< 0.0 then
        F = 0.0
    else
        T = (X / B) ** A,                  % (x/β)^α
        LBeta = lgamma(P) + lgamma(Q) - lgamma(P + Q),
        % log f = log(A) - log(B) - log B(p,q)
        %         + (A*P - 1)*log(X/B) - (P+Q)*log(1 + (X/B)^A)
        LogF = log(A) - log(B) - LBeta
               + (A * P - 1.0) * log(X / B)
               - (P + Q) * log(1.0 + T),
        F = exp(LogF)
    end.

% Aliases:
beta_prime_dist_pdf(P, Q, B, X) = beta_prime_dist_pdf(P, Q, 1.0, B, X).
beta_prime_dist_pdf(P, Q, X)    = beta_prime_dist_pdf(P, Q, 1.0, 1.0, X).


% BetaPrime / GB2 CDF (Mathematica-compatible)
% F(x) = I_{ z }(P,Q),  z = (x/B)^A / (1 + (x/B)^A)
% Parameters: P>0, Q>0, A>0 (alpha), B>0 (beta);  Support: X>0

beta_prime_dist_cdf(P, Q, A, B, X) = F =>
    if P =< 0.0 ; Q =< 0.0 ; A =< 0.0 ; B =< 0.0 then
        throw(invalid_parameters_for_beta_prime_cdf)
    end,
    if X =< 0.0 then
        F = 0.0
    else
        % L = A * (log(X) - log(B))
        L = A * (log(X) - log(B)),
        % Stable logistic for z = (X/B)^A / (1 + (X/B)^A)
        if L >= 0.0 then
            Z = 1.0 / (1.0 + exp(-L))
        else
            EL = exp(L),
            Z = EL / (1.0 + EL)
        end,
        % Delegate to Beta CDF
        F = beta_dist_cdf(P, Q, Z)
    end.

% Aliases (matching Mathematica’s shorter forms):
beta_prime_dist_cdf(P, Q, B, X) = beta_prime_dist_cdf(P, Q, 1.0, B, X).
beta_prime_dist_cdf(P, Q, X)    = beta_prime_dist_cdf(P, Q, 1.0, 1.0, X).


% BetaPrime / GB2 Quantile (Mathematica-compatible)
% Returns x = Quantile[BetaPrimeDistribution(P,Q,A,B), U]
% Parameters: P>0, Q>0, A>0 (alpha), B>0 (beta)
% Input: U in (0,1)

beta_prime_dist_quantile(P, Q, A, B, U) = X =>
    if P =< 0.0 ; Q =< 0.0 ; A =< 0.0 ; B =< 0.0 then
        throw(invalid_parameters_for_beta_prime_quantile)
    end,
    if U =< 0.0 then
        X = 0.0
    elseif U >= 1.0 then
        X = 1.0e300  % inf
    else
        Z = beta_dist_quantile(P, Q, U),
        if Z =< 0.0 then
            X = 0.0
        elseif Z >= 1.0 then
            X = 1.0e300 % inf
        else
            X = B * ((Z / (1.0 - Z)) ** (1.0 / A))
        end
    end.

% Aliases (consistent with Mathematica shorter forms)
beta_prime_dist_quantile(P, Q, B, U) = beta_prime_dist_quantile(P, Q, 1.0, B, U).
beta_prime_dist_quantile(P, Q, U)    = beta_prime_dist_quantile(P, Q, 1.0, 1.0, U).


% BetaPrime / GB2 mean (Mathematica-compatible)
% Mean = B * Beta(p + 1/α, q - 1/α) / Beta(p,q)
% Requires q > 1/α

beta_prime_dist_mean(P, Q, A, B) = M =>
    if P =< 0.0 ; Q =< 0.0 ; A =< 0.0 ; B =< 0.0 then
        throw(invalid_parameters_for_beta_prime_mean)
    elseif Q =< 1.0 / A then
        throw(mean_not_defined_for_beta_prime)
    else
        LBeta0 = lgamma(P) + lgamma(Q) - lgamma(P + Q),
        LBeta1 = lgamma(P + 1.0 / A) + lgamma(Q - 1.0 / A)
                 - lgamma(P + Q),
        M = B * exp(LBeta1 - LBeta0)
    end.

% Variance = β² * [B(p+2/α, q-2/α)/B(p,q) - (B(p+1/α, q-1/α)/B(p,q))²]
% Requires q > 2/α

beta_prime_dist_variance(P, Q, A, B) = V =>
    if P =< 0.0 ; Q =< 0.0 ; A =< 0.0 ; B =< 0.0 then
        throw(invalid_parameters_for_beta_prime_variance)
    elseif Q =< 2.0 / A then
        throw(variance_not_defined_for_beta_prime)
    else
        LBeta0 = lgamma(P) + lgamma(Q) - lgamma(P + Q),
        LBeta1 = lgamma(P + 1.0 / A) + lgamma(Q - 1.0 / A)
                 - lgamma(P + Q),
        LBeta2 = lgamma(P + 2.0 / A) + lgamma(Q - 2.0 / A)
                 - lgamma(P + Q),
        Term1 = exp(LBeta2 - LBeta0),
        Term2 = exp(2.0 * (LBeta1 - LBeta0)),
        V = B * B * (Term1 - Term2)
    end.

% Aliases (consistent shorter forms)
beta_prime_dist_mean(P, Q, B)      = beta_prime_dist_mean(P, Q, 1.0, B).
beta_prime_dist_variance(P, Q, B)  = beta_prime_dist_variance(P, Q, 1.0, B).
beta_prime_dist_mean(P, Q)         = beta_prime_dist_mean(P, Q, 1.0, 1.0).
beta_prime_dist_variance(P, Q)     = beta_prime_dist_variance(P, Q, 1.0, 1.0).


/*
  Rademacher distribution

  https://en.wikipedia.org/wiki/Rademacher_distribution
  """
  In probability theory and statistics, the Rademacher distribution (which is named after 
  Hans Rademacher) is a discrete probability distribution where a random variate X has a 50% 
  chance of being +1 and a 50% chance of being -1.
  """
*/
rademacher_dist() = Res =>
  U = bernoulli_dist(1/2),
  Res = cond(U < 1/2,-1,1).
rademacher_dist_n(N) = [rademacher_dist() : _ in 1..N].

rademacher_dist_pdf(K) = Res =>
  if K == -1 ; K == 1 then
    Res = 1/2
  else
    Res = 0
  end.

rademacher_dist_cdf(K) = Res =>
  if K < -1 then
    Res = 0
  else
    Res = cond( (K>= -1,K < 1), 1/2, 1)
  end.
  
% https://statslib.readthedocs.io/en/latest/api/rademacher.html#quantile-function
rademacher_dist_quantile(K,Q) = Res =>
  if K <= (1-Q) then
    Res = -1
  else
    Res = 1
  end.




/*
 Birthday / Coincidence distribution

  Port of R's birthday/coincidence functions (from stats library)
  """
  Probability of coincidences

  Description:

     Computes answers to a generalised _birthday paradox_ problem.
     ‘pbirthday’ computes the probability of a coincidence and
     ‘qbirthday’ computes the smallest number of observations needed to
     have at least a specified probability of coincidence.

  Usage:

     qbirthday(prob = 0.5, classes = 365, coincident = 2)
     pbirthday(n, classes = 365, coincident = 2)
     
  Arguments:

 classes: How many distinct categories the people could fall into

    prob: The desired probability of coincidence

       n: The number of people

    coincident: The number of people to fall in the same category

  Details:

     The birthday paradox is that a very small number of people, 23,
     suffices to have a 50-50 chance that two or more of them have the
     same birthday.  This function generalises the calculation to
     probabilities other than 0.5, numbers of coincident events other
     than 2, and numbers of classes other than 365.

     The formula used is approximate for ‘coincident > 2’.  The
     approximation is very good for moderate values of ‘prob’ but less
     good for very small probabilities.

  Value:

  qbirthday: Minimum number of people needed for a probability of at
             least ‘prob’ that ‘k’ or more of them have the same one out
             of ‘classes’ equiprobable labels.

  pbirthday: Probability of the specified coincidence.  
  """

  And I added rbirthday for generating random variates.

*/

% Generate random variates
rbirthday(Classes,Coincident) = qbirthday(Classes,Coincident,P) =>
  P = uniform(0,1).

expm1(X) = exp(X) - 1.

pbirthday(Classes,Coincident,N) = Res =>
  K = Coincident,
  C = Classes,
  if K < 2 then
    Res = 1
  elseif K == 2 then
    % For K == 2, compute the complement of the product probability
    Res = 1 - prod([(C-I) / C  : I in 0..N-1])
  elseif K > N then
    Res = 0
  elseif N > C * (K - 1) then
    Res = 1
  else
    % General case for K >= 3
    LHS = (N * exp( (-N) / (C*K) )) / (1 - (N / (C * (K+1))))**(1/K),
    LXX = (K * log(LHS)) - (K-1)*log(C) - lgamma(K+1),
    Res = - expm1(- exp(LXX))
  end.



log1p(X) = log(X+1).

adjust_up(C,K,N,P) = Res =>
  if pbirthday(C,K,N) < P then
    Res = adjust_up(C,K,N+1,P)
  else
    Res = N
  end.
adjust_down(C,K,N,P) = Res =>
  if pbirthday(C,K,N-1) >= P then
    Res = adjust_down(C,K,N-1,P)
  else
    Res = N
  end.
  
qbirthday(Classes,Coincident,Prob) = Res =>
  K = Coincident,
  C = Classes,
  P = Prob,

  if P <= 0 then
    Res = 1
  elseif P >= 1 then
    Res = C * (K-1) + 1
  else
    N = ceiling( exp( ((K-1)*log(C) + lgamma(K+1) + log(-log1p(-P))) / K)),
    if pbirthday(C,K,N) < P then
      Res = adjust_up(C,K,N,P)
    elseif pbirthday(C,K,N-1) >= P then
      Res = adjust_down(C,K,N,P)
    else
      Res = N
    end
  end.


% And with a more standard names
birthday_dist() = rbirthday(365,2).
birthday_dist_n(N) = [rbirthday(365,2) : _ in 1..N].

birthday_dist(Classes,Coincident) = rbirthday(Classes,Coincident).
birthday_dist_n(Classes,Coincident,N) = [rbirthday(Classes,Coincident) : _ in 1..N].

birthday_dist_pdf(X) = birthday_dist_pdf(365,2,X).
birthday_dist_pdf(Classes,Coincident,X) = pbirthday(Classes,Coincident,X).

% Nope: This does not make sense
% birthday_dist_cdf(X) = birthday_dist_cdf(365,2,X).
% birthday_dist_cdf(Classes,Coincident,X) = Res =>
%   Res = sum([birthday_dist_pdf(Classes,Coincident,I) : I in 0..X]).

birthday_dist_quantile(X) = birthday_dist_quantile(365,2,X).
birthday_dist_quantile(Classes,Coincident,X) = qbirthday(Classes,Coincident,X).

birthday_dist_median(Classes,Coincident) = birthday_dist_quantile(Classes,Coincident,0.5).


/*
  
  Probability of an r-sized run for n trial with probability p

  From see https://math.stackexchange.com/questions/417762/probability-of-20-consecutive-success-in-100-runs
  """
  There's a pretty simple recurrence relation that yields the same result. Let f(n) be the 
  possibility of getting a string of at least r successes in n trials where the possibility of 
  success in one trial is p.

  The next value of f(n+1) is the possibility of starting a new string of r successes (p^r)
  preceded by a failure (1−p), but you don't want to double count any sequence that already 
  had a string of r successes. That double counting is removed by multiplying by (1−f(n+1−r−1)).

  So the recurrence relation is:

  f(n+1)=f(n)+(1−p)∗(p^r)∗(1−f(n−r)) for n > r
  f(r)                               p^r
                                     0 otherwise

  Substituting r = 20, and p = .9 yields the same answer as @awkward above. i.e. f(100) = 0.7752991959
  """

  And we get the same result:

  Picat> X = probability_of_run_size(100,0.9,20) 
  X = 0.775299195879503

  Also see https://www.reddit.com/r/askmath/comments/rwy34y/if_i_flip_a_coin_100_times_what_is_the/
  which shows a different way for coins.

*/

/*
  Recurrence formula:
  f(n+1) = f(n)+(1−p)∗(p^r)∗(1−f(n−r))
  f(r)   = p^r
  f(_)   = 0

*/
table
probability_of_run_size(N,P,R) = Res =>
  if N < 0 ; N < R then
    Res = 0
  elseif N == R then
    Res = P**R
  else
    N1 = N-1,
    T1 = probability_of_run_size(N1,P,R),
    T2 = 1-probability_of_run_size(N1-R,P,R),
    Res = T1 + (1-P)* P**R * T2
  end.

% % Probability of exactly the length X
% probability_of_run_size_pdf(N,P,X) = Res =>
%    Ps = [0] ++ [-probability_of_run_size(N,P,V) : V in 0..N+1].differences_,
%    Res = Ps[X+1].

% % Mean
% probability_of_run_size_mean(N,P) = Res =>
%    Res = [V*probability_of_run_size_pdf(N,P,V) : V in 0..N].sum.

% probability_of_run_size_cdf(N,P,X) = Res =>
%  Res = [probability_of_run_size_pdf(N,P,V) : V in 0..X].sum.



/*
  prob-n-heads-after-k-in-max-m-tosses

  Probability of getting n heads in a row after exact k tosses,
  with max m tosses.

*/

/*

  Port of Python code from
  https://www.reddit.com/r/math/comments/4kj27s/probability_of_getting_n_heads_in_a_row_after_m/
  M: number of tosses
  N: number of heads in a row
  P: probability of success

  The last value in P is the probability of getting N heads in a row in M tosses

*/
table
prob_n_heads_after_k_in_max_m_tosses_list(P,M,N) = Ps =>
  Ps = [ 0 : _ in 1..M+1],
  Ps[N+1] := P**N,
  foreach(I in N..M-1)
    Ps[I+1+1] := Ps[I+1] + (1-Ps[I-N+1]) * (1-P) * P**N
  end.


% Generating random number of k
prob_n_heads_after_k_in_max_m_tosses_dist(P,M,N) = Res => 
  U = uniform(0,1),
  Res = prob_n_heads_after_k_in_max_m_tosses_dist_quantile(P,M,N,U).

/*

  PDF
  
  Probability of getting n tosses (with probability p of getting heads)
  in a row in m tosses in the k'th toss,

  Note that the range of k is 0..m+1.
  where k=m+1 contains the rest of the probabilities of k=0..m

  Example
  For getting 3 heads in a row in max 21 tosses after just 3 tosses is 0.125.
  For getting 3 heads in a row in max 21 tosses after 21 tosses is 0.017113685607910156
  This is the weird part of this PDF:
  The probability of not getting 3 heads in a row in max 21 tosses is when k=21+1
    Picat> prob_n_heads_after_k_in_max_m_tosses_pdf(1/2,21,3,22)
    0.19585800170898438

  The full probability of getting n tosses in a row in m tosses
  is given by the CDF (see below)
    (prob_n_heads_after_k_in_max_m_tosses_cdf(P,M,N,M)

*/
% differences_(L) = [L[I+1]-L[I] : I in 1..L.len-1].
prob_n_heads_after_k_in_max_m_tosses_dist_pdf(P,M,N,K) = Res =>
  CDF = prob_n_heads_after_k_in_max_m_tosses_list(P,M,N),
  Diffs = [0] ++ differences(CDF),
  if K > M then
    Res = 1-Diffs.sum
  else
    Res = Diffs[K+1]
  end.

/*

  CDF
  The cumulative probability of getting n heads in a row in max m tosses after k tosses
  is prob_n_heads_after_k_in_max_m_tosses_cdf(P,M,N,K)

  The total probability of getting n heads in a row in max m tosses is when k=m:
   prob_n_heads_after_k_in_max_m_tosses_cdf(P,M,N,M)

  The total probability of not getting n heads in a row in max m tosses is when k=m+1
   (prob_n_heads_after_k_in_max_m_tosses_cdf(P,M,N,1+M)
  
*/
prob_n_heads_after_k_in_max_m_tosses_dist_cdf(P,M,N,K) = Res =>
  CDF = prob_n_heads_after_k_in_max_m_tosses_list(P,M,N),
  if K > M then
    Res = 1
  else
    Res = CDF[K+1]
  end.
  
%  Quantile
prob_n_heads_after_k_in_max_m_tosses_dist_quantile(P,M,N,Q) = Res =>
  Val = _,
  foreach(K in 0..M+1,break(nonvar(Val)))
    if prob_n_heads_after_k_in_max_m_tosses_dist_cdf(P,M,N,K) >= Q then
      Val = K
    end
  end,
  Res = Val.
  

% The mean value of tosses needed to get n heads
expected_tosses_needed_for_n_heads(N) = 2 * (2**N - 1).

% Generalized version of (expected-tosses-needed-for-n-heads n)
% The mean value of tosses needed to get n successes when the
% probabilisty of success is p
expected_tosses_needed_for_n_successes(N, P) = Res =>
  Res = (1 - P**N) / ((1-P) * (P**N)).



/*
  Coupon collector distribution

  From Siegrist "Probability Mathematical Statisics and Stochastic Processes"

*/

% "Exact" probability
% from https://en.wikipedia.org/wiki/Coupon_collector%27s_problem (footnote [b])
coupon_collectors_problem_theoretical(N) = N * sum([1/(I+1) : I in 0..N-1]).

coupon_collector_dist(M,K) = coupon_collector_dist_quantile(M,K,U) => 
   U = uniform(0,1).
coupon_collector_dist_n(M,K,N) = [coupon_collector_dist(M,K) : _ in 1..N].

/*
  For k in 1..m
  The PDF of Wk is given by
  ...
  n in (k+1,k+2,k+3,...)

  What is the probability of having to buy n coupons
  in order to get k different coupons of m different coupons.
  The normal coupon collector problem is thus
    coupon_collector_dist_pdf(M,M,N)

  I.e. that we want all the M different coupons

*/
table
coupon_collector_dist_pdf(M,K,N) = Res =>
  if N < K then
    Res = 0
  else
    Res = binomialf(M-1, K-1) *
          sum([ (-1)**J *
                binomialf(K-1,J ) *
                % ( (K - J - 1) / M) ** (N-1)
                safe_pow(( (K - J - 1) / M),(N-1))
                : J in 0..K-1])
  end.

coupon_collector_dist_cdf(M,K,N) = sum([coupon_collector_dist_pdf(M,K,I) :  I in 0..N]).

% Note: This is slow for larger m and large q
coupon_collector_dist_quantile(M,K,Q) = Res =>
  I = 0, T = 0,
  OK = false,
  while (OK == false)
    if T >= Q then
      OK := I
    else
      I := I + 1,
      T := T + coupon_collector_dist_pdf(M,K,I)
    end
   
  end,
  Res = cond(OK == false, 0, OK).

coupon_collector_dist_mean(M,K) = Res =>
  Res = sum([ M / (M - I + 1)  : I in 1..K]).
  
%% Euler's approximation
coupon_collector_dist_mean_euler(N) = (N * log(N)) + EulersGamma * N + 1/2 =>
  EulersGamma = 0.5772156649.

coupon_collector_dist_variance(M,K) = Res =>
  Res = sum([ (M * (I-1)) / (M -I + 1)**2   : I in 1..K]).



/*
  Record statistics distributions


  From
  * Alexei Stepanov "On the Mathematical Theory of Records"
    https://cm.episciences.org/9528/pdf
    page 153
  * Blom, Holst, Sandell: "Problems and Snapshots from the World of Probability", p108ff

  This works fairly well for permutations and
  continuous distributions, but not for
  (discrete) distributions with (higher) probability
  of duplicates.

*/

/*
  k_record_dist(K,N),

  Probability that the kth record is n

  Note the parameter order: K is the "constant" and N is the
  parameter we search for.

  (In Gamble the parameters are (k_record_pdf n k) ) 
  
*/
% Note: Since quantile is restricted we have an restricted upper limit
k_record_dist(K) = k_record_dist_quantile(K,U) =>
  U = uniform(0,0.9).


table
k_record_dist_pdf(K,N) = Res =>
  Res = stirling1(N-1,K-1) / factorial_mem(N).


table
k_record_dist_cdf(K,N) = [k_record_dist_pdf(K,I)  : I in 1..N].sum.

/*
  Warning: For higher quantiles, e.g. 0.999 this might
  overflow. A warning is printed and the last "valid" N is
  returned.

*/
k_record_dist_quantile(K,Q) = Res =>
  OK = false,
  N = 0,
  while (OK == false)
    Val = k_record_dist_cdf(K,N),
    ValS = Val.to_string,   
    if not number(Val) ; ValS == "-nan" ; ValS == 'nan' then
      printf("k_record_dist_quantile(%d,%f): overflow. This is an estimate!\n",K,Q),
      OK := N-1,
    elseif OK == false, Val >= Q then
      OK := N
    else
      N := N + 1
    end
  end,
  Res = cond(OK == false,0,OK).

% mean value
k_record_dist_mean(K,N) = [V*k_record_dist_pdf(K,V) : V in 1..N].sum.


% Float - and perhaps faster - versions
k_record_distf_pdf(K,N) = stirling1(K-1,N-1) / exp(lgamma(N+1)).
k_record_distf_cdf(K,N) = [k_record_distf_pdf(K,I)  : I in 1..N].sum.


% Note: This is not faster than the plain version
/*
k_record_distf_quantile(K,Q) = Res =>
  OK = false,
  N = 0,
  while (OK == false)
    Val = k_record_distf_cdf(K,N),
    ValS = Val.to_string,   
    if not number(Val) ; ValS == "-nan" ; ValS == 'nan' then
      println("k_record_dist_quantile: overflow. This is an estimate!"),
      OK := N-1,
    elseif OK == false, Val >= Q then
      OK := N
    else
      N := N + 1
    end
  end,
  Res = cond(OK == false,0,OK).
*/


/*
  num_records_dist(N,K)

  Probability of K records in N values.

  Note: This works fairly well for permutations and
        continuous distributions, but not for
        (discrete) distributions with (higher) probability
        of duplicates.
*/

% Generate records
num_records_dist(N) = num_records_dist_quantile(N,U) =>
  U = uniform(0,1).

num_records_dist_pdf(N,K) = stirling1(N,K) / factorial_mem(N).
num_records_dist_cdf(N,K) = [ num_records_dist_pdf(N,I) : I in 1..K].sum.
num_records_dist_quantile(N,Q) = Res =>
  K = 1,
  OK = false,
  while (OK == false)
    Val = num_records_dist_cdf(N,K),
    if Val >= Q then
      OK := K
    else
      K := K + 1
    end,
  end,
  Res = cond(OK == false, 0, OK).

num_records_dist_mean(N) = [V*num_records_dist_pdf(N,V)  : V in 1..N].sum.

% Using harmonic number
num_records_dist_mean_h_n(N) = harmonic_number(N).




/*
  Multivariate Hypergeometric dist

  PDF and Mean from Mathematica MultivariateHypergeometricDistribution

*/
%
% Generate random instances (cf ppl_urn_model_generalized.pi)
%
multivariate_hypergeometric_dist(N,NumBalls) = Res =>
  Len = NumBalls.len,
  Res = multivariate_hypergeometric_dist1(N,NumBalls,Len,[ 0 : _ in 1..Len]).

multivariate_hypergeometric_dist1(RemainingTrials,RemainingBalls,NumBalls,Aux) = Res =>
  if RemainingTrials == 0 ;  [1 : B in RemainingBalls, B <= 0].len > 0 then
    Res = Aux
  else
    V = 1..NumBalls,
    Ball = categorical(RemainingBalls, V),
    Res = multivariate_hypergeometric_dist1(RemainingTrials-1,
                                            [ cond( Ball == I, B-1,B) : I in V, B = RemainingBalls[I]],
                                            NumBalls,
                                            [ cond( Ball == I, B+1,B) : I in V, B = Aux[I]],                                            
                                            )    
  end.

multivariate_hypergeometric_dist_n(N,NumBalls,Num) = [multivariate_hypergeometric_dist(N,NumBalls) : _ in 1..Num].

multivariate_hypergeometric_dist_pdf(N,NumBalls,Ps) = Res =>
  Res = prod([ binomialf(NumBalls[C],Ps[C]) : C in 1..NumBalls.len  ]) / binomialf(NumBalls.sum, N).


/*
  According to Mathematica: CDF = PDF
*/
multivariate_hypergeometric_dist_cdf(N,NumBalls,Ps) = multivariate_hypergeometric_dist_pdf(N,NumBalls,Ps).

multivariate_hypergeometric_dist_mean(N,NumBalls) = Res =>
  Res = [NumBalls[C] * N / NumBalls.sum :  C in 1..NumBalls.len ].

  
  
  
/*
  Generalized extreme value distribution

  Wikipedia 
  https://en.wikipedia.org/wiki/Generalized_extreme_value_distribution
  GEV(mu, sigma, xi)

*/

generalized_extreme_value_dist(Mu,Sigma,Xi) = generalized_extreme_value_dist_quantile(Mu,Sigma,Xi,U) =>
  U = uniform(0,1).
generalized_extreme_value_dist_n(Mu,Sigma,Xi,N) = [generalized_extreme_value_dist(Mu,Sigma,Xi) : _ in 1..N].


% Reduced domain for avoiding extreme values
generalized_extreme_value_dist_reduced(Mu,Sigma,Xi) = generalized_extreme_value_dist_quantile(Mu,Sigma,Xi,U) =>
  U = uniform(0.01,0.99).

generalized_extreme_value_dist_reduced_n(Mu,Sigma,Xi,N) = [generalized_extreme_value_dist_reduced(Mu,Sigma,Xi) : _ in 1..N].

generalized_extreme_value_dist_pdf(Mu,Sigma,Xi,X) = Res =>
  if Xi == 0 then
    T = exp(-(X-Mu)/Sigma)
  else
    T = (1+Xi*((X-Mu)/Sigma ))**(-1/Xi)
  end,
  if Xi == 0 ; (Xi > 0, X >= (Mu - Sigma/Xi)) ; (Xi < 0, X <= (Mu - Sigma/Xi))  then
    Res = 1/Sigma * T**(Xi+1) * exp(-T)
  else
    Res = 0
  end.



generalized_extreme_value_dist_cdf(Mu,Sigma,Xi,X) = Res =>
  if Xi == 0 then
    T = exp(-(X-Mu)/Sigma)
  else
    T = (1 + Xi * ((X-Mu)/Sigma))**(-1/Xi)
  end,
  if Xi == 0 ; (Xi > 0, X >= (Mu - Sigma/Xi)) ; (Xi < 0, X <= (Mu - Sigma/Xi))  then
    Res = exp(-T)
  else
    Res = 0
  end.

generalized_extreme_value_dist_quantile(Mu,Sigma,Xi,Q) = Res =>
  if Xi == 0 then
    Res = Mu - (Sigma * (log(- log(Q))))
  else
    Res = Mu + (Sigma/Xi) * ( (-log(Q))**(-Xi) - 1)
  end.

generalized_extreme_value_dist_mean(Mu,Sigma,Xi) = Res =>
  if Xi != 0, Xi < 1 then
    Res = Mu + (Sigma * (gamma_func(1-Xi) - 1) / Xi)
  elseif Xi == 0 then
    Res = Mu + Sigma * euler_gamma()
  else
    Res = 1.0e100    
  end.

/*
  Poisson process

  From Mathematica PoissonProcess
  PoissonProcess[Mu] 
  represents a Poisson process with rate Mu.

  PoissonProcess[Mu][t]  
  represents a Poisson process with rate Mu of time t


*/

poisson_process_dist(Mu,T) = poisson_process_dist_quantile(Mu,T,U) =>
  U = uniform(0,1). % u01().

poisson_process_dist_n(Mu,T,N) = [poisson_process_dist(Mu,T) : _ in 1..N].

poisson_process_dist_pdf(Mu,T,X) = Res =>
  if X >= 0 then
    Res = (exp(-T*Mu) * (T * Mu)**X) / factorial(X)
  else
    Res = 0
  end.

poisson_process_dist_cdf(Mu,T,X) = Res =>
  if X >= 0 then
    Res = 1-gamma_regularized(1 + floor(X),T*Mu)
  else
    Res = 0
  end.

% Reversing the CDF
poisson_process_dist_quantile(Mu,T,Q) = Res =>
  if Q == 1.0 then
    Res = 1.0e100 % infinity
  else
    N = 0,
    OK = false,
    while(OK == false)
      Val = poisson_process_dist_cdf(Mu,T,N),
      if Val >= Q then
        OK := N
      else
        N := N + 1
      end
    end,
    Res = cond(OK == false,0,OK)
  end.

poisson_process_dist_mean(Mu,T) = T*Mu.
poisson_process_dist_variance(Mu,T) = T*Mu.



/*
  Random walk process

  From Mathematica RandomWalkProcess
  RandomWalkProcess[p]
  represents a random walk on a line with the probability of a positive unit 
  step p and the probability of a negative unit step 1-p.

  RandomWalkProcess[p][t] 
  represents a random walk on a line with the probability of a positive unit 
  step p and the probability of a negative unit step 1-p at time t.

*/

random_walk_process_dist(P,T) = random_walk_process_dist_quantile(P,T,U) =>
  U = uniform(0,1).
random_walk_process_dist_n(P,T,N) = [random_walk_process_dist(P,T) : _ in 1..N].

random_walk_process_dist_pdf(P,T,X) = Res =>
  if T+X >= 0, T-X >= 0 then
    Res = (1/2) * (1 + (-1)**(T+X)) * (1-P)**((T-X)/2) * P**((T+X)/2) * binomialf_float(T,(T+X)/2)
  else
    Res = 0
  end.

random_walk_process_dist_cdf(P,T,X) = Res =>
  TX2 = (T+X)/2,
  if 0 <= TX2, TX2 < T then
    Res = reg_incomplete_beta(T-floor(TX2),1+floor(TX2),1-P)
  elseif TX2 >= T then
    Res = 1
  else
    Res = 0
  end.
  
random_walk_process_dist_quantile(P,T,Q) = Res =>
  OK = false,
  foreach(X in -T..T, break(OK != false))
    Val = random_walk_process_dist_cdf(P,T,X),
    if Val >= Q then
      OK := X
    end
  end,
  Res = cond(OK == false, 0, OK).

random_walk_process_dist_mean(P,T) = (2 * P - 1) * T.
random_walk_process_dist_variance(P,T) = 4 * (1 - P) * P * T.


/*
  Sum probability

  Distribution of summing discrete uniform numbers to a specific sum.

  Inspired by Fletcher Thompson "What’s the Probability Ten Dice Add Up To 12?"
  https://medium.com/puzzle-sphere/whats-the-probability-ten-dice-add-up-to-12-83f637205505

*/

% Number of ways to sum n integers from a to b to get s
table
count_partitions_sum(A,B,N,S) = Res =>
  if N == 1 then
    Res = cond((S >= A, S <= B),1,0)
  elseif S < N*A ; S > N * B then
    Res = 0
  else
    Res = sum([ count_partitions_sum(A,B,N-1,S-X) : X in A..B])
  end.

% Probability of getting the sum s for n random samples in the range a to b
sum_prob_dist(A,B,N) = sum_prob_dist_quantile(A,B,N,U) =>
  U = uniform(0,1).

sum_prob_dist_n(A,B,N,Num) = [sum_prob_dist(A,B,N) : _ in 1..Num].

sum_prob_dist_pdf(A,B,N,S) = Res =>
  if S < N*A ; S > N * B then
    Res = 0  
  else
    Res = count_partitions_sum(A,B,N,S) / (B - A + 1)**N
  end.

sum_prob_dist_cdf(A,B,N,S) = Res =>
  if S < A*N then
    Res = 0
  elseif S > B*N then
    Res = 1
  else
    Res = sum([sum_prob_dist_pdf(A,B,N,I)  : I in A*N..S])
  end.

sum_prob_dist_quantile(A,B,N,Q) = Res =>
  OK = false,
  foreach(I in A*N..B*N,break(OK != false))
    Val = sum_prob_dist_cdf(A,B,N,I),
    if Val >= Q then
      OK := I
    end
  end,
  Res = cond(OK == false, 0, OK).
  
sum_prob_dist_mean(A,B,N) = N * (A+B) / 2.
sum_prob_dist_variance(A,B,N) = N * ((B - A + 1)**2 -1) / 12.



/*
  Binomial_process

  From Mathematica BinomialProcess
  BinomialProcess[p] 
  represents a binomial process with event probability p.

  BinomialProcess[p][t] 
  represents a binomial process with event probability p at time t

*/
binomial_process_dist(P,T) = binomial_process_dist_quantile(P,T,U) =>
  U = uniform(0,1).

binomial_process_dist_pdf(P,T,X) = Res =>
  if 0 <= X, X <= T then
    Res = (1-P)**(T-X) * P**X * binomialf(T,X)
  else
    Res = 0
  end.

binomial_process_dist_cdf(P,T,X) = Res =>
  if 0 <= X, X <= T then
    Res = beta_reg_incomplete(T-floor(X),1+floor(X),1-P)
  elseif X >= T then
    Res = 1
  else
    Res = 0
  end.

binomial_process_dist_quantile(P,T,Q) = Res =>
  OK = false,
  foreach(X in -T..T, break(OK != false))
    Val = binomial_process_dist_cdf(P,T,X),
    if Val >= Q then
      OK := X
    end
  end,
  Res = cond(OK == false, 0, OK).

binomial_process_dist_mean(P,T) = P * T.
binomial_process_dist_variance(P,T) = (1-P) * P * T.


/*
  Wiener process, a,k.a. Brownian motion

  From Mathematica's WienerProcess

  WienerProcess[Mu,Sigma] 
  represents a Wiener process with a drift Mu and volatility Sigma.

  WienerProcess[Mu,Sigma][t]
  represents a Wiener process with a drift Mu and volatility Sigma 
  at time t.

  WienerProcess[]
  represents a standard Wiener process with drift 0 and volatility 1.
  WienerProcess[][t]
  represents a standard Wiener process with drift 0 and volatility 1
  at time t.

*/
wiener_process_dist(Mu,Sigma,T) = wiener_process_dist_quantile(Mu,Sigma,T,U) =>
  U = uniform(0,1).
wiener_process_dist_n(Mu,Sigma,T,N) = [wiener_process_dist(Mu,Sigma,T) : _ in 1..N].

wiener_process_dist_pdf(Mu,Sigma,T,X) = normal_dist_pdf(Mu*T,Sigma*sqrt(T),X).

wiener_process_dist_cdf(Mu,Sigma,T,X) = normal_dist_cdf(Mu*T,Sigma*sqrt(T),X).

wiener_process_dist_quantile(Mu,Sigma,T,Q) = normal_dist_quantile(Mu*T,Sigma*sqrt(T),Q).

wiener_process_dist_mean(Mu,_Sigma,T) = T * Mu.
wiener_process_dist_variance(_Mu,Sigma,T) = T * Sigma**2.

% wiener_process_dist/1
wiener_process_dist(T) = wiener_process_dist(0,1,T).
wiener_process_dist_n(T,N) = [wiener_process_dist(T) : _ in 1..N].

wiener_process_dist_pdf(T,X) = wiener_process_dist_pdf(0,1,T,X).
wiener_process_dist_cdf(T,X) = wiener_process_dist_cdf(0,1,T,X).
wiener_process_dist_quantile(T,X) = wiener_process_dist_quantile(0,1,T,X).

wiener_process_dist_mean(T) = wiener_process_dist_mean(0,1,T).
wiener_process_dist_variance(T) = wiener_process_dist_variance(0,1,T).



/*

  Discrete Markov process

  Inspired by Mathematica's DiscreteMarkovProcess

*/

% Generate random states for a discrete Markov process
discrete_markov_process_dist(TM,Init,T) = discrete_markov_process_dist_quantile(TM,Init,T,U) =>
  U = uniform(0,1).
discrete_markov_process_dist_n(TM,Init,T,N) = [discrete_markov_process_dist(TM,Init,T) : _ in 1..N].

  
% Probability of being in state state after t steps
% Note: This is a proper PDF over the states, but not over the time.
table
discrete_markov_process_dist_pdf_loop(TM,Init,T,State,N,K) = Res =>
  if N == 0 then
    Res = Init[K]
  else
    % Compute probability recursively, summing over previous states
    Res = sum([ % Prob of reaching state j at previous step
                discrete_markov_process_dist_pdf_loop(TM,Init,T,State,N-1,J)
                *
                TM[J,K]
               : J in 1..TM.len]),
  end.  
discrete_markov_process_dist_pdf(TM,Init,T,State) = Res =>
  Res = discrete_markov_process_dist_pdf_loop(TM,Init,T,State,T,State).
  
% Cumulative probability of being in state state after t steps
table
discrete_markov_process_dist_cdf(TM,Init,T,State) = Res =>
  StateProbs = [discrete_markov_process_dist_pdf(TM,Init,T,S)   : S in 1..TM.len].asum,
  Res = StateProbs[State].
  
% Which state have the probability q after step steps?
discrete_markov_process_dist_quantile(TM,Init,Step,Q) = Res =>
  A = [discrete_markov_process_dist_cdf(TM,Init,Step,State) : State in 1..TM.len],
  OK = false,
  foreach(State in 1..TM.len, break(OK != false))
    if A[State] >= Q then
      OK := State
    end
  end,
  Res = cond(OK == false,0,OK).


/*
  Stationary distributions of a transition matrix.

*/

/*
======================================================================
stationary_dist(P)
======================================================================
Computes the stationary distribution of transition matrix P.

This is a wrapper that:
  - validates the matrix shape
  - runs power iteration
  - checks for convergence failure
  - verifies the fixed-point condition Pi * P ≈ Pi

Throws:
  error(convergence_error,P) if matrix does not converge
  error(fixed_point_error,P) if Pi is not a true stationary point
======================================================================
*/

stationary_dist(P) = stationary_dist(P,1.0e-12,200000).
stationary_dist(P,Tolerance,MaxIter) = Pi =>
    % Tolerance = 1.0e-12,
    % MaxIter   = 200000,
    Pi0 = [1.0 / length(P) : _ in 1..length(P)],

    [Pi,Iter,_LastDiff] = stationary_power_iteration(P,Pi0,Tolerance,MaxIter),

    if Iter >= MaxIter then
        throw($error('convergence_error,P),stationary_dist/1'))
    end,

    % Verify fixed point: Pi * P ≈ Pi
    PiCheck = row_vec_times_matrix(Pi,P),
    Diff2 = max([abs(PiCheck[I] - Pi[I]) : I in 1..length(Pi)]),

    if Diff2 >= 10 * Tolerance then
        throw($error('fixed_point_error,P),stationary_dist/1'))
    end.


/*
======================================================================
stationary_power_iteration(P,Pi,Tolerance,MaxIter)
======================================================================
Performs the actual iteration.

Returns a triple:
  [PiFinal, IterationsUsed, LastDifference]
======================================================================
*/

stationary_power_iteration(P,Pi,Tolerance,MaxIter) =
  stationary_power_iteration(P,Pi,Tolerance,MaxIter,0).


stationary_power_iteration(P,Pi,Tolerance,MaxIter,Iter) = Result =>
    N = length(P),

    if Iter >= MaxIter then
        Result = [Pi,Iter,1.0]    % did not converge
    else
        PiNext0 = row_vec_times_matrix(Pi,P),
        PiNext  = normalize_dist(PiNext0),

        Diff = max([abs(PiNext[I] - Pi[I]) : I in 1..N]),

        if Diff < Tolerance then
            Result = [PiNext,Iter,Diff]
        else
            Result = stationary_power_iteration(P,PiNext,Tolerance,MaxIter,Iter+1)
        end
    end.


/*
======================================================================
row_vec_times_matrix(V,P)
----------------------------------------------------------------------
Computes W = V * P for row vector V and matrix P.
======================================================================
*/
row_vec_times_matrix(V,P) = W =>
    N = length(P),
    W0 = [0.0 : _ in 1..N],
    W = row_vec_times_matrix_acc(V,P,1,W0).


row_vec_times_matrix_acc(V,P,I,W0) = W =>
    N = length(P),
    if I > N then
        W = W0
    else
        Vi = V[I],
        Row = P[I],
        W1 = add_scaled_row(W0,Row,Vi,1),
        W = row_vec_times_matrix_acc(V,P,I+1,W1)
    end.


/*
======================================================================
add_scaled_row(W,Row,Scale,J)
======================================================================
*/
add_scaled_row(W,Row,Scale,J) = WOut =>
    N = length(Row),
    if J > N then
        WOut = W
    else
        NewVal = W[J] + Scale * Row[J],
        W1 = replace_nth(W,J,NewVal),
        WOut = add_scaled_row(W1,Row,Scale,J+1)
    end.


/*
======================================================================
normalize_dist(V)
======================================================================
*/
normalize_dist(V) = VNorm =>
    Total = sum(V),
    N = length(V),
    if Total =:= 0.0 then
        VNorm = [1.0 / N : _ in 1..N]
    else
        VNorm = [X / Total : X in V]
    end.


/*
======================================================================
replace_nth(List,N,Val)
======================================================================
*/
replace_nth([_|Rest],1,Val) = [Val|Rest].
replace_nth([H|Rest],N,Val) = [H|replace_nth(Rest,N-1,Val)] => N > 1.


/*
  This is a smaller variant
*/
stationary_dist1(P) = stationary_dist1(P, 1.0e-14).
stationary_dist1(P, Epsilon) = Pi_final =>
    N = P.length,
    Pi_initial = [ 1.0/N : _ in 1..N ],
    Pi_final = iterate_power(Pi_initial, P, Epsilon).

calculate_pi_next(Pi_current, P, N) = Pi_next =>
    Pi_next = [ sum([Pi_current[I] * P[I, J] : I in 1..N]) : J in 1..N ].

iterate_power(Pi_k, P, Epsilon) = Pi_final =>
    N = P.length,
    Pi_next = calculate_pi_next(Pi_k, P, N),
    
    % Check L1 norm for convergence
    Diff = sum([ abs(Pi_next[I] - Pi_k[I]) : I in 1..N ]),
    
    if Diff < Epsilon then
        Pi_final = Pi_next
    else
        Pi_final = iterate_power(Pi_next, P, Epsilon)
    end.



/*
  Chinese restaurant process (CRP) distribution
  https://en.wikipedia.org/wiki/Chinese_restaurant_process
  """
  In probability theory, the Chinese restaurant process is a discrete-time 
  stochastic process, analogous to seating customers at tables in a restaurant. 
  Imagine a restaurant with an infinite number of circular tables, each with infinite 
  capacity. Customer 1 sits at the first table. The next customer either sits at the 
  same table as customer 1, or the next table. This continues, with each customer 
  choosing to either sit at an occupied table with a probability proportional to the 
  number of customers already there (i.e., they are more likely to sit at a table with 
  many customers than few), or an unoccupied table.

  The definition can be generalized by introducing a parameter theta > 0 which modifies the 
  probability of the new customer sitting at a new table to theta/(n+theta) and 
  correspondingly modifies the probability of them sitting at a table of size 
    |b| / (n + theta)

  The vanilla process introduced above can be recovered by setting theta = 1.
  Intuitively,  theta can be interpreted as the effective number of customers sitting 
  at the first empty table.

  Alternative defition

  An equivalent, but subtly different way to define the Chinese restaurant process, is to 
  let new customers choose companions rather than tables. Customer n+1 chooses to sit at 
  the same table as any one of the n seated customers with probability 
    1 / (n + theta)
  or chooses to sit at a new, unoccupied table with probability 
    theta / (n + theta)
  Notice that in this formulation, the customer chooses a table without having to count 
  table occupancies---we don't need |b|.

  ...

  It can be understood as the sum of n independent Bernoulli random variables, each with 
  a different parameter:
  
    K = sum(bi,i=1..n)
    bi ~ Bernoulli theta / (i - 1 + theta)
    PMF:  gamma(theta)
          ------------     * |s(n,k)*theta^k   k = 1..N
          gamma(n + theta)

  where s denotes Stirling number of the first kind
  """

*/
crp_dist(Theta,N) = crp_dist_quantile(Theta,N,U) =>
  U = uniform(0,1).
crp_dist_n(Theta,N,Num) = [crp_dist(Theta,N) : _ in 1..Num].

crp_dist_pdf(Theta,N,K) = gamma_func(Theta) * stirling1(N,K) * Theta**K / gamma_func(N + Theta).

crp_dist_cdf(Theta,N,K) = [crp_dist_pdf(Theta,N,R) : R in 0..K].sum.

crp_dist_quantile(Theta,N,Q) = Res =>
  OK = false,
  foreach(I in 0..N, break(OK != false))
    Val = crp_dist_cdf(Theta,N,I),
    if Val >= Q then
      OK := I
    end
  end,
  Res = cond(OK == false,0,OK).

crp_dist_mean(Theta,N) = Theta * (digammaf(Theta+N) - digammaf(Theta)).