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

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

#include <vector>

using std::vector;

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

/*
 * 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. []
 */

class SqrtInteger : public CfracSource {
    struct mda {
        bigint m, d, a;
        mda(const bigint &m_, const bigint &d_, const bigint &a_);
    };

    bigint n, m, d, a, a0;
    vector<mda> list;
    size_t i, j;
    int crState { -1 };

    virtual Spigot clone() override;
    virtual bool gen_term(bigint *term) override;

  public:
    SqrtInteger(const bigint &radicand);
};

SqrtInteger::mda::mda(const bigint &m_, const bigint &d_, const bigint &a_)
    : m(m_), d(d_), a(a_)
{
}

MAKE_CLASS_GREETER(SqrtInteger);

SqrtInteger::SqrtInteger(const bigint &radicand)
{
    n = radicand;
    dgreet(n);
}

Spigot SqrtInteger::clone()
{
    return spigot_clone(this, n);
}

bool SqrtInteger::gen_term(bigint *term)
{
    crBegin;

    if (n < 0) {
        throw domain_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);

        list.push_back(mda(m, d, a));

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

        for (i = 0; i < list.size(); i++) {
            if (m == list[i].m && d == list[i].d && a == list[i].a) {
                /*
                 * 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 < list.size(); j++) {
                        *term = list[j].a;
                        crReturn(true);
                    }
                }
            }
        }
    }
    crEnd;
}

/*
 * For the square root of an arbitrary _real_, 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.)
 */
class SqrtReal : public Source {
    Spigot x_orig;
    StaticGenerator sg;
    bigint r, e;
    int crState { -1 };

    // 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

    virtual Spigot clone() override;
    virtual bool gen_interval(bigint *low, bigint_or_inf *high) override;
    virtual bool gen_matrix(Matrix &matrix) override;

  public:
    SqrtReal(Spigot ax);
};

MAKE_CLASS_GREETER(SqrtReal);

SqrtReal::SqrtReal(Spigot ax)
    : x_orig(ax->clone()), sg(move(ax))
{
    dgreet();
}

Spigot SqrtReal::clone()
{
    return spigot_clone(this, x_orig->clone());
}

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

bool SqrtReal::gen_matrix(Matrix &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, nullptr, nullptr, 0, nullptr, true);
        if (hi < 0)
            throw domain_error("square root of negative number");

        dprint("initial integer part range (", 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=", bits);

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

        // And emit an initial matrix which maps the starting
        // interval [0,2] into the range [0, 2^(bits/2)].
        matrix = { 1_bi << (bits/2 - 1), 0, 0, 1 };
        dprint("returning initial scale matrix ", 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, nullptr, &hi_open,
                                 2, nullptr, false);
            dprint("integer part range (", 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 domain_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=", r, " e=", e, " d=", d_scale2, "/2"
#ifdef SQRT_STORE_K
                   " 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].
            {
                Matrix outmatrix = { 4, -4*lo, 0, 1 };
                dprint("premultiplying digit-extract matrix ", 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 = { 1, 0, 0, 2 };
                e = 2 * e_plus_d_scale2;
                r = r*2;
                dprint("output digit 0: ", 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 = { 1, 1, 0, 2 };
                e = 2 * e_plus_d_scale2 - 4*r - 1;
                r = r*2 + 1;
                dprint("output digit 1/2: ", 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 = { 1, 2, 0, 2 };
                e = 2 * e_plus_d_scale2 - 8*r - 4;
                r = r*2 + 2;
                dprint("output digit 1: ", 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;
}

/*
 * 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.
 */

class CbrtReal : public Source {
    Spigot x_orig;
    StaticGenerator sg;
    bigint r, r2, e;
    int crState { -1 };

    // 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

    virtual Spigot clone() override;
    virtual bool gen_interval(bigint *low, bigint_or_inf *high) override;
    virtual bool gen_matrix(Matrix &matrix) override;

  public:
    CbrtReal(Spigot ax);
};

MAKE_CLASS_GREETER(CbrtReal);

CbrtReal::CbrtReal(Spigot ax)
    : x_orig(ax->clone()), sg(move(ax))
{
    dgreet();
}

Spigot CbrtReal::clone()
{
    return spigot_clone(this, x_orig->clone());
}

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

bool CbrtReal::gen_matrix(Matrix &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, nullptr, nullptr, 0, nullptr, true);
        dprint("initial integer part range (", 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=", bits);

        // Scale the input down into the range [-8,8].
        {
            Matrix outmatrix = { 1, 0, 0, 1_bi << (bits - 3) };
            dprint("premultiplying initial scale matrix ", 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 = { 1_bi << (bits/3 - 1), 0, 0, 1 };
        dprint("returning initial scale matrix ", 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, nullptr, nullptr,
                                 2, nullptr, false);
            dprint("integer part range (", 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=", r, " e=", e, " d=", d
#ifdef CBRT_STORE_K
                   , " 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].
            {
                Matrix outmatrix = { 8, -8*d, 0, 1 };
                dprint("premultiplying digit-extract matrix ", 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 = { 1, -2, 0, 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: ", matrix);
            } else if (e_plus_d > 0) {
                // Output digit +1.
                matrix = { 1, 2, 0, 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: ", matrix);
            } else {
                // Output digit 0.
                matrix = { 1, 0, 0, 2 };
                e = 8 * e_plus_d;
                r2 = 4 * r2;
                r = 2 * r;
                dprint("output digit 0: ", 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)) {
        if (ad == 1) {
            /*
             * Square root of integer.
             */
            return make_unique<SqrtInteger>(an);
        } else {
            bigint ans, ads;

            /*
             * Special cases involving perfect squares.
             */
            if (an < 0)
                throw domain_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(make_unique<SqrtInteger>(ad)), ans, 1);
            else if (dexact)
                return spigot_rational_mul(
                    make_unique<SqrtInteger>(an), 1, ads);

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

static const UnaryFnWrapper expr_sqrt{"sqrt", spigot_sqrt};

static Spigot spigot_hypot(argvector<Spigot> &args)
{
    args.syntax.enforce_trivial("hypot");
    if (args.size() == 0)
        throw expr_error("expected at least one argument to 'hypot'");
    else if (args.size() == 1)
        return spigot_abs(move(args[0]));

    Spigot sumofsquares = spigot_square(move(args[0]));
    for (size_t i = 1; i < args.size(); i++)
        sumofsquares = spigot_add(move(sumofsquares),
                                  spigot_square(move(args[i])));
    return spigot_sqrt(move(sumofsquares));
}
    
static const VariadicFnWrapper expr_hypot{"hypot", spigot_hypot};

static 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 make_unique<CbrtReal>(move(a));
}

static const UnaryFnWrapper expr_cbrt{"cbrt", spigot_cbrt};
