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

#include <arrayfire.h>
#include <af/traits.hpp>

using af::array;
using af::randomEngine;
using af::randomEngineType;

template <typename T>
T chi2_statistic(array input, array expected) {
    expected *= af::sum<T>(input) / af::sum<T>(expected);
    array diff = input - expected;
    return af::sum<T>((diff * diff) / expected);
}

// should be used only for x <= 5 (roughly)
array cnd(array x)
{
    return 0.5 * erfc(- x * sqrt(0.5));
}

template <typename T>
bool testRandomEngineNormalChi2(randomEngineType type)
{
    af::dtype ty = (af::dtype)af::dtype_traits<T>::af_type;

    int elem = 1024*1024;
    int steps = 256*32;
    int bins = 100;
    T lower_edge = -7.0;
    T upper_edge = 7.0;

    array total_hist = af::constant(0.0, 2 * bins, ty);
    array edges = af::seq(bins + 1)/bins * lower_edge;
    array expected = -af::diff1(cnd(edges));
    expected = af::join(0, expected(af::seq(bins-1, 0, -1)), expected);

    af::randomEngine r(type, 0);

    // R> qchisq(c(5e-6, 1 - 5e-6), 197)
    // [1] 121.3197 297.2989
    T lower = 121.3197;
    T upper = 297.2989;

    bool prev_step = true;
    bool prev_total = true;
    af::setSeed(0x76fa214467690e3c);
    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) {
        array step_hist = af::histogram(af::randn(elem, ty, r), 2 * bins, lower_edge, upper_edge);
        T step_chi2 = chi2_statistic<T>(step_hist, expected);
        std::cout << std::setw(4) << i
                  << std::setw(7) << step_chi2;
        bool step = step_chi2 > lower && step_chi2 < upper;
        if (!(step || prev_step)) return false;
        prev_step = step;

        total_hist += step_hist;
        T total_chi2 = chi2_statistic<T>(total_hist, expected);
        std::cout << std::setw(7) << total_chi2
                  << std::endl;
        bool total = total_chi2 > lower && total_chi2 < upper;
        if (!(total || prev_total)) return false;
        prev_total = total;
    }
    return true;
}



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();

    std::cout << "PHILOX" << std::endl;
    if (!testRandomEngineNormalChi2<float>(AF_RANDOM_ENGINE_PHILOX))
        std::cout << "FAILURE" << std::endl;
    std::cout << "THREEFRY" << std::endl;
    if (!testRandomEngineNormalChi2<float>(AF_RANDOM_ENGINE_THREEFRY))
        std::cout << "FAILURE" << std::endl;
    std::cout << "MERSENNE" << std::endl;
    if (!testRandomEngineNormalChi2<float>(AF_RANDOM_ENGINE_MERSENNE))
        std::cout << "FAILURE" << std::endl;

    if (af::isDoubleAvailable(device)) {
        std::cout << "PHILOX" << std::endl;
        if (!testRandomEngineNormalChi2<double>(AF_RANDOM_ENGINE_PHILOX))
            std::cout << "FAILURE" << std::endl;
        std::cout << "THREEFRY" << std::endl;
        if (!testRandomEngineNormalChi2<double>(AF_RANDOM_ENGINE_THREEFRY))
            std::cout << "FAILURE" << std::endl;
        std::cout << "MERSENNE" << std::endl;
        if (!testRandomEngineNormalChi2<double>(AF_RANDOM_ENGINE_MERSENNE))
            std::cout << "FAILURE" << std::endl;
    }
}