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

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

/*
 * Riemann zeta function zeta(s), for real s
 * -----------------------------------------
 *
 * The basic method is taken from [1], section 1.2, Proposition 2,
 * which defines a formula for computing zeta(s) with a good error
 * bound by summing 2n terms of the obvious series n^{-s}, except that
 * terms n+1,...,2n are multiplied by gradually dwindling correction
 * factors derived from binomial coefficients. This tapering-off of
 * the sum causes the output value to be wrong by a factor of
 * 1-2^{1-s}, but that's easily multiplied back in afterwards; and 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 < 1! 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() {}
    IntInterval(const bigint &alo, const bigint &ahi) : lo(alo), hi(ahi) {
        assert(lo <= hi);
    }
    IntInterval(const bigint &n) : lo(n), hi(n) {}
    IntInterval(int n) : lo(n), hi(n) {}
    IntInterval(unsigned n) : lo(n), hi(n) {}
    static IntInterval with_error(const bigint &n, const bigint &err) {
        return IntInterval(n - err, n + err);
    }
    void increase_error(const bigint &moreerr) {
        lo -= moreerr;
        hi += moreerr;
    }
    void get_median_and_error(bigint *median, bigint *error) {
        *median = fdiv(lo + hi, 2);
        *error = hi - *median;     // fdiv guarantees hi-med >= med-lo
    }
    bigint max_magnitude() {
        bigint ret = bigint_abs(lo);
        bigint ret2 = bigint_abs(hi);
        return ret > ret2 ? ret : ret2;
    }
    const IntInterval &operator=(const IntInterval &that) {
        lo = that.lo;
        hi = that.hi;
        return *this;
    }
    const IntInterval &operator+=(const IntInterval &that) {
        lo += that.lo;
        hi += that.hi;
        return *this;
    }
    const IntInterval &operator-=(const IntInterval &that) {
        lo -= that.lo;
        hi -= that.hi;
        return *this;
    }
    const 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;
        return *this;
    }
    const IntInterval &operator*=(const bigint &that) {
        lo *= that;
        hi *= that;
        return *this;
    }
    const IntInterval &operator/=(const bigint &that) {
        lo = fdiv(lo, that);
        hi = -fdiv(-hi, that);
        return *this;
    }
    const IntInterval &operator>>=(const bigint &that) {
        lo >>= that;
        hi = -((-hi) >> that);
        return *this;
    }
    IntInterval operator+(const IntInterval &that) const {
        IntInterval ret = *this; ret += that; return ret;
    }
    IntInterval operator-(const IntInterval &that) const {
        IntInterval ret = *this; ret -= that; return ret;
    }
    IntInterval operator*(const IntInterval &that) const {
        IntInterval ret = *this; ret *= that; return ret;
    }
    IntInterval operator*(const bigint &that) const {
        IntInterval ret = *this; ret *= that; return ret;
    }
    IntInterval operator/(const bigint &that) const {
        IntInterval ret = *this; ret /= that; return ret;
    }
    IntInterval operator>>(const bigint &that) const {
        IntInterval ret = *this; ret >>= that; return ret;
    }
    friend IntInterval operator*(const bigint &a, const IntInterval &b);
};
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 = (bigint)1 << i;
    cn = (bigint)1 << (1U << log2_n);
    termlimit = cn;

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

    n = (bigint)1 << 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.
     */
    std::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.
     */
    for (unsigned k = 1; (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 -= (bigint)1 << 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 ((bigint)k >= n) {
            cn -= bcoeff;
            bcoeff = bcoeff * (n2-(bigint)k) / ((bigint)k-n+1);
        }
        sign = -sign;
    }

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

    sum.get_median_and_error(retn, err);
}

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

  public:
    ZetaSeries(Spigot *as) : s(as->clone()),
                             log2_n(3),
                             gen_minus_s(spigot_neg(as->clone())),
                             gen_2_to_the_minus_s(spigot_exp2
                                                  (spigot_neg(as->clone())))
    {
        dprint("hello ZetaSeries");
    }

    virtual ~ZetaSeries()
    {
        delete s;
    }

    virtual ZetaSeries *clone()
    {
        return new ZetaSeries(s);
    }

    virtual void 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);
        *ret_bits = dshift;
        *ret_lo = n - err;
        *ret_hi = n + err;
        log2_n++;
    }
};

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(new ZetaSeries(s), multiplier);
}
