#include <iostream>
#include <iomanip>
#include <cassert>

#include <arrayfire.h>

// normalized chi^2 statistic for uniform distribution
// multiply with number of samples to get chi^2 statistic
float uchi2(af::array input) {
    assert(input.isvector());
    double expected = 1.0 / input.elements();
    af::array diff = input - expected;
    return af::sum<float>((diff * diff) / expected);
}

int main(int argc, char ** argv) {
    int backend = argc > 1 ? atoi(argv[1]) : 0;
    af::setBackend(static_cast<af::Backend>(backend));
    int device = argc > 2 ? atoi(argv[2]) : 0;
    af::setDevice(device);
    af::info();

    double samples = 1e8;
    int steps = 300;
    int bins = 100;

    af::array total_hist = af::constant(0, bins, f32);
    af::array step_hist = af::constant(0, bins, f32);

    af::setDefaultRandomEngineType(AF_RANDOM_ENGINE_PHILOX);
    //af::setDefaultRandomEngineType(AF_RANDOM_ENGINE_THREEFRY);
    //af::setDefaultRandomEngineType(AF_RANDOM_ENGINE_MERSENNE);
    std::cout << std::setw(4) << "step"
              << std::setw(7) << "chi2_i"
              << std::setw(7) << "chi2_t"
              << std::setprecision(2) << std::fixed << std::endl;
    for(int i = 0; i < steps; ++i) {
        step_hist = af::histogram(af::randu(samples), bins, 0.0, 1.0) / samples;
        total_hist = (total_hist * i + step_hist) / (i + 1);
        std::cout << std::setw(4) << i
                  << std::setw(7) << uchi2(step_hist) * samples
                  << std::setw(7) << uchi2(total_hist) * samples * (i +1)
                  << std::endl;
    }
}