/*
 * sqrt.cpp: Square roots.
 */

#include <assert.h>
#include <stdlib.h>

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

class SqrtInteger : public CfracSource {
    /*
     * To compute the square root of an arbitrary non-square
     * integer, we use the following magic algorithm to find its
     * continued fraction. The continued fraction is periodic, so
     * this costs us a bounded amount of work, after which we can
     * cheerfully return the same sequence of matrices over and over
     * again.
     *
     * The algorithm, taken from Wikipedia
     * (http://en.wikipedia.org/wiki/Methods_of_computing_square_roots
     * as of 2007-04-28), is as follows:
     *
     *  - Let m <- 0, d <- 1, and a <- floor(sqrt(n)).
     *  - Repeatedly:
     *     + output a as a continued fraction term.
     *     + let m <- d*a-m.
     *     + let d <- (n-m^2)/d. (This division always yields an
     *       exact integer; see below.)
     *     + let a <- floor((sqrt(n)+m)/d). (We can safely replace
     *       sqrt(n) with floor(sqrt(n)) here without affecting the
     *       result.)
     *     + check if (m,d,a) repeats a set of values which it
     *       previously had at this point. If so, we need not
     *       continue calculating.
     *
     * Proof by induction that every value of d is an integer:
     *  - Base case: d_0 is an integer, because it's defined to be
     *    1.
     *  - Suppose d_i and all previous d_j are integers.
     *     + d_i was computed as (n-m_i^2)/d_{i-1}, and hence if
     *       it's an integer then d_{i-1} is a factor of (n-m_i^2),
     *       and so d_i is _also_ a factor of that.
     *     + Now m_{i+1} = d_i a_i - m_i. So n - m_{i+1}^2 is equal to
     *       n - (m_i - d_i a_i)^2 = n - m_i^2 + 2 d_i a_i m_i - d_i^2 a_i^2.
     *     + And d_i divides n - m_i^2 (by construction), and
     *       divides any multiple of d_i (obviously), so it divides
     *       that. Hence d_{i+1} will be an integer too. []
     */

    bigint n, m, d, a, a0;
    bigint **list;
    int listlen, listsize;
    int i, j;
    int crState;

  public:
    SqrtInteger(const bigint &radicand)
    {
        n = radicand;
        crState = -1;
        listsize = 16;
        listlen = 0;
        list = (bigint **)malloc(listsize * 3 * sizeof(bigint *));
    }

    ~SqrtInteger()
    {
        for (int i = 0; i < listlen; i++) {
            delete list[3*i+0];
            delete list[3*i+1];
            delete list[3*i+2];
        }
        free(list);
    }

    virtual SqrtInteger *clone() { return new SqrtInteger(n); }

    virtual bool gen_term(bigint *term)
    {
        crBegin;

        if (n < 0) {
            throw spigot_error("square root of negative number");
        }
        a = a0 = bigint_sqrt(n);
        if (a * a == n) {
            /*
             * Special case of an exactly square number.
             */
            *term = a;
            crReturn(true);
            crReturn(false);
            assert(!"Should never get here");
        }
        m = 0;
        d = 1;

        while (1) {
            *term = a;
            crReturn(true);

            if (listlen >= listsize) {
                listsize = listlen * 5 / 4 + 16;
                list = (bigint **)realloc(list, listsize * 3 * sizeof(bigint *));
            }
            list[listlen*3+0] = new bigint(m);
            list[listlen*3+1] = new bigint(d);
            list[listlen*3+2] = new bigint(a);
            listlen++;

            m = d * a - m;
            d = (n - m*m) / d;
            a = (a0 + m) / d;

            for (i = 0; i < listlen; i++) {
                if (m == *list[i*3+0] &&
                    d == *list[i*3+1] &&
                    a == *list[i*3+2]) {
                    /*
                     * The current m,d,a repeats at position i.
                     * Repeat from there until the end of the list
                     * for ever.
                     */
                    while (1) {
                        for (j = i; j < listlen; j++) {
                            *term = *list[j*3+2];
                            crReturn(true);
                        }
                    }
                }
            }
        }
        crEnd;
    }
};

class SqrtReal : public Source {
    /*
     * We use a variant of the usual iterative binary square root
     * algorithm, adjusted for this environment by removing exactness
     * hazards.
     *
     * Suppose, at a given point in the algorithm, we know our input
     * to be in the interval [k,k+4], and we've already given a
     * promise that our output is in the interval [r,r+2]. Hence, we
     * must have arranged that r^2 <= k and that (r+2)^2 >= k+4.
     *
     * Now we get more data, which tells us our input is in the
     * interval [k+d,k+d+1] for some input 'digit' d, which is at
     * least 0 and at most 3. We want to correspondingly narrow our
     * output to an interval [r+x,r+x+1], for some x which is at least
     * 0 and at most 1 and which preserves the above invariants. That
     * is, we want
     *
     *   (r+x  )^2 <= (k+d  )  <=>  r^2 + 2xr + x^2               <= k + d
     *   (r+x+1)^2 >= (k+d+1)  <=>  r^2 + 2xr + x^2 + 2r + 2x + 1 >= k + d + 1
     *
     * or, simplifying a bit,
     *
     *   r^2 + 2xr + x^2 <= k + d <= r^2 + 2xr + x^2 + 2r + 2x
     *
     * We usually optimise slightly by not keeping a copy of the
     * largest number involved, namely k; instead we store the error
     * term e = k - r^2. So we start each step with the invariant that
     *
     *   0 <= e <= 4r
     *
     * and we seek to refine this by finding x such that
     *
     *   2xr + x^2 <= e + d <= 2xr + x^2 + 2r + 2x
     *
     * In the standard binary square root algorithm, the input digit d
     * would be one of {0,1,2,3} and the output bit x would be either
     * 0 or 1, so that the above inequality would become one of:
     *
     *        0 <= e + d <= 2r           (if we choose x = 0)
     *   2r + 1 <= e + d <= 4r + 3       (if we choose x = 1)
     *
     * and it's clear that given the input invariants and given that
     * r, e and d are all integers, exactly one of those possibilities
     * must hold, so we always have a unique choice of x.
     *
     * But in _this_ scenario, we have to avoid exactness hazards on
     * input, which means we have to permit some extra values of d.
     * Namely, we'll add in the half-integers between 0 and 3, so that
     * d can be any of {0,0.5,1,1.5,2,2.5,3}. This raises the
     * possibility that e+d might manage to fall _between_ 2r and
     * 2r+1, in which case we wouldn't be able to choose an integer
     * output digit. Hence, we must permit the output digit 1/2 as
     * well, which gives us one more possible inequality on output:
     *
     *   r + 1/4 <= e + d <= 3r + 3/4    (if we choose x = 1/2)
     *
     * Finally, after each step, we adjust r to its new value r+x, and
     * e to its new value e + d - 2xr - x^2 (in which 'r' still means
     * the old value). Then we scale up r by 2, and e by 4. (Because r
     * is in the same units as the output, and e the input.)
     */
    Spigot *x_orig;
    StaticGenerator sg;
    bigint r, e;
    int crState;

    // For debugging the algorithm, we can store k itself if
    // necessary, and check that e, r, k remain in the right
    // relationship throughout. Uncomment this #define to do so.

//#define SQRT_STORE_K
#ifdef SQRT_STORE_K
    bigint k;
#endif

  public:
    SqrtReal(Spigot *ax)
        : x_orig(ax->clone()), sg(ax)
    {
        crState = -1;
        dprint("hello SqrtReal %p", ax);
    }

    ~SqrtReal() {
        delete x_orig;
    }

    virtual SqrtReal *clone() { return new SqrtReal(x_orig->clone()); }

    virtual bool gen_interval(bigint *low, bigint *high) {
        *low = 0;
        *high = 2;
        return true;               // force absorption of first matrix
    }

    virtual bool gen_matrix(bigint *matrix) {
        crBegin;

        /*
         * We kick off by finding out any old information we can about
         * the input number's integer part, deciding on a starting
         * scale factor, and mapping the number to something in the
         * interval [0,8]. (Conceptually [0,4], but we'll need the
         * next iteration of the spigot to give us output at
         * half-integer granularity.)
         *
         * We can throw an immediate error if the _high_ end of the
         * interval is negative, but if the low end of the interval is
         * negative, we ignore that fact and just pretend it's zero.
         * If the number does subsequently turn out to be negative,
         * we'll find that we've extracted a negative digit later in
         * the algorithm, and we can throw our error then. This
         * strategy permits us to take the square root of a non-
         * obviously zero number without any exactness hazard.
         */
        {
            // Get some info on the integer part.
            bigint lo, hi;
            sg.iterate_to_bounds(&lo, &hi, NULL, NULL, 0, NULL, true);
            if (hi < 0)
                throw spigot_error("square root of negative number");

            dprint("initial integer part range (%b,%b)", &lo, &hi);

            // Find a power of 4 at least as big as hi.
            unsigned bits = bigint_approxlog2(hi-1);
            bits -= (bits % 2);
            bits += 2;
            if (bits < 4)
                bits = 4;              // bound below at 16

            dprint("bits=%d", (int)bits);

            // Scale the input down into the range [0,8) (plus
            // possible negative numbers we haven't spotted yet).
            bigint outmatrix[4];
            outmatrix[0] = 1;
            outmatrix[1] = outmatrix[2] = 0;
            outmatrix[3] = (bigint)1 << (bits - 3);
            dprint("premultiplying initial scale matrix %4m", outmatrix);
            sg.premultiply(outmatrix);

            // And emit an initial matrix which maps the starting
            // interval [0,2] into the range [0, 2^(bits/2)].
            matrix[0] = (bigint)1 << (bits/2 - 1);
            matrix[1] = matrix[2] = 0;
            matrix[3] = 1;
            dprint("returning initial scale matrix %4m", matrix);
        }
        crReturn(false);

        /*
         * Now we can initialise the main algorithm.
         */
        r = 0;
        e = 0;
#ifdef SQRT_STORE_K
        k = 0;
#endif

        while (1) {
            {
                // Iterate the input spigot until we find a bounding
                // interval of width at most 2. We're starting off in
                // the interval range [0,8], so narrowing down to a
                // width-2 subinterval of that will reliably give us a
                // half-integer-valued digit.
                bigint lo, hi;
                bool hi_open;
                sg.iterate_to_bounds(&lo, &hi, NULL, &hi_open, 2, NULL, false);
                dprint("integer part range (%b,%b)", &lo, &hi);

                // If the number has turned out negative, bomb out. If
                // it still only _might_ be negative, then just take
                // our digit to be zero, and we'll throw the error
                // later on when (if) we have final proof.
                if (hi < 0 || (hi == 0 && hi_open))
                    throw spigot_error("square root of negative number");
                if (lo < 0)
                    lo = 0;
                if (lo > 6)            // clip digits at the top as well
                    lo = 6;

                // OK. 'lo' is our actual digit value. (Or rather,
                // twice it: it's stored as 1-bit fixed point, so that
                // it ranges from 0 to 6 for our possible digits 0,
                // 1/2, ..., 5/2, 3.)
                bigint &d_scale2 = lo;
                dprint("step: r=%b e=%b d=%b/2"
#ifdef SQRT_STORE_K
                       " k=%b"
#endif
                       , &r, &e, &d_scale2
#ifdef SQRT_STORE_K
                       , &k
#endif
                    );
                assert(0 <= e);
                assert(e <= 4*r);
#ifdef SQRT_STORE_K
                assert(e == k - r*r);
#endif

                // Premultiply in a matrix which zooms the new input
                // interval back up to [0,8].
                {
                    bigint outmatrix[4];
                    outmatrix[0] = 4;
                    outmatrix[1] = -4*lo;
                    outmatrix[2] = 0;
                    outmatrix[3] = 1;
                    dprint("premultiplying digit-extract matrix %4m",
                           outmatrix);
                    sg.premultiply(outmatrix);
                }

                // Now figure out what digit to output. Note that r
                // and e are as described above, but d has twice its
                // nominal value.
                bigint e_plus_d_scale2 = 2*e + d_scale2;
                if (e_plus_d_scale2 <= 4*r) { // e+d <= 2r
                    // Output digit 0.
                    matrix[0] = 1;
                    matrix[1] = 0;
                    matrix[2] = 0;
                    matrix[3] = 2;
                    e = 2 * e_plus_d_scale2;
                    r = r*2;
                    dprint("output digit 0: %4m", matrix);
                } else if (2*r < e_plus_d_scale2 && // r+1/4 <= e+d
                           e_plus_d_scale2 < 6*r+2) { // e+d <= 3r+3/4
                    // Output digit 1/2.
                    matrix[0] = 1;
                    matrix[1] = 1;
                    matrix[2] = 0;
                    matrix[3] = 2;
                    e = 2 * e_plus_d_scale2 - 4*r - 1;
                    r = r*2 + 1;
                    dprint("output digit 1/2: %4m", matrix);
                } else if (4*r+2 <= e_plus_d_scale2 && // 2r+1 <= e+d
                           e_plus_d_scale2 <= 8*r+6) { // e+d <= 4r+3
                    // Output digit 1.
                    matrix[0] = 1;
                    matrix[1] = 2;
                    matrix[2] = 0;
                    matrix[3] = 2;
                    e = 2 * e_plus_d_scale2 - 8*r - 4;
                    r = r*2 + 2;
                    dprint("output digit 1: %4m", matrix);
                } else {
                    assert(!"SqrtReal: no output digit worked!");
                }

#ifdef SQRT_STORE_K
                k = 4*k + 2*d_scale2;
#endif
            }

            crReturn(false);
        }

        crEnd;
    }
};

/*
 * Integer cube root function.
 */
bigint cuberoot(bigint n)
{
    int sign = 1;
    if (n < 0) {
        sign = -1;
        n = -n;
    }

    bigint T1 = 0;
    bigint T2 = 0;
    bigint x = 0;
    int as = bigint_approxlog2(n) / 3 + 2;
    bigint T0 = bigint_power(2, 3*as);
    
    while (as >= 0) {
        bigint T = T0 + T1 + T2;
        if (n >= T) {
            n = n - T;
            x = x + bigint_power(2, as);
            T2 = T2 + T1;
            T1 = T1 + 3*T0;
            T2 = T2 + T1;
        }
        T0 >>= 3;
        T1 >>= 2;
        T2 >>= 1;
        --as;
    }

    return x * sign;
}

class CbrtReal : public Source {
    /*
     * Cube root algorithm, basically like the square root one above
     * but with all the details fiddled. In particular, I think this
     * time it'll be convenient to put our base value k at the centre
     * of the interval, since cube roots can handle either positive or
     * negative numbers, and also because then most of the cases
     * become symmetric and hence easy to check against each other.
     * And since that will involve a lot of annoying 1/2s, I'll also
     * scale the whole thing up by a factor of two.
     *
     * So suppose, at a given point in the algorithm, we know our
     * input to be in the interval [k-8,k+8], and we've already given
     * a promise that our output is in the interval [r-2,r+2]. Hence,
     * we must have arranged that (r-2)^3 <= k-8 and that (r+2)^3 >=
     * k+8.
     *
     * Now we get more data, which tells us our input is in the
     * interval [k+d-1,k+d+1] for some input 'digit' d, which is at
     * least -7 and at most +7. We want to correspondingly narrow our
     * output to an interval [r+x-1,r+x+1], for some x which is at
     * least -1 and at most 1, and which preserves the above
     * invariants. That is, we want
     *
     *   (r+x-1)^3 <= (k+d-1)
     *   (r+x+1)^3 >= (k+d+1)
     *
     * which, similarly to the above, we rewrite as
     *
     *   (r+x-1)^3+1 <= k+d <= (r+x+1)^3-1
     *
     * and, as with the square root, we store e = k-r^3 in place of k,
     * so that becomes
     *
     *   (r+x-1)^3 - r^3 + 1 <= e + d <= (r+x+1)^3 - r^3 - 1      (*)
     *
     * So we start each step with the invariant that
     *
     *   -6r^2 + 12r <= e <= 6r^2 + 12r
     *
     * and we seek to refine this by finding an x which satisfies the
     * inequality marked (*) above.
     *
     * Again similarly to the square root above, we're permitted three
     * output digits: the two most extreme values -1 and +1, and 0 in
     * the middle. Expanding out (*) for each of those options gives
     * us the following inequalities:
     *
     *   -6r^2 + 12r - 7 <= e + d <= -1               (if we choose x = -1)
     *   -3r^2 +  3r     <= e + d <= 3r^2 +  3r       (if we choose x =  0)
     *                 1 <= e + d <= 6r^2 - 12r + 7   (if we choose x = +1)
     *
     * Finally, after each step, we adjust r to its new value r+x, and
     * e to its new value e + d - 3xr^2 - 3x^2r - x^3 (in which 'r'
     * still means the old value). Then we scale up r by 2, and e by
     * 8.
     */
    Spigot *x_orig;
    StaticGenerator sg;
    bigint r, r2, e;
    int crState;

    // For debugging the algorithm, we can store k itself if
    // necessary, and check that e, r, k remain in the right
    // relationship throughout. Uncomment this #define to do so.

//#define CBRT_STORE_K
#ifdef CBRT_STORE_K
    bigint k;
#endif

  public:
    CbrtReal(Spigot *ax)
        : x_orig(ax->clone()), sg(ax)
    {
        crState = -1;
        dprint("hello CbrtReal %p", ax);
    }

    ~CbrtReal() {
        delete x_orig;
    }

    virtual CbrtReal *clone() { return new CbrtReal(x_orig->clone()); }

    virtual bool gen_interval(bigint *low, bigint *high) {
        *low = -2;
        *high = 2;
        return true;               // force absorption of first matrix
    }

    virtual bool gen_matrix(bigint *matrix) {
        crBegin;

        /*
         * We kick off by finding out any old information we can about
         * the input number's integer part, deciding on a starting
         * scale factor, and mapping the number to something in the
         * interval [-8,8].
         */
        {
            // Get some info on the integer part.
            bigint lo, hi;
            sg.iterate_to_bounds(&lo, &hi, NULL, NULL, 0, NULL, true);
            dprint("initial integer part range (%b,%b)", &lo, &hi);

            // Find a power of 2 at least as big as the absolute value
            // of either limit.
            unsigned bits = 0;
            if (hi >= 0) {
                unsigned bits2 = bigint_approxlog2(hi - 1);
                if (bits < bits2)
                    bits = bits2;
            }
            if (lo <= 0) {
                unsigned bits2 = bigint_approxlog2(1 - lo);
                if (bits < bits2)
                    bits = bits2;
            }

            // Round up to the next power of 8, and to at least 8.
            bits -= (bits % 3);
            bits += 3;
            if (bits < 3)
                bits = 3;
            dprint("bits=%d", (int)bits);

            // Scale the input down into the range [-8,8].
            bigint outmatrix[4];
            outmatrix[0] = 1;
            outmatrix[1] = outmatrix[2] = 0;
            outmatrix[3] = (bigint)1 << (bits - 3);
            dprint("premultiplying initial scale matrix %4m", outmatrix);
            sg.premultiply(outmatrix);

            // And emit an initial matrix which maps the starting
            // interval [-2,2] into the range [-2^(bits/3), 2^(bits/3)].
            matrix[0] = (bigint)1 << (bits/3 - 1);
            matrix[1] = matrix[2] = 0;
            matrix[3] = 1;
            dprint("returning initial scale matrix %4m", matrix);
        }
        crReturn(false);

        /*
         * Now we can initialise the main algorithm. To avoid lots of
         * needless multiplication, we keep running track of r^2 as
         * well as r.
         */
        r = 0;
        r2 = 0;
        e = 0;
#ifdef CBRT_STORE_K
        k = 0;
#endif

        while (1) {
            {
                // Iterate the input spigot until the range of integer
                // parts has narrowed to at most two.
                bigint lo, hi;
                sg.iterate_to_bounds(&lo, &hi, NULL, NULL, 2, NULL, false);
                dprint("integer part range (%b,%b)", &lo, &hi);

                // Our input digit will be the midpoint of that
                // width-2 interval. (If it is width 2; if it's
                // already narrower than that, then either integer
                // will do. The important thing is that
                // [digit-1,digit+1] contains the number.)
                bigint d = (lo + hi) / 2U;

                // We know the input was in in the range [-8,+8], so
                // if the spigot gave us an interval jammed up against
                // the edge of that range so that d has ended up
                // outside the range {-7,...,7}, clip it into range
                // because we know better.
                if (d < -7)
                    d = -7;
                if (d > +7)
                    d = +7;

                dprint("step: r=%b e=%b d=%b"
#ifdef CBRT_STORE_K
                       " k=%b"
#endif
                       , &r, &e, &d
#ifdef CBRT_STORE_K
                       , &k
#endif
                    );
                assert(-6*r2 + 12*r <= e);
                assert(e <= 6*r2 + 12*r);
#ifdef CBRT_STORE_K
                assert(r2 == r*r);
                assert(e == k - r*r*r);
#endif

                // Premultiply in a matrix which zooms the new input
                // interval back up to [-8,8].
                {
                    bigint outmatrix[4];
                    outmatrix[0] = 8;
                    outmatrix[1] = -8*d;
                    outmatrix[2] = 0;
                    outmatrix[3] = 1;
                    dprint("premultiplying digit-extract matrix %4m",
                           outmatrix);
                    sg.premultiply(outmatrix);
                }

                // Now figure out what digit to output. We can almost
                // always output a -1 or 1 digit; only that
                // inconvenient zero in the middle needs to use the 0
                // digit.
                bigint e_plus_d = e + d;
                if (e_plus_d < 0) {
                    // Output digit -1.
                    matrix[0] = 1;
                    matrix[1] = -2;
                    matrix[2] = 0;
                    matrix[3] = 2;
                    e = 8 * (e_plus_d + 3*r2 - 3*r + 1);
                    r2 = 4 * (r2 - 2*r + 1);
                    r = 2 * (r-1);
                    dprint("output digit -1: %4m", matrix);
                } else if (e_plus_d > 0) {
                    // Output digit +1.
                    matrix[0] = 1;
                    matrix[1] = 2;
                    matrix[2] = 0;
                    matrix[3] = 2;
                    e = 8 * (e_plus_d - 3*r2 - 3*r - 1);
                    r2 = 4 * (r2 + 2*r + 1);
                    r = 2 * (r+1);
                    dprint("output digit +1: %4m", matrix);
                } else {
                    // Output digit 0.
                    matrix[0] = 1;
                    matrix[1] = 0;
                    matrix[2] = 0;
                    matrix[3] = 2;
                    e = 8 * e_plus_d;
                    r2 = 4 * r2;
                    r = 2 * r;
                    dprint("output digit 0: %4m", matrix);
                }

#ifdef CBRT_STORE_K
                k = 8 * (k + d);
#endif
            }

            crReturn(false);
        }

        crEnd;
    }
};

Spigot *spigot_sqrt(Spigot *a)
{
    bigint an, ad;

    if (a->is_rational(&an, &ad)) {
        delete a;
        if (ad == 1) {
            /*
             * Square root of integer.
             */
            return new SqrtInteger(an);
        } else {
            bigint ans, ads;

            /*
             * Special cases involving perfect squares.
             */
            if (an < 0)
                throw spigot_error("square root of negative number");
            ans = bigint_sqrt(an);
            ads = bigint_sqrt(ad);
            bool nexact = (ans*ans == an);
            bool dexact = (ads*ads == ad);
            if (nexact && dexact)
                return spigot_rational(ans, ads);
            else if (nexact)
                return spigot_rational_mul(spigot_reciprocal
                                           (new SqrtInteger(ad)), ans, 1);
            else if (dexact)
                return spigot_rational_mul(new SqrtInteger(an), 1, ads);

            /*
             * Square root of rational, computed by dividing two
             * square roots of integers.
             */
            return spigot_div(new SqrtInteger(an), new SqrtInteger(ad));
        }
    } else {
        return new SqrtReal(a);
    }
}

Spigot *spigot_hypot(Spigot *a, Spigot *b)
{
    Spigot *ac = a->clone();
    Spigot *bc = b->clone();
    return spigot_sqrt(spigot_add(spigot_mul(a, ac), spigot_mul(b, bc)));
}
    
Spigot *spigot_cbrt(Spigot *a)
{
    bigint an, ad;

    if (a->is_rational(&an, &ad)) {
        /*
         * Special case of the cube of a rational.
         */
        bigint anr = cuberoot(an);
        bigint adr = cuberoot(ad);
        bool nexact = (anr*anr*anr == an);
        bool dexact = (adr*adr*adr == ad);
        if (nexact && dexact)
            return spigot_rational(anr, adr);
    }

    return new CbrtReal(a);
}
