#include <iostream>
#include <cmath>
#include <boost/math/distributions/normal.hpp>
#include <boost/math/constants/constants.hpp>
#include <boost/multiprecision/cpp_dec_float.hpp>

using namespace boost::math::constants;
using namespace boost::math;
using namespace boost::multiprecision;

//using FLOAT = double;
using FLOAT = cpp_dec_float_50;

FLOAT N_f(const FLOAT &x) {
    return exp(-.5*x*x);
}
FLOAT N_finv(const FLOAT &x) {
    return sqrt(-2*log(x));
}
FLOAT N_tail_area(const FLOAT &x) {
    return root_half_pi<FLOAT>()*erfc(x * half_root_two<FLOAT>());
}

FLOAT E_f(const FLOAT &x) {
    return exp(-x);
}
FLOAT E_finv(const FLOAT &x) {
    return -log(x);
}
FLOAT E_tail_area(const FLOAT &x) {
    return exp(-x);
}

int main(int argc, char *argv[]) {
    std::function<FLOAT(const FLOAT&)> f, finv, tail_area;
    std::string which;
    unsigned size;
    bool good = false;
    if (argc == 3) {
        which = argv[1];
        size = std::stoul(argv[2]);
        if (which == "E" or which == "e") {
            f = E_f; finv = E_finv; tail_area = E_tail_area;
            good = true;
        }
        else if (which == "N" or which == "n") {
            f = N_f; finv = N_finv; tail_area = N_tail_area;
            good = true;
        }
        else {
            std::cerr << "Invalid distribution option `" << which << "'\n\n";
        }
        if (size == 0) {
            std::cerr << "Invalid size option `" << size << "'\n\n";
        }
        else {
            good = true;
            if (size < 16) std::cerr << "Warning: size seems small (" << size << ")\n\n";
            if (size % 2 == 0) std::cerr << "Warning: size is not a power of 2\n\n";
        }
    }

    if (not good) {
        std::cerr << "Usage: " << argv[0] << " {E,N} SIZE -- generate ziggurat tables for the (E)xponential or (N)ormal distribution\n\n";
        exit(1);
    }

    std::vector<FLOAT> x(size+1);
    std::vector<FLOAT> y(size+1);

    std::cout.precision(std::numeric_limits<FLOAT>::max_digits10);
    FLOAT left = 0, right = 10, last = -1;
    while (left != right) {
        FLOAT r = 0.5*(left+right);
        if (r == last) // We're at our precision limit, so stop
            break;
        last = r;
        std::cout << "trying " << r << "\n";
        x[1] = r;
        y[1] = f(x[1]);

        FLOAT A = x[1]*y[1] + tail_area(x[1]);

        x[0] = A/y[1];
        y[0] = f(x[0]);
        for (unsigned i = 2; i <= size; i++) {
            y[i] = y[i-1] + A/x[i-1];
            if (y[i] > f(0)) { // x[1] guess was too low
                left = r;
                goto NEXTGUESS;
            }
            x[i] = finv(y[i]);
        }
        // If final is negative, r was too big; if positive, too small.
        if (y[size] < f(0)) right = r;
        else left = r;
NEXTGUESS:
        ;
    }

    // Make sure y[SIZE] =~ 1, and x[SIZE] =~ 0
    if (abs(y[size] - 1) > 1e-12) throw "Error: y_n != 1";
    if (abs(x[size]) > 1e-20) throw "Error: x_n != 0";
    y[size] = 1;
    x[size] = 0;
    y[0] = 0; // This value never gets used, so just set it to 0 to slightly save space in the code

    std::string structname(which == "N" ? "normal_table" : "exponential_table");
    std::cout << "\n\n\n// tables for the ziggurat algorithm\n" <<
        "template<class RealType>\n" <<
        "struct " << structname << " {\n" <<
        "    static const RealType table_x[" << size+1 << "];\n" <<
        "    static const RealType table_y[" << size+1 << "];\n" <<
        "};\n\n";

    std::cout << "template<class RealType>\nconst RealType " << structname << "<RealType>::table_x[" << size+1 << "] = {";
    std::cout.precision(20);
    std::cout.setf(std::ios_base::showpoint);
    for (unsigned i = 0; i < size; i++) {
        std::cout << (i % 4 == 0 ? "\n    " : " ");
        if (x[i] == 0) std::cout << "0,";
        else if (x[i] == 1) std::cout << "1,";
        else std::cout << x[i] << ",";
    }
    std::cout.unsetf(std::ios_base::showpoint);
    std::cout << "\n    " << x[size] << "\n};\n\n";

    std::cout << "template<class RealType>\nconst RealType " << structname << "<RealType>::table_y[" << size+1 << "] = {";
    std::cout.precision(20);
    std::cout.setf(std::ios_base::showpoint);
    for (unsigned i = 0; i < size; i++) {
        std::cout << (i % 4 == 0 ? "\n    " : " ");
        if (y[i] == 0) std::cout << "0,";
        else if (y[i] == 1) std::cout << "1,";
        else std::cout << y[i] << ",";
    }
    std::cout.unsetf(std::ios_base::showpoint);
    std::cout << "\n    " << y[size] << "\n};\n\n";
}