#include <stdio.h>
#include <stdlib.h>

#include <vector>
#include <algorithm>
using std::vector;
using std::swap;

#include "spigot.h"
#include "funcs.h"
#include "cr.h"
#include "error.h"

/*
 * Riemann zeta function zeta(s), for real s
 * -----------------------------------------
 *
 * We begin with the alternating-series representation of zeta
 *
 *                1         -s    -s    -s    -s    -s    -s
 * zeta(s) = ----------- ( 1   - 2   + 3   - 4   + 5   - 6   + ... )
 *           1 - 2^{1-s}
 *
 * (which is easy to derive from the obvious definition by observing
 * that 2^{-s} zeta(s) contains only the even terms of the series for
 * zeta(s) itself, so that subtracting twice the former from the latter
 * has the effect of _negating_ the even terms).
 *
 * Proposition 2 in [1] (section 1.2) defines a formula for computing
 * zeta(s) with a good error bound, by summing 2n terms of that
 * alternating series, except that terms n+1,...,2n are multiplied by
 * gradually dwindling correction factors derived from binomial
 * coefficients.
 *
 * The absolute error in approximating (1-2^{1-s})zeta(s) with a
 * 2n-term version of this series is bounded by a formula that works
 * out (when s is real) to a maximum absolute error of 8^{-n}, which
 * is very nice indeed compared to the incredibly slow convergence of
 * just summing the obvious infinite series for zeta and waiting.
 *
 * This isn't a very spigot-shaped algorithm, because of the need to
 * commit up front to the value of n, i.e. how much precision you
 * want. To spigotise it, we must loop round for larger and larger n;
 * things fall out reasonably nicely if we make n a power of two, in
 * that we get approximations in the form of dyadic rationals which we
 * can feed to a BinaryIntervalSource.
 *
 * But that inelegance is more than made up for by the fact that the
 * formula works even for s < 0! In spite of the fact that the full
 * infinite sum of n^{-s} diverges in that part of the domain, the
 * tapering-off series defined by Gourdon & Sebah still keeps working.
 *
 * [1] "Numerical evaluation of the Riemann Zeta-function", Xavier
 * Gourdon and Pascal Sebah, 2003. As of 2015-11-05, available at
http://numbers.computation.free.fr/Constants/Miscellaneous/zetaevaluations.pdf
 */

class IntInterval {
    /*
     * Small class to permit error-bound tracking via interval
     * arithmetic.
     */
    bigint lo, hi;
  public:
    IntInterval() = default;
    IntInterval(const bigint &alo, const bigint &ahi);
    IntInterval(const bigint &n);
    IntInterval(int n);
    IntInterval(unsigned n);
    static IntInterval with_error(const bigint &n, const bigint &err);
    void increase_error(const bigint &moreerr);
    void get_median_and_error(bigint *median, bigint *error);
    bigint max_magnitude() const;
    const IntInterval &operator=(const IntInterval &that);
    const IntInterval &operator+=(const IntInterval &that);
    const IntInterval &operator-=(const IntInterval &that);
    const IntInterval &operator*=(const IntInterval &that);
    const IntInterval &operator*=(const bigint &that);
    const IntInterval &operator/=(const bigint &that);
    const IntInterval &operator>>=(const bigint &that);
    IntInterval operator+(const IntInterval &that) const;
    IntInterval operator-(const IntInterval &that) const;
    IntInterval operator*(const IntInterval &that) const;
    IntInterval operator*(const bigint &that) const;
    IntInterval operator/(const bigint &that) const;
    IntInterval operator>>(const bigint &that) const;
    friend IntInterval operator*(const bigint &a, const IntInterval &b);
};

IntInterval::IntInterval(const bigint &alo, const bigint &ahi)
    : lo(alo), hi(ahi)
{
    assert(lo <= hi);
}

IntInterval::IntInterval(const bigint &n)
    : lo(n), hi(n)
{
}

IntInterval::IntInterval(int n)
    : lo(n), hi(n)
{
}

IntInterval::IntInterval(unsigned n)
    : lo(n), hi(n)
{
}

IntInterval IntInterval::with_error(const bigint &n, const bigint &err)
{
    return IntInterval(n - err, n + err);
}

void IntInterval::increase_error(const bigint &moreerr)
{
    lo -= moreerr;
    hi += moreerr;
    assert(lo <= hi);
}

void IntInterval::get_median_and_error(bigint *median, bigint *error)
{
    *median = fdiv(lo + hi, 2);
    *error = hi - *median;     // fdiv guarantees hi-med >= med-lo
}

bigint IntInterval::max_magnitude() const
{
    bigint ret = bigint_abs(lo);
    bigint ret2 = bigint_abs(hi);
    return ret > ret2 ? ret : ret2;
}

const IntInterval &IntInterval::operator=(const IntInterval &that)
{
    lo = that.lo;
    hi = that.hi;
    return *this;
}

const IntInterval &IntInterval::operator+=(const IntInterval &that)
{
    lo += that.lo;
    hi += that.hi;
    assert(lo <= hi);
    return *this;
}

const IntInterval &IntInterval::operator-=(const IntInterval &that)
{
    lo -= that.hi;
    hi -= that.lo;
    assert(lo <= hi);
    return *this;
}

const IntInterval &IntInterval::operator*=(const IntInterval &that)
{
    bigint newlo, newhi, n;
    newlo = newhi = lo * that.lo;
    n = lo * that.hi; if (newlo > n) newlo = n; if (newhi < n) newhi = n;
    n = hi * that.lo; if (newlo > n) newlo = n; if (newhi < n) newhi = n;
    n = hi * that.hi; if (newlo > n) newlo = n; if (newhi < n) newhi = n;
    lo = newlo;
    hi = newhi;
    assert(lo <= hi);
    return *this;
}

const IntInterval &IntInterval::operator*=(const bigint &that)
{
    lo *= that;
    hi *= that;
    if (lo > hi)
        swap(lo, hi);
    return *this;
}

const IntInterval &IntInterval::operator/=(const bigint &that)
{
    lo = fdiv(lo, that);
    hi = -fdiv(-hi, that);
    if (lo > hi)
        swap(lo, hi);
    return *this;
}

const IntInterval &IntInterval::operator>>=(const bigint &that)
{
    lo >>= that;
    hi = -((-hi) >> that);
    assert(lo <= hi);
    return *this;
}

IntInterval IntInterval::operator+(const IntInterval &that) const
{
    IntInterval ret = *this; ret += that; return ret;
}

IntInterval IntInterval::operator-(const IntInterval &that) const
{
    IntInterval ret = *this; ret -= that; return ret;
}

IntInterval IntInterval::operator*(const IntInterval &that) const
{
    IntInterval ret = *this; ret *= that; return ret;
}

IntInterval IntInterval::operator*(const bigint &that) const
{
    IntInterval ret = *this; ret *= that; return ret;
}

IntInterval IntInterval::operator/(const bigint &that) const
{
    IntInterval ret = *this; ret /= that; return ret;
}

IntInterval IntInterval::operator>>(const bigint &that) const
{
    IntInterval ret = *this; ret >>= that; return ret;
}

IntInterval operator*(const bigint &a, const IntInterval &b)
{
    return b * a;
}

/*
 * Internal helper function returning a dyadic rational approximation
 * to zeta(s) * (1-2^(1-s)).
 */
static void spigot_zeta_series_approx(StaticGenerator *gen_minus_s,
                                      StaticGenerator *gen_2_to_the_minus_s,
                                      unsigned log2_n,
                                      bigint *retn, bigint *err,
                                      unsigned *dshift)
{
    bigint cn, termlimit, bcoeff, d, n, n2, r, sign;
    IntInterval sum;

    /*
     * Gourdon & Sebah's formula combines 2n values of n^{-s} to give
     * an approximation to zeta(s) with absolute error at most
     * 2^{-3n}. There will be further error arising from the way we
     * compute the n^{-s} values, which we'll track carefully as we go
     * along below; but right now we need to choose a denominator to
     * multiply all of this by, and it should be 2^{3n} plus a bit.
     */
    unsigned three_n = 3U << log2_n;
    unsigned i = three_n + log2_n;
    // Add another factor of 2^n from the series coefficients
    unsigned d_extra = (1U << log2_n);
    *dshift = i + d_extra;
    d = 1_bi << i;
    cn = 1_bi << (1U << log2_n);
    termlimit = cn;

    sign = +1;
    sum = 0;
    bcoeff = 1;

    n = 1_bi << log2_n;
    n2 = 2*n;

    /*
     * Array to hold the (-s)th powers of 1,2,...,n2. We leave the 0th
     * array element uninitialised (or rather default-initialised) and
     * will never care about it.
     */
    vector<IntInterval> powers((unsigned)n2 + 1);

    /*
     * Initialise the first two elements of powers[] before we start
     * our main loop summing the series. The rest will be computed as
     * we go along. We'll also need the value of -s itself.
     */
    powers[1] = d;
    /* get_approximate_approximant returns a value within 2 of the
     * true one. Set up the initial IntIntervals accordingly. */
    powers[2] = IntInterval::with_error
        (gen_2_to_the_minus_s->get_approximate_approximant(d), 2);
    IntInterval minus_s = IntInterval::with_error
        (gen_minus_s->get_approximate_approximant(d), 2);

    /*
     * Main loop which constructs each powers[] value and adds them
     * together with the Gourdon & Sebah series coefficients. Most
     * addition is done using the IntInterval class defined above.
     */
    bigint kb;
    for (unsigned k = 1; (kb = bigint(k)) <= n2; ++k) {
        /*
         * Construct powers[k].
         */
        if (k <= 2) {
            /* Already done before this loop started */
        } else if (k % 2 == 0) {
            /*
             * We handle even numbers by multiplying 2^{-s} by
             * (k/2)^{-s}.
             *
             * We could do absolutely everything using the binomial
             * theorem technique in the else clause below, but that
             * would accumulate error linearly along the powers[]
             * array. Handling even numbers using this more accurate
             * technique means that any given value in powers[] will
             * be derived from at most O(log n) successive binomial-
             * theorem adjustments, instead of O(n).
             *
             * (We could go in completely the other direction, by
             * computing accurate values for all primes via spigotry,
             * and then extend this simple multiplication approach to
             * all composite numbers. But that would slow things down
             * a lot, according to my experimentation, because spigots
             * are very slow. This seems a reasonable compromise.)
             */
            powers[k] = (powers[k/2] * powers[2]) >> i;
        } else {
            /*
             * Odd number. In this case we adjust from the previous
             * value, by using the generalised binomial theorem to
             * compute a correction factor. We want to turn (k-1)^{-s}
             * into k^{-s}; the ratio between those values works out
             * to (1 + 1/(k-1))^{-s}, which the binomial theorem
             * expands (writing 'prev' for k-1) as
             *
             *            -s          -s      -s(-s-1)   -s(-s-1)(-s-2)
             *   (1+prev)    = 1 + ------- + --------- + -------------- + ...
             *                     1! prev   2! prev^2      3! prev^3
             *
             * We sum that series until either a term becomes zero (if
             * -s was a positive integer), meaning the binomial series
             * was finite and has terminated, or because the terms of
             * an infinite one gets small enough that we're confident
             * we can bound the error introduced by leaving out the
             * tail of the series. In the latter case, we must work
             * out that error bound.
             *
             * The jth term of this series has the form
             *
             *    -s(-s-1)(-s-2)...(-s-j+1)
             *    -------------------------
             *            j! prev^j
             *
             * If we can show that there's some x, with |x|<1, such
             * that each term's magnitude is at most x times the
             * previous one, then we can bound the truncation error by
             * the sum of a GP, which makes the truncation error at
             * most the size of the first omitted term times 1/(1-x).
             *
             * Well, each new term differs from the previous one by
             * the factor (-s-j+1)/(j prev), so we need that to be
             * bounded by something less than 1, i.e. we need to know
             * |j prev| > |-s-j+1|. But if those two numbers are still
             * _almost_ equal, the 1/(1-x) factor might still be a bit
             * big, so we might want to carry on summing further
             * terms. I think a reasonable stopping point is to wait
             * until the inter-term ratio is at most 2/3, so that the
             * series tail is at most 3 times the size of the last
             * included term; that means I need to have 2 |j prev| > 3
             * |-s-j+1|.
             */
            unsigned prev = k-1;
            IntInterval term = powers[prev]; // first term of the series
            IntInterval total = 0;           // sum of the series
            IntInterval nfactor = minus_s;   // inter-term ratio numerator
            bigint dfactor = prev;           // inter-term ratio denominator
            bool ok_to_terminate = false;
            while (!ok_to_terminate || term.max_magnitude() > termlimit) {
                total += term;
                term = (term * nfactor / dfactor) >> i;
                nfactor -= 1_bi << i;
                dfactor += prev;
                if (!ok_to_terminate &&
                    2*dfactor*d > 3*nfactor.max_magnitude())
                    ok_to_terminate = true;
            }

            /* Max error from truncating the binomial series, as
             * analysed above. */
            total.increase_error(term.max_magnitude() * 3);

            powers[k] = total;
        }

        sum += sign * cn * powers[k];
        if (kb >= n) {
            cn -= bcoeff;
            bcoeff = bcoeff * (n2-kb) / (kb-n+1);
        }
        sign = -sign;
    }

    /* 2^{-3n} error inherent in this formula for zeta */
    sum.increase_error(1_bi << (i - three_n + d_extra));

    sum.get_median_and_error(retn, err);
}

class ZetaSeries : public BinaryIntervalSource {
    unsigned log2_n;
    StaticGenerator gen_minus_s, gen_2_to_the_minus_s;
    Spigot s;

    virtual Spigot clone() override;
    virtual void gen_bin_interval(bigint *ret_lo, bigint *ret_hi,
                                  unsigned *ret_bits) override;

  public:
    ZetaSeries(Spigot as);
};

MAKE_CLASS_GREETER(ZetaSeries);

ZetaSeries::ZetaSeries(Spigot as)
    : log2_n(3),
      gen_minus_s(spigot_neg(as->clone())),
      gen_2_to_the_minus_s(spigot_exp2(spigot_neg(as->clone()))),
      s(move(as))
{
    dgreet();
}

Spigot ZetaSeries::clone()
{
    return spigot_clone(this, s->clone());
}

void ZetaSeries::gen_bin_interval(bigint *ret_lo, bigint *ret_hi,
                                  unsigned *ret_bits)
{
    bigint n, err;
    unsigned dshift;
    spigot_zeta_series_approx(&gen_minus_s, &gen_2_to_the_minus_s,
                              log2_n, &n, &err, &dshift);
    dprint("zeta approx = (", n, " +- ", err, ") / 2^", dshift);
    *ret_bits = dshift;
    *ret_lo = n - err;
    *ret_hi = n + err;
    log2_n++;
}

static Spigot spigot_zetaseries(Spigot s)
{
    return make_unique<ZetaSeries>(move(s));
}

static const UnaryDebugFnWrapper
expr_ZetaSeries("ZetaSeries", spigot_zetaseries);

static Spigot spigot_zeta(Spigot s)
{
    Spigot multiplier = spigot_sub(spigot_integer(1),
                                   spigot_pow(spigot_integer(2),
                                              spigot_sub(spigot_integer(1),
                                                         s->clone())));
    return spigot_div(spigot_zetaseries(move(s)), move(multiplier));
}

static const UnaryFnWrapper expr_zeta("zeta", spigot_zeta);
