/*
 * gamma.cpp: Gamma functions (lgamma, tgamma/gamma).
 */

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

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

/*
 * Computing the gamma, or even log-gamma, function in a spigoty
 * fashion (i.e. generating more and more digits until told to stop)
 * is _hard_.
 *
 * I know of four basic strategies for gamma computation, which are
 * listed below.
 *
 * The Lanczos approximation has a generally good reputation on the
 * Internet. It computes the gamma function proper, by means of
 * taking some terms looking like Stirling's approximation
 * (sqrt(2*pi), z plus a small constant to the power of z plus
 * another small constant, exp(-z)) and then multiplying in a
 * corrective term generated by a series. Supposedly this yields
 * good convergence right across the plausible positive range of the
 * function (which, for these purposes, I'm defining as up to about
 * 180-200). However, for this spigot framework it's completely
 * unviable because the coefficients of the series that gives the
 * Lanczos corrective term are irrational: to translate any kind of
 * series into a sequence of spigot matrices I need the coefficients
 * to be rational.
 *
 * There's another technique which I pulled off Wikipedia that works
 * on the primary interval [1,2]: you pick a large x, and then sum
 * an infinite series involving powers of x divided by terms looking
 * like z(z+1)(z+2)...(z+n). Unfortunately, this series is not exact
 * - even after the series has been summed to infinity and has
 * converged, there's _still_ a fixed error term remaining, given by
 * an integral. One can make that error term arbitrarily small by
 * choosing x sufficiently large, but again, this is an unviable
 * approach for spigot computation since we need the series itself
 * to sum to the _right_ value.
 *
 * A third approach is based on the lower and upper _incomplete_ gamma
 * functions. These are two-argument functions defined as the integral
 * from 0 to x, and x to infinity (respectively), of the same function
 * of x,t of which gamma itself is the integral from 0 to inf. In
 * other words, lower_gamma(x,t) + upper_gamma(x,t) = gamma(t) for any
 * x. And these functions have convenient expressions as continued
 * fractions, which are readily translated into spigot form, and so
 * you can pick the simplest possible x (I like 1) and add them
 * together. This does well for small input values of gamma, but
 * becomes painfully slow for large ones.
 *
 * That leaves Stirling's approximation. Wikipedia provides a
 * series, described in detail below, which uses rational
 * coefficients and converges exactly to the difference between the
 * well-known Stirling approximation and the exact value of the
 * log(gamma) function. Just what we wanted, right?
 *
 * ... sort of. In fact, even _this_ series isn't usable in the
 * obvious way, because the form of the terms of the series consists
 * of coefficients which grow roughly as n!, divided by denominators
 * which are successively z+1, (z+1)(z+2), (z+1)(z+2)(z+3) and so
 * on. For large z this starts off converging impressively fast
 * (unsurprisingly since Stirling's approximation is asymptotic, and
 * differs less and less from the real value of log(gamma) as z
 * increases), but if you look a little more closely you find that
 * the nth term is about n/(z+n) times the (n-1)th term - and while
 * that starts off looking very attractive for large z and small n,
 * for any z you will eventually reach large enough n that one term
 * is really not much smaller than the previous one, whereupon
 * convergence grinds to a halt. (Wikipedia assures me the series
 * does converge in theory, but then so does the series
 * 1-1/2+1/3-1/4... for log(2) and I wouldn't want to sum that one
 * in practice either.)
 *
 * Of course, you can _make_ z as large as you like, if it wasn't:
 * simply use the recurrence relation gamma(z+1)=z*gamma(z) to
 * convert the problem of computing lgamma of a small number into
 * computing lgamma of a large number and subtracting off some nice
 * simple logarithms. But this still has the irritating feature that
 * you have to decide _in advance_ what precision you want the
 * answer to. The precision limit is imposed by the convergence
 * slowing down in this case, whereas in the other approach above
 * it's imposed by an unavoidable residual error, but either way,
 * it's not very practical for spigot computation.
 *
 * My solution is to do a similar trick to the one I did in
 * monotone.cpp, involving periodically junking the computation and
 * starting it again with better parameters. The plan is: we
 * construct a spigot class that computes the Stirling correction
 * series in the obvious way - except that it tracks the ratio
 * between successive terms of the series, and when that ratio
 * starts getting too small, it gives up on the computation and
 * returns something blatantly untrue. It also sets a flag in itself
 * to indicate that it has done so.
 *
 * So we then have a wrapper class which computes the full lgamma of
 * a rational: it sets up a complicated spigot tree to compute all
 * the obvious Stirling bits and pieces and add them together, and
 * one of the additive terms in this tree is an instance of the
 * above-described class. Then we fetch data from the front of the
 * whole spigot, and after retrieving each refinement we _check_ to
 * see whether the object that was summing the series has set its "I
 * lied" flag. If not, then we can be confident of the term we're
 * holding, and can safely output it.
 *
 * When we find that the "I lied" flag _is_ set, then we know the
 * coefficient we're holding was derived from false information, so
 * we throw it away. Then we also throw away the entire spigot tree,
 * discard the coefficient we're holding, and start again from
 * scratch by constructing another spigot tree based on lgamma of a
 * much larger number (hence a more convergent correction series)
 * plus some log terms to range-reduce back to the thing we actually
 * wanted lgamma of in the first place. And when that in turn sets its
 * "I lied" flag, we throw it away and build another one even further
 * out, and so on.
 *
 * If you think this is a really horrible strategy, I wouldn't
 * disagree. If nothing else, it still converges rather slowly, and
 * quickly reaches the point where just handling all those
 * corrective logs is too much effort. But I challenge anyone
 * reading this to find a better one.
 *
 * TL;DR: of the four gamma strategies, two are not viable within
 * spigot's constraints, one works well for small inputs, and the
 * other (really disgusting and complicated) works well for large.
 */

/* ----------------------------------------------------------------------
 * Continued-fraction expansions of the lower and upper incomplete
 * gamma functions.
 *
 * Source for the actual expansions:
 * http://en.wikipedia.org/wiki/Incomplete_gamma_function#Evaluation_formulae
 * as of 2015-01-10.
 */

class LowerGammaBaseRational : public Source {
    /*
     * The lower incomplete gamma function is computed using the
     * following continued fraction, in which the even and odd terms
     * each have an obvious pattern and the two patterns interleave:
     *
     *                z^s exp(-z)
     *   gamma(s,z) = -------------------------------
     *                s - sz
     *                    ---------------------------
     *                    (s+1) + z
     *                            -------------------
     *                            (s+2) - (s+1)z
     *                                    -----------
     *                                    (s+3) + 2z
     *                                            ---
     *                                            ...
     *
     * This class computes the whole of the above formula, with z
     * fixed at 1, except that the numerator z^s exp(-z) is replaced
     * with 1; the caller must multiply in that factor afterwards.
     * (But since with z=1 it works out to just 1/e, that shouldn't be
     * hard.)
     */

    bigint sn, sd, i;
    int crState;

  public:
    LowerGammaBaseRational(const bigint &asn, const bigint &asd)
        : sn(asn), sd(asd)
    {
        crState = -1;
    }

    virtual LowerGammaBaseRational *clone() {
        return new LowerGammaBaseRational(sn, sd);
    }

    bool gen_interval(bigint *low, bigint *high)
    {
        *low = 0;
        *high = 0;
        return false;
    }

    bool gen_matrix(bigint *matrix)
    {
        crBegin;

        /*
         * An initial matrix representing x |-> 1/x.
         */
        matrix[0] = 0;
        matrix[1] = 1;
        matrix[2] = 1;
        matrix[3] = 0;
        crReturn(false);

        /*
         * Then the regular series.
         */
        i = 0;

        while (1) {
            matrix[0] = sn + 2*i*sd;
            matrix[1] = -(sn + i*sd);
            matrix[2] = sd;
            matrix[3] = 0;
            crReturn(false);

            matrix[0] = sn + (2*i+1) * sd;
            matrix[1] = (i+1) * sd;
            matrix[2] = sd;
            matrix[3] = 0;
            crReturn(false);

            ++i;
        }

        crEnd;
    }
};

class UpperGammaBaseRational : public Source {
    /*
     * The upper incomplete gamma function is computed using the
     * following continued fraction, in which the even and odd terms
     * each have an obvious pattern and the two patterns interleave:
     *
     *                z^s exp(-z)
     *   gamma(s,z) = -------------------
     *                z + 1-s
     *                    ---------------
     *                    1 + 1
     *                        -----------
     *                        z + 2-s
     *                            -------
     *                            1 + 2
     *                                ---
     *                                ...
     *
     * As above, this class computes this formula with z fixed at 1
     * and the top-level numerator (namely 1/e) missing.
     */

    bigint sn, sd, i;
    int crState;

  public:
    UpperGammaBaseRational(const bigint &asn, const bigint &asd)
        : sn(asn), sd(asd)
    {
        crState = -1;
    }

    virtual UpperGammaBaseRational *clone() {
        return new UpperGammaBaseRational(sn, sd);
    }

    bool gen_interval(bigint *low, bigint *high)
    {
        *low = 0;
        *high = 0;
        return false;
    }

    bool gen_matrix(bigint *matrix)
    {
        crBegin;

        /*
         * An initial matrix representing x |-> 1/x.
         */
        matrix[0] = 0;
        matrix[1] = 1;
        matrix[2] = 1;
        matrix[3] = 0;
        crReturn(false);

        /*
         * Then the regular series.
         */
        i = 1;

        while (1) {
            matrix[0] = sd;
            matrix[1] = i*sd - sn;
            matrix[2] = sd;
            matrix[3] = 0;
            crReturn(false);

            matrix[0] = 1;
            matrix[1] = i;
            matrix[2] = 1;
            matrix[3] = 0;
            crReturn(false);

            ++i;
        }

        crEnd;
    }
};

struct LowerBaseConstructor : MonotoneConstructor {
    MonotoneConstructor *clone() { return new LowerBaseConstructor(); }
    Spigot *construct(const bigint &n, const bigint &d) {
        return new LowerGammaBaseRational(n, d);
    }
};
struct UpperBaseConstructor : MonotoneConstructor {
    MonotoneConstructor *clone() { return new UpperBaseConstructor(); }
    Spigot *construct(const bigint &n, const bigint &d) {
        return new UpperGammaBaseRational(n, d);
    }
};

static Spigot *spigot_gamma_by_incomplete(Spigot *z)
{
    Spigot *lower = spigot_monotone(new LowerBaseConstructor, z->clone());
    Spigot *upper = spigot_monotone(new UpperBaseConstructor, z);
    return spigot_div(spigot_add(lower, upper), spigot_e());
}

static Spigot *spigot_lgamma_by_incomplete(Spigot *z)
{
    return spigot_log(spigot_abs(spigot_gamma_by_incomplete(z)));
}

/* ----------------------------------------------------------------------
 * Approach based on Stirling's approximation with periodic restarts.
 */

class GammaResidualBase : public Source {
    /*
     * Stirling's approximation for the gamma function involves some
     * logarithmic terms, which we can compute using functions
     * defined elsewhere in this program, plus a series sum which
     * converges to the error.
     *
     * This series, which I'm calling the 'gamma residual function',
     * converges to
     *
     *   log(gamma(z)) - [ (z-1/2) log(z) - z + log(2*pi)/2 ]
     *
     * and the series itself is
     *
     *   1/12       1/12         59/360               29/60
     *  ----- + ---------- + --------------- + -------------------- + ...
     *  (z+1)   (z+1)(z+2)   (z+1)(z+2)(z+3)   (z+1)(z+2)(z+3)(z+4)
     *
     * where those obscure fractions on the top are defined as an
     * expression involving Stirling numbers of the first kind.
     *
     * Stirling numbers of the first kind are defined to be the
     * coefficients you get if you multiply out
     * x(x-1)(x-2)...(x-n+1). Formally, we have s(n,k) such that
     *
     *   s(1,k) = 1 if k==1, 0 otherwise
     *   s(n+1,k) = s(n,k-1) - n s(n,k) for n>=1
     *
     * We're only interested in the absolute value, so we can turn
     * that minus sign into a plus. Now the coefficient for the nth
     * term of the above series is defined to be 1/2n times the sum,
     * over all k with s(n,k) nonzero, of k |s(n,k)| / (k+1)(k+2).
     *
     * Hideous, no? But it does converge.
     *
     * In order to convert this bizarrely defined series into a
     * spigot description, we need to be able to bound the
     * contribution from the remaining terms. To do this we observe
     * that the definition of the Stirling numbers tells us that the
     * sum of |s(n,k)| over all k for a given n is equal to n!; so
     * the corresponding series coefficient is bounded by the
     * smallest value of k/(k+1)(k+2) - which is given by k=1 and
     * hence is 1/6 - times n!.
     *
     * So we need to arrange that we multiply the interval size, in
     * each successive function in the spigot stream, by n
     * (reflecting the increasing size of the coefficient) and
     * divide by (z+n). This will clearly decrease it every time, so
     * we can reliably bound the amount which the series has left to
     * run.
     *
     * So our series is now rewritten (with the actual coefficient
     * values mercifully wrapped up as c_n, so that c_1=c_2=1/12,
     * c_3=59/360, c_4=29/60 and so on) as
     *
     *    1    c_1    2    c_2    3    c_3    4
     *   --- ( --- + --- ( --- + --- ( --- + --- ( ... ) ) ) )
     *   z+1    1!   z+2    2!   z+3    3!   z+4
     *
     * and from this format we can readily translate it into a
     * stream of spigot matrices. If c_n/n! is represented as the
     * rational a_n/b_n, and z itself is represented as p/q (I would
     * normally write 'n/d' for a rational here but I've already
     * used n in the above discussion), our matrices are
     *
     *  ( b_1.q a_1.q     ) ( 2.b_2.q 2.a_2.q      ) ( 3.b_3.q 3.a_3.q      )
     *  (   0   b_1.(p+q) ) (     0     b_2.(p+2q) ) (   0     3.b_3.(p+3q) )
     *
     * and so on.
     */

    struct listnode {
        bigint s;
        struct listnode *next;
        listnode() { next = NULL; }
        listnode(bigint as) : s(as) { next = NULL; }
    };

    bigint p, q, pnq;
    bigint nfact, nplus2fact;
    unsigned n, k;
    bigint coeff_num, coeff_denom;
    listnode *shead, *node;
    int crState;
    bool I_lied;

  public:
    GammaResidualBase(const bigint &ap, const bigint &aq)
        : p(ap), q(aq), pnq(ap), nfact(1), nplus2fact(6), n(1), shead(NULL)
    {
        crState = -1;
        shead = new listnode(1);
        I_lied = false;
    }

    virtual ~GammaResidualBase()
    {
        while (shead) {
            listnode *tmp = shead->next;
            delete shead;
            shead = tmp;
        }
    }

    virtual GammaResidualBase *clone() { return new GammaResidualBase(p, q); }

    bool gen_interval(bigint *low, bigint *high)
    {
        *low = 0;
        *high = 2;
        return true;
    }

    bool gen_matrix(bigint *matrix)
    {
        crBegin;

        while (1) {
            /*
             * Compute the coefficient c_n/n!.
             *
             * Our Stirling numbers are stored backwards, so that
             * shead represents s(n,n), shead->next is s(n,n-1) and
             * the final element in the list is s(n,1). This is just
             * because it happens to be more convenient to update
             * the list that way round.
             */
            k = n;
            coeff_num = 0;
            for (node = shead; node; node = node->next, k--)
                coeff_num += nplus2fact / (k+1) / (k+2) * k * node->s;
            coeff_denom = nfact * nplus2fact * (2*n);
            {
                bigint a = coeff_num, b = coeff_denom;
                while (1) {
                    if (b == 0) {
                        coeff_num /= a;
                        coeff_denom /= a;
                        break;
                    }
                    a %= b;
                    if (a == 0) {
                        coeff_num /= b;
                        coeff_denom /= b;
                        break;
                    }
                    b %= a;
                }
            }
//printf("[%d: ", n);
//bigint_print(coeff_num);printf("/");bigint_print(coeff_denom);printf(" ");

            /*
             * Return a matrix.
             */
            pnq += q;
            matrix[0] = coeff_denom * q * n;
            matrix[1] = coeff_num * q * n;
            matrix[2] = 0;
            matrix[3] = coeff_denom * pnq;

#if 1
            if (4*matrix[0] > 3*matrix[3]) {
                /*
                 * Give up and return a constant matrix instead.
                 */
                I_lied = true;
                matrix[0] = matrix[2] = 0;
                matrix[1] = matrix[3] = 1;
                crReturn(false);
            }
#else
            printf("<");
            bigint_print(100*matrix[0]/matrix[3]);
            printf("/100>");
#endif
//bigint_print(matrix[0]);printf(" ");
//bigint_print(matrix[1]);printf(" ");
//bigint_print(matrix[2]);printf(" ");
//bigint_print(matrix[3]);printf("]\n");
            crReturn(false);

            /*
             * Transform our linked list of Stirling numbers into
             * the next row.
             */
            node = new listnode(1);
            node->next = shead;
            shead = node;
            while ((node = node->next) != NULL) {
                node->s *= n;
                if (node->next)
                    node->s += node->next->s;
            }

            ++n;
            nfact *= n;
            nplus2fact *= (n+2);
        }

        crEnd;
    }

    friend class GammaWrapper;         /* so it can get at the I_lied flag */
};

class GammaWrapper : public BinaryIntervalSource {
    bigint p, q;

    bigint n, nextn, currp, constn, constd;

    GammaResidualBase *b;
    BracketingGenerator *bg;
    bool sret;
    unsigned outbits;

    int crState;

  public:
    GammaWrapper(const bigint &ap, const bigint &aq)
        : p(ap), q(aq), constn(1), constd(1), b(NULL), bg(NULL)
    {
        crState = -1;
        n = p / q;
        dprint("hello GammaWrapper %b/%b", &p, &q);
    }

    virtual ~GammaWrapper()
    {
        if (bg)
            delete bg;
    }

    virtual GammaWrapper *clone()
    {
        return new GammaWrapper(p, q);
    }

    virtual void gen_bin_interval(bigint *ret_lo, bigint *ret_hi,
                                  unsigned *ret_bits)
    {
        crBegin;

        if (n < 80)
            nextn = 80;
        else
            nextn = n;
        outbits = 0;
        currp = p;

        while (1) {
            while (n < nextn) {
                constn *= q;
                constd *= currp;
                currp += q;
                ++n;
            }
            /*
             * Instantiate the gamma residual function for currp/q.
             */
            dprint("initialising for %b/%b", &currp, &q);
            b = new GammaResidualBase(currp, q);

            /*
             * Now construct lgamma(z=p/q) out of that.
             */
            {
                Spigot *z = spigot_rational(currp, q);

                /* (z-1/2) log z */
                Spigot *t1 =
                    spigot_mul(spigot_sub(z->clone(), spigot_rational(1,2)),
                               spigot_log(spigot_abs(z->clone())));
                /* minus z */
                Spigot *t12 = spigot_sub(t1, z);
                /* plus log(2*pi)/2 */
                Spigot *t3 = spigot_rational_mul
                    (spigot_log(spigot_rational_mul(spigot_pi(),
                                                    2, 1)), 1, 2);
                /* plus log of all the stuff between p/q and currp/q */
                Spigot *t34 = t3;
                if (constn != 1 || constd != 1) {
                    Spigot *t4 = spigot_log
                        (spigot_abs(spigot_rational(constn, constd)));
                    t34 = spigot_add(t34, t4);
                };
                /* plus the residual function. */
                Spigot *t1234 = spigot_add(t12, t34);
                Spigot *s = spigot_add(t1234, b);
                bg = new BracketingGenerator(s);
                bg->set_denominator_lower_bound_shift(outbits);
            }

            while (1) {
                bg->get_bracket_shift(ret_lo, ret_hi, ret_bits);
                dprint("got bracket (%b,%b) / 2^%d",
                       ret_lo, ret_hi, (int)*ret_bits);
                if (b->I_lied) {
                    dprint("but it was a lie!");
                    break;
                }
                outbits = *ret_bits;
                crReturnV;
            }

            delete bg;
            bg = NULL;

            nextn = nextn * 3U / 2U;
        }

        crEnd;
    }
};

struct GammaConstructor : MonotoneConstructor {
    MonotoneConstructor *clone() { return new GammaConstructor(); }
    Spigot *construct(const bigint &n, const bigint &d) {
        return new GammaWrapper(n, d);
    }
};

static Spigot *spigot_lgamma_by_stirling(Spigot *z)
{
    return spigot_monotone(new GammaConstructor, z);
}

static Spigot *spigot_gamma_by_stirling(Spigot *z)
{
    bigint n, d;

    /*
     * lgamma computes log(|gamma|), so taking exp of that will always
     * be positive. So check whether this output ought to be negative.
     * gamma(x) < 0 if floor(x) is a negative odd integer.
     */
    bigint i;
    {
        CfracGenerator cfg(z->clone());
        cfg.get_term(&i);
    }
    bool negate = (i < 0 && (i % (unsigned)2) != 0);

    Spigot *absret = spigot_exp(spigot_lgamma_by_stirling(z));

    if (negate)
        return spigot_neg(absret);
    else
        return absret;
}

/* ----------------------------------------------------------------------
 * Top-level wrappers which handle special cases and then decide which
 * of the other two strategies to hand off to.
 *
 * In a comparative speed test (evaluating gamma to 40 sig figs for
 * lots of different inputs using both methods) I found that the
 * cross-over point was about |z|=35, for both positive and negative
 * inputs.
 */

static Spigot *spigot_gamma_special_cases(Spigot *z)
{
    bigint n, d;

    if (z->is_rational(&n, &d) && d == 1) {
        if (n <= 0) {
            throw spigot_error("gamma of non-positive integer");
        }

        /*
         * Special case for gamma of a positive integer: just
         * compute its factorial in the obvious way.
         */
        bigint ret = 1, k = 1;
        while (k < n) {
            ret *= k;
            ++k;
        }
        return spigot_integer(ret);
    }

    /* No special case was possible */
    return NULL;
}

Spigot *spigot_gamma(Spigot *z)
{
    Spigot *ret;

    if ((ret = spigot_gamma_special_cases(z)) != NULL)
        return ret;

    bool small;
    {
        StaticGenerator test(z->clone());
        bigint approx = test.get_approximate_approximant(2);
        small = (bigint_abs(approx) < 2*35);
    }

    if (small)
        return spigot_gamma_by_incomplete(z);
    else
        return spigot_gamma_by_stirling(z);
}

Spigot *spigot_lgamma(Spigot *z)
{
    Spigot *ret;

    if ((ret = spigot_gamma_special_cases(z)) != NULL)
        return spigot_log(ret);

    bool small;
    {
        StaticGenerator test(z->clone());
        bigint approx = test.get_approximate_approximant(2);
        small = (bigint_abs(approx) < 2*35);
    }

    if (small)
        return spigot_lgamma_by_incomplete(z);
    else
        return spigot_lgamma_by_stirling(z);
}
