1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126

// Copyright 2019-2024 ChainSafe Systems
// SPDX-License-Identifier: Apache-2.0, MIT

use std::f64::consts::E;

use statrs::function::gamma::ln_gamma;

const MAX_BLOCKS: usize = 15;
const MU: f64 = 5.0;

fn poiss_pdf(x: f64, mu: f64, cond: f64) -> f64 {
    let ln_gamma = ln_gamma(x + 1.0);
    let log_mu = mu.log(E);
    let exponent = log_mu * x - ln_gamma - cond;
    E.powf(exponent)
}

/// Calculate the number of winners for each block number, up to [`MAX_BLOCKS`].
// * This will be needed for optimal message selection
#[cfg(test)]
fn no_winners_prob() -> Vec<f64> {
    (0..MAX_BLOCKS)
        .map(|i| poiss_pdf(i as f64, MU, MU))
        .collect()
}

/// Calculate the number of winners for each block number, up to [`MAX_BLOCKS`],
/// assuming at least one winner.
fn no_winners_prob_assuming_more_than_one() -> Vec<f64> {
    let cond = (E.powf(5.0) - 1.0).log(E);
    (0..MAX_BLOCKS)
        .map(|i| poiss_pdf(i as f64, MU, cond))
        .collect()
}

fn binomial_coefficient(mut n: f64, k: f64) -> Result<f64, ()> {
    if k > n {
        return Err(());
    }

    let mut r = 1.0;
    let mut d = 1.0;
    while d <= k {
        r *= n;
        r /= d;
        n -= 1.0;
        d += 1.0;
    }
    Ok(r)
}

fn bino_pdf(x: f64, trials: f64, p: f64) -> f64 {
    // based on https://github.com/atgjack/prob
    if x > trials {
        return 0.0;
    }
    if p == 0.0 {
        if x == 0.0 {
            return 1.0;
        }
        return 0.0;
    }
    if (p - 1.0).abs() < f64::EPSILON {
        if (x - trials).abs() < f64::EPSILON {
            return 1.0;
        }
        return 0.0;
    }
    let coef = if let Ok(v) = binomial_coefficient(trials, x) {
        v
    } else {
        return 0.0;
    };

    let pow = p.powf(x) * (1.0 - p).powf(trials - x);
    coef * pow
}

pub fn block_probabilities(tq: f64) -> Vec<f64> {
    let no_winners = no_winners_prob_assuming_more_than_one();
    let p = 1.0 - tq;
    (0..MAX_BLOCKS)
        .map(|place| {
            no_winners
                .iter()
                .enumerate()
                .map(|(other_winner, p_case)| {
                    p_case * bino_pdf(place as f64, other_winner as f64, p)
                })
                .sum()
        })
        .collect()
}

#[test]
fn test_block_probability() {
    let bp = block_probabilities(1.0 - 0.15);
    for i in 0..bp.len() - 1 {
        assert!(bp[i] >= bp[i + 1]);
    }
}

#[test]
fn test_winner_probability() {
    use rand::{thread_rng, Rng};
    let n = 1_000_000;
    let winner_prob = no_winners_prob();
    let mut sum = 0.0;

    // Generates a radnom number from 0 to not including 1
    let mut rng = thread_rng();

    for _ in 0..n {
        let mut miners_rand: f64 = rng.gen::<f64>() * f64::MAX;
        for prob in winner_prob.iter().take(MAX_BLOCKS) {
            miners_rand -= prob;
            if miners_rand < 0.0 {
                break;
            }
            sum += 1.0;
        }
    }

    let avg = sum / (n as f64);
    assert!((avg - 5.0).abs() > 0.01, "Average too far off ");
}