/**
 * @file
 * @brief [Geometric
 * Distribution](https://en.wikipedia.org/wiki/Geometric_distribution)
 *
 * @details
 * The geometric distribution models the experiment of doing Bernoulli trials
 * until a sucess was observed. There are two formulations of the geometric
 * distribution: 1) The probability distribution of the number X of Bernoulli
 * trials needed to get one success, supported on the set { 1, 2, 3, ... } 2)
 * The probability distribution of the number Y = X − 1 of failures before the
 * first success, supported on the set { 0, 1, 2, 3, ... } Here, the first one
 * is implemented.
 *
 * Common variables used:
 * p - The success probability
 * k - The number of tries
 *
 * @author [Domenic Zingsheim](https://github.com/DerAndereDomenic)
 */

#include <cassert>   /// for assert
#include <cmath>     /// for math functions
#include <cstdint>   /// for fixed size data types
#include <ctime>     /// for time to initialize rng
#include <iostream>  /// for std::cout
#include <limits>    /// for std::numeric_limits
#include <random>    /// for random numbers
#include <vector>    /// for std::vector

/**
 * @namespace probability
 * @brief Probability algorithms
 */
namespace probability {
/**
 * @namespace geometric_dist
 * @brief Functions for the [Geometric
 * Distribution](https://en.wikipedia.org/wiki/Geometric_distribution) algorithm
 * implementation
 */
namespace geometric_dist {
/**
 * @brief Returns a random number between [0,1]
 * @returns A uniformly distributed random number between 0 (included) and 1
 * (included)
 */
float generate_uniform() {
    return static_cast<float>(rand()) / static_cast<float>(RAND_MAX);
}

/**
 * @brief A class to model the geometric distribution
 */
class geometric_distribution {
 private:
    float p;  ///< The succes probability p

 public:
    /**
     * @brief Constructor for the geometric distribution
     * @param p The success probability
     */
    explicit geometric_distribution(const float& p) : p(p) {}

    /**
     * @brief The expected value of a geometrically distributed random variable
     * X
     * @returns E[X] = 1/p
     */
    float expected_value() const { return 1.0f / p; }

    /**
     * @brief The variance of a geometrically distributed random variable X
     * @returns V[X] = (1 - p) / p^2
     */
    float variance() const { return (1.0f - p) / (p * p); }

    /**
     * @brief The standard deviation of a geometrically distributed random
     * variable X
     * @returns \sigma = \sqrt{V[X]}
     */
    float standard_deviation() const { return std::sqrt(variance()); }

    /**
     * @brief The probability density function
     * @details As we use the first definition of the geometric series (1),
     * we are doing k - 1 failed trials and the k-th trial is a success.
     * @param k The number of trials to observe the first success in [1,\infty)
     * @returns A number between [0,1] according to p * (1-p)^{k-1}
     */
    float probability_density(const uint32_t& k) const {
        return std::pow((1.0f - p), static_cast<float>(k - 1)) * p;
    }

    /**
     * @brief The cumulative distribution function
     * @details The sum of all probabilities up to (and including) k trials.
     * Basically CDF(k) = P(x <= k)
     * @param k The number of trials in [1,\infty)
     * @returns The probability to have success within k trials
     */
    float cumulative_distribution(const uint32_t& k) const {
        return 1.0f - std::pow((1.0f - p), static_cast<float>(k));
    }

    /**
     * @brief The inverse cumulative distribution function
     * @details This functions answers the question: Up to how many trials are
     * needed to have success with a probability of cdf? The exact floating
     * point value is reported.
     * @param cdf The probability in [0,1]
     * @returns The number of (exact) trials.
     */
    float inverse_cumulative_distribution(const float& cdf) const {
        return std::log(1.0f - cdf) / std::log(1.0f - p);
    }

    /**
     * @brief Generates a (discrete) sample according to the geometrical
     * distribution
     * @returns A geometrically distributed number in [1,\infty)
     */
    uint32_t draw_sample() const {
        float uniform_sample = generate_uniform();
        return static_cast<uint32_t>(
                   inverse_cumulative_distribution(uniform_sample)) +
               1;
    }

    /**
     * @brief This function computes the probability to have success in a given
     * range of tries
     * @details Computes P(min_tries <= x <= max_tries).
     * Can be used to calculate P(x >= min_tries) by not passing a second
     * argument. Can be used to calculate P(x <= max_tries) by passing 1 as the
     * first argument
     * @param min_tries The minimum number of tries in [1,\infty) (inclusive)
     * @param max_tries The maximum number of tries in [min_tries, \infty)
     * (inclusive)
     * @returns The probability of having success within a range of tries
     * [min_tries, max_tries]
     */
    float range_tries(const uint32_t& min_tries = 1,
                      const uint32_t& max_tries =
                          std::numeric_limits<uint32_t>::max()) const {
        float cdf_lower = cumulative_distribution(min_tries - 1);
        float cdf_upper = max_tries == std::numeric_limits<uint32_t>::max()
                              ? 1.0f
                              : cumulative_distribution(max_tries);
        return cdf_upper - cdf_lower;
    }
};
}  // namespace geometric_dist
}  // namespace probability

/**
 * @brief Tests the sampling method of the geometric distribution
 * @details Draws 1000000 random samples and estimates mean and variance
 * These should be close to the expected value and variance of the given
 * distribution to pass.
 * @param dist The distribution to test
 */
void sample_test(
    const probability::geometric_dist::geometric_distribution& dist) {
    uint32_t n_tries = 1000000;
    std::vector<float> tries;
    tries.resize(n_tries);

    float mean = 0.0f;
    for (uint32_t i = 0; i < n_tries; ++i) {
        tries[i] = static_cast<float>(dist.draw_sample());
        mean += tries[i];
    }

    mean /= static_cast<float>(n_tries);

    float var = 0.0f;
    for (uint32_t i = 0; i < n_tries; ++i) {
        var += (tries[i] - mean) * (tries[i] - mean);
    }

    // Unbiased estimate of variance
    var /= static_cast<float>(n_tries - 1);

    std::cout << "This value should be near " << dist.expected_value() << ": "
              << mean << std::endl;
    std::cout << "This value should be near " << dist.variance() << ": " << var
              << std::endl;
}

/**
 * @brief Self-test implementations
 * @returns void
 */
static void test() {
    probability::geometric_dist::geometric_distribution dist(0.3);

    const float threshold = 1e-3f;

    std::cout << "Starting tests for p = 0.3..." << std::endl;
    assert(std::abs(dist.expected_value() - 3.33333333f) < threshold);
    assert(std::abs(dist.variance() - 7.77777777f) < threshold);
    assert(std::abs(dist.standard_deviation() - 2.788866755) < threshold);
    assert(std::abs(dist.probability_density(5) - 0.07203) < threshold);
    assert(std::abs(dist.cumulative_distribution(6) - 0.882351) < threshold);
    assert(std::abs(dist.inverse_cumulative_distribution(
                        dist.cumulative_distribution(8)) -
                    8) < threshold);
    assert(std::abs(dist.range_tries() - 1.0f) < threshold);
    assert(std::abs(dist.range_tries(3) - 0.49f) < threshold);
    assert(std::abs(dist.range_tries(5, 11) - 0.2203267f) < threshold);
    std::cout << "All tests passed" << std::endl;
    sample_test(dist);

    dist = probability::geometric_dist::geometric_distribution(0.5f);

    std::cout << "Starting tests for p = 0.5..." << std::endl;
    assert(std::abs(dist.expected_value() - 2.0f) < threshold);
    assert(std::abs(dist.variance() - 2.0f) < threshold);
    assert(std::abs(dist.standard_deviation() - 1.4142135f) < threshold);
    assert(std::abs(dist.probability_density(5) - 0.03125) < threshold);
    assert(std::abs(dist.cumulative_distribution(6) - 0.984375) < threshold);
    assert(std::abs(dist.inverse_cumulative_distribution(
                        dist.cumulative_distribution(8)) -
                    8) < threshold);
    assert(std::abs(dist.range_tries() - 1.0f) < threshold);
    assert(std::abs(dist.range_tries(3) - 0.25f) < threshold);
    assert(std::abs(dist.range_tries(5, 11) - 0.062011f) < threshold);
    std::cout << "All tests passed" << std::endl;
    sample_test(dist);

    dist = probability::geometric_dist::geometric_distribution(0.8f);

    std::cout << "Starting tests for p = 0.8..." << std::endl;
    assert(std::abs(dist.expected_value() - 1.25f) < threshold);
    assert(std::abs(dist.variance() - 0.3125f) < threshold);
    assert(std::abs(dist.standard_deviation() - 0.559016f) < threshold);
    assert(std::abs(dist.probability_density(5) - 0.00128) < threshold);
    assert(std::abs(dist.cumulative_distribution(6) - 0.999936) < threshold);
    assert(std::abs(dist.inverse_cumulative_distribution(
                        dist.cumulative_distribution(8)) -
                    8) < threshold);
    assert(std::abs(dist.range_tries() - 1.0f) < threshold);
    assert(std::abs(dist.range_tries(3) - 0.04f) < threshold);
    assert(std::abs(dist.range_tries(5, 11) - 0.00159997f) < threshold);
    std::cout << "All tests have successfully passed!" << std::endl;
    sample_test(dist);
}

/**
 * @brief Main function
 * @return 0 on exit
 */
int main() {
    srand(time(nullptr));
    test();  // run self-test implementations
    return 0;
}