/*
 * Generate an arbitrary algebraic number.
 */

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

#include <vector>
#include <utility>
using std::vector;
using std::pair;

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

class Algebraic : public Source {
    /*
     * This class assumes that the root we're looking for is the
     * unique one in the interval (0,1).
     */
    vector<bigint> P;
    vector<bigint> P_orig;

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

  public:
    Algebraic(const vector<bigint> &aP);
};

MAKE_CLASS_GREETER(Algebraic);

Algebraic::Algebraic(const vector<bigint> &aP)
{
    P = aP;
    P_orig = aP;
    dgreet(P);
}

Spigot Algebraic::clone()
{
    return spigot_clone(this, P_orig);
}

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

bool Algebraic::gen_matrix(Matrix &matrix)
{
    /*
     * Scale up our polynomial P by replacing it with Q such that
     * Q(x) = P(x/2) (up to a constant). This means our root will
     * now be in the interval [0,2].
     */
    bigint factor = 1;
    for (int i = P.size(); i-- > 0 ;) {
        P[i] *= factor;
        factor *= 2;
    }

    /*
     * Now evaluate P at 1 (i.e. just sum the terms) to find out
     * whether the root lies in [0,1] or [1,2], i.e. which side of
     * 1/2 it was on before we scaled up.
     */
    bigint sum = 0;
    for (int i = 0; i < (int)P.size(); ++i)
        sum += P[i];
    int digit = (bigint_sign(sum) == bigint_sign(P[0])) ? 1 : 0;

    if (digit) {
        /*
         * If the root is in [1,2], we need to transform P again,
         * by finding Q such that Q(x) = P(x+1), which brings the
         * root back down into [0,1].
         */
        for (int j = P.size() - 1; j-- > 0 ;)
            for (int i = j; i+1 < (int)P.size(); ++i)
                P[i] += P[i+1];
    }

    /*
     * Return an appropriate binary-digit matrix.
     */
    matrix = { 1, digit, 0, 2 };
    return false;
}

typedef pair<bigint, bigint> qval;
typedef vector<qval> qpoly;

static int qval_sign(const qval &v) {
    return bigint_sign(v.first) * bigint_sign(v.second);
}

static qval qpoly_eval(qpoly p, qval x) {
    qval acc = { 0, 1 };
    for (size_t i = p.size(); i-- > 0 ;) {
        qval &c = p[i];
        acc = { acc.first * x.first, acc.second * x.second };
        acc = { acc.first * c.second + acc.second * c.first,
            acc.second * c.second };
    }
    return acc;
}

static void debug_print_item(FILE *fp, const qval &a)
{
    debug_print_item(fp, a.first);
    if (a.second != 1) {
        debug_print_item(fp, "/");
        debug_print_item(fp, a.second);
    }
}

template<typename T>
static void debug_print_item(FILE *fp, const vector<T> &a)
{
    debug_print_item(fp, "( ");
    for (size_t i = 0; i < a.size(); ++i) {
        if (i > 0)
            debug_print_item(fp, " + ");
        debug_print_item(fp, a[i]);
        if (i == 1) {
            debug_print_item(fp, " x");
        } else if (i > 1) {
            debug_print_item(fp, " x^");
            debug_print_item(fp, i);
        }
    }
    debug_print_item(fp, " )");
}

/*
 * This class holds the code that sets up a polynomial to pass to the
 * above Algebraic class. This has no ongoing state, so it could
 * perfectly well have been a bunch of standalone free functions,
 * except that then we couldn't assign it a debugging id and permit
 * the diagnostics to be enabled and disabled at the command line.
 */
class AlgebraicPrep : public Debuggable {
    bigint nhi, dhi, nlo, dlo, d;
    vector<bigint> P;
    vector<qpoly> sturm;

    static qval a_minus_b_times_c(const qval &a, const qval &b, const qval &c);
    qpoly divide_poly(qpoly &n, const qpoly &d);
    void scale_poly(qpoly &P);
    void reduce_poly(vector<bigint> &Porig);
    void sturm_calculate(vector<bigint> &Porig);
    bool sturm_eval(const bigint &x, unsigned &output);

  public:
    AlgebraicPrep(const bigint &nhi, const bigint &dhi,
                  const bigint &nlo, const bigint &dlo,
                  const vector<bigint> &P);
    Spigot spigot();
};

MAKE_CLASS_GREETER(AlgebraicPrep);

AlgebraicPrep::AlgebraicPrep(const bigint &anhi, const bigint &adhi,
                             const bigint &anlo, const bigint &adlo,
                             const vector<bigint> &aP)
    : nhi{anhi}, dhi{adhi}, nlo{anlo}, dlo{adlo}, P{aP}
{
    dgreet("[", nlo, "/", dlo, ",", nhi, "/", dhi, "] P=", P);
}

qval AlgebraicPrep::a_minus_b_times_c(
    const qval &a, const qval &b, const qval &c)
{
    bigint n = a.first*b.second*c.second - a.second*b.first*c.first;
    bigint d = a.second*b.second*c.second;
    bigint g = gcd(n, d);
    return qval(n / g, d / g);
}

/*
 * Divide two polynomials over Q. Returns the quotient; modifies n to
 * be the remainder.
 */
qpoly AlgebraicPrep::divide_poly(qpoly &n, const qpoly &d)
{
    assert(d.size() > 0);
    assert(d[d.size()-1].first != 0);

    dprint("dividing qpoly ", n, " / ", d);

    qpoly ret;

    while (n.size() >= d.size()) {
        size_t shift = n.size() - d.size();
        qval &nfirst = n[n.size() - 1];
        const qval &dfirst = d[d.size() - 1];
        qval mult(nfirst.first * dfirst.second, nfirst.second * dfirst.first);
        for (size_t i = 0; i < d.size(); ++i) {
            n[i+shift] = a_minus_b_times_c(n[i+shift], d[i], mult);
        }
        while (n.size() > 0 && n[n.size() - 1].first == 0)
            n.resize(n.size() - 1);

        if (mult.first != 0) {
            if (ret.size() < shift+1)
                ret.resize(shift+1, qval(0, 1));
            ret[shift] = mult;
        }
        dprint("divide: now quot = ", ret, " rem = ", n);
    }
    dprint("divide done");

    return ret;
}

/*
 * Scale a polynomial by a constant to make the coefficients all
 * integers and as small as possible.
 */
void AlgebraicPrep::scale_poly(qpoly &P)
{
    bigint lcm_of_denoms = 1;
    for (size_t i = 0; i < P.size(); ++i) {
        lcm_of_denoms = lcm_of_denoms * P[i].second /
            gcd(lcm_of_denoms, P[i].second);
    }
    for (size_t i = 0; i < P.size(); ++i) {
        P[i].first = P[i].first * lcm_of_denoms / P[i].second;
        P[i].second = 1;
    }
    bigint gcd_of_nums = P[0].first;
    for (size_t i = 1; i < P.size(); ++i) {
        gcd_of_nums = gcd(gcd_of_nums, P[i].first);
    }
    for (size_t i = 0; i < P.size(); ++i) {
        P[i].first /= gcd_of_nums;
    }
}

/*
 * Reduce a polynomial to one with no repeated roots, by dividing off
 * the gcd of P and its derivative.
 */
void AlgebraicPrep::reduce_poly(vector<bigint> &Porig)
{
    qpoly P, Pprime;

    // Copy P into a qpoly.
    for (size_t i = 0; i < Porig.size(); ++i) {
        P.push_back(qval(Porig[i], 1));
    }

    // And differentiate P.
    for (size_t i = 1; i < Porig.size(); ++i) {
        Pprime.push_back(qval((int)i * Porig[i], 1));
    }

    dprint("reduce_poly: ", P);
    dprint("derivative = ", Pprime);

    // Now do Euclid's algorithm between P and Pprime, i.e. repeatedly
    // reduce one of them mod the other.
    qpoly *a = &P, *b = &Pprime;
    while (b->size() > 0) {
        dprint("Euclid step: ", *a, " / ", *b);
        divide_poly(*a, *b);
        scale_poly(*a);
        qpoly *tmp = a; a = b; b = tmp;
    }

    dprint("gcd = ", *a);

    // Now *a is the gcd. Make another copy of the original P, and
    // divide that off.
    qpoly Pcopy;
    for (size_t i = 0; i < Porig.size(); ++i) {
        Pcopy.push_back(qval(Porig[i], 1));
    }
    qpoly ret = divide_poly(Pcopy, *a);
    dprint("reduced poly before scaling = ", ret);
    scale_poly(ret);

    // Done. Copy into Porig for return.
    Porig.resize(0);
    for (size_t i = 0; i < ret.size(); ++i) {
        Porig.push_back(ret[i].first);
    }
    dprint("final reduced poly = ", Porig);
}

/*
 * Calculate the Sturm sequence for a (square-free) polynomial, which
 * can be used to count its roots in an interval.
 *
 * Source: https://en.wikipedia.org/wiki/Sturm%27s_theorem . The idea
 * is to form a sequence of polynomials
 *
 *   P_0 = P
 *   P_1 = P'
 *   P_{n+2} = -P_n mod P_{n+1}
 *
 * terminating when one becomes zero. Then define an integer-valued
 * function V(x) by evaluating all those polynomials at x and counting
 * the number of times the sign changes. A zero, or run of zeroes, in
 * the middle of the sequence, is ignored - the question is whether
 * the nonzero values before and after the run have opposite sign.
 *
 * Sturm's theorem says that if P has no repeated roots (which we've
 * already ensured), then V(x) goes down by one exactly when x crosses
 * a real root of P. Hence, the number of roots between a and b is
 * exactly V(a)-V(b).
 *
 * (What if there's a root _at_ a or b? Sturm's theorem does give an
 * answer to that, but we don't care, because in our application, we
 * call that an input validation error. Since the first polynomial in
 * the Sturm sequence is P itself, we can detect this situation
 * easily, and return failure immediately without bothering to
 * evaluate the rest of the sequence.)
 */
void AlgebraicPrep::sturm_calculate(vector<bigint> &Porig)
{
    // Initialise the sequence with P and its derivative.
    {
        qpoly P, Pprime;

        for (size_t i = 0; i < Porig.size(); ++i) {
            P.push_back(qval(Porig[i], 1));
        }
        dprint("sturm[0]: ", P);

        for (size_t i = 1; i < Porig.size(); ++i) {
            Pprime.push_back(qval((int)i * Porig[i], 1));
        }
        dprint("sturm[1]: ", Pprime);

        sturm.push_back(std::move(P));
        sturm.push_back(std::move(Pprime));
    }

    // Now repeatedly reduce the second-last poly mod the last.
    while (true) {
        qpoly next = sturm[sturm.size() - 2];
        divide_poly(next, sturm[sturm.size() - 1]);
        while (next.size() > 0 && next[next.size() - 1].first == 0)
            next.pop_back();
        if (next.size() == 0)
            break;
        // The Sturm sequence recurrence uses the _negative_ residue.
        for (auto &c: next)
            c.first = -c.first;
        dprint("sturm[", sturm.size(), "]: ", next);
        sturm.push_back(std::move(next));
    }
}

bool AlgebraicPrep::sturm_eval(const bigint &x, unsigned &output)
{
    qval qx { x, 1_bi };

    // Find the sign of P_0(x). If it's zero, that means x is a root
    // of P, so return failure (the user should have given better
    // bounds).
    int prev_sign;
    {
        qval v = qpoly_eval(sturm[0], qx);
        dprint("evaluating sturm[0] at ", qx, " -> ", v);
        prev_sign = qval_sign(v);
    }
    if (prev_sign == 0)
        return false;

    output = 0;

    // Evaluate all the other P_i(x) and count sign changes, skipping
    // any zeroes.
    for (size_t i = 1, e = sturm.size(); i < e; i++) {
        qval v = qpoly_eval(sturm[i], qx);
        dprint("evaluating sturm[", i, "] at ", qx, " -> ", v);
        int sign = qval_sign(v);
        if (sign != 0) {
            if (sign != prev_sign)
                output++;
            prev_sign = sign;
        }
    }

    return true;
}

Spigot AlgebraicPrep::spigot()
{
    if (dlo != dhi) {
        // Reduce dlo and dhi to a single common d.
        {
            bigint a = bigint_abs(dlo), b = bigint_abs(dhi);
            while (b != 0) {
                bigint t = b;
                b = a % b;
                a = t;
            }
            d = dlo * dhi / a;
            nlo *= dhi / a;
            nhi *= dlo / a;
        }
        dprint("common-denominator interval = [",
               nlo, "/", d, ",", nhi, "/", d, "]");
    } else {
        d = dlo;
    }

    // Reduce the polynomial to one without repeated roots.
    reduce_poly(P);

    dprint("reduced: P = ", P);

    int n = P.size();
    assert(n > 0);
    assert(P[n-1] != 0);

    /*
     * Scale P so that the root is in an integer interval (nlo,nhi)
     * rather than a rational interval (nlo/d,nhi/d). That is, we
     * transform the polynomial P(x) into P(x/d), and scale the
     * coefficients back up to be integers. We leave d containing the
     * original denominator: we'll transform the output of the
     * Algebraic class by a factor of 1/d to compensate for this
     * transformation.
     */
    {
        bigint factor = 1;
        for (int i = P.size(); i-- > 0 ;) {
            P[i] *= factor;
            factor *= d;
        }
    }

    dprint("scaled #1: P = ", P);

    /*
     * Scale up again similarly, this time to make P monic. That way,
     * if the root we're looking for is rational, it will be an
     * integer, so the initial binary search of the integer interval
     * (nlo, nhi) will find it, and we can return it in an exact form.
     *
     * Here, we're transforming P(x) into P(x / L) where L is the
     * leading coefficient. But this time we _do_ adjust d for the
     * output scaling, because this change to the denominator wasn't
     * accounted for by the user - it's one we came up with ourselves.
     */
    d *= P[n-1];
    nlo *= P[n-1];
    nhi *= P[n-1];
    {
        bigint factor = 1;
        for (int i = n-1; i-- > 0 ;) {
            P[i] *= factor;
            factor *= P[n-1];
        }
    }
    P[n-1] = 1;

    /*
     * Swap round nlo and nhi if they're backwards. Partly in case the
     * user specified them the wrong way round, and also because the
     * scaling above might have negated both.
     */
    if (nlo > nhi) {
        bigint tmp = nhi;
        nhi = nlo;
        nlo = tmp;
    }

    dprint("scaled #2: P = ", P);

    /*
     * Check that there's exactly one root of P in the interval.
     */
    {
        sturm_calculate(P);

        unsigned shi, slo;
        if (!sturm_eval(nhi, shi) || !sturm_eval(nlo, slo))
            throw domain_error("interval for algebraic() has a root "
                               "exactly at one end");

        unsigned nroots = slo - shi;
        if (nroots != 1)
            throw domain_error("interval for algebraic() contains ", nroots,
                               " roots instead of 1");
    }

    /*
     * Evaluate P at nhi, to find the sign we expect it to have.
     */
    bigint phi = 0;
    for (int i = n; i-- > 0 ;)
        phi = phi * nhi + P[i];

    /*
     * Binary-search between nlo and nhi to find a unit interval
     * containing the root. If we hit the root exactly, return a
     * rational.
     */
    while (nhi - nlo > 1) {
        bigint nmid = fdiv(nhi + nlo, 2);
        bigint p = 0;
        for (int i = n; i-- > 0 ;)
            p = p * nmid + P[i];
        if (p == 0) {
            dprint("rational root: ", nmid, "/", d);
            return spigot_rational(nmid, d);
        }
        if (bigint_sign(p) == bigint_sign(phi))
            nhi = nmid;
        else
            nlo = nmid;
    }

    /*
     * Now we're searching for a root in the range (nlo,nlo+1).
     * Transform P again so that the root is translated to (0,1).
     */
    dprint("integer part: nlo=", nlo);

    for (int j = n-1; j-- > 0 ;)
        for (int i = j; i+1 < n; ++i)
            P[i] += nlo * P[i+1];

    dprint("translated: P = ", P);

    Spigot ret = make_unique<Algebraic>(P);
    if (nlo != 0 || d != 1)
        ret = spigot_mobius(move(ret), 1, nlo, 0, d);
    return ret;
}

/*
 * Semantically sensible interface, which takes a vector of integer
 * polynomial coefficients along with a pair of rational endpoints.
 */
Spigot spigot_algebraic(const vector<bigint> &P,
                        bigint nlo, bigint nhi, bigint d)
{
    AlgebraicPrep prep(nlo, d, nhi, d, P);
    return prep.spigot();
}

/*
 * Interface convenient for expr.cpp, which will pass us a vector of
 * Coreables which include the endpoints _and_ the polynomial
 * coefficients.
 */
static Spigot spigot_algebraic_wrapper(argvector<Spigot> &args)
{
    vector<bigint> P;
    bigint nlo, dlo, nhi, dhi;

    if (args.size() < 2)
        throw expr_error("expected at least two arguments to 'algebraic'");
    if (args.syntax.semicolons.size() > 1)
        throw expr_error("expected at most one semicolon in 'algebraic'");
    if (args.syntax.semicolons.size() == 1 &&
        args.syntax.semicolons[0] != 2)
        throw expr_error("expected exactly two interval bounds before "
                         "semicolon in 'algebraic'");

    if (!args[0]->is_rational(&nlo, &dlo) ||
        !args[1]->is_rational(&nhi, &dhi)) {
        // We could, I suppose, attempt to tolerate any old real
        // number as an interval bound, by being prepared to fetch
        // convergents to it until we got one on the right side of the
        // input real at which P had the right sign.
        throw expr_error("expected rational bounds for 'algebraic'");
    }

    // Diagnose zero-sized intervals. We specify that the polynomial
    // has to have a unique root in the _open_ interval (a,b), so that
    // interval certainly has to be non-empty!
    if (nlo * dhi == nhi * dlo)
        throw domain_error("empty interval for 'algebraic'");

    // Get the polynomial coefficients themselves. We permit them to
    // be rationals on input, and then scale up to integers for our
    // internal representation; so to start with, we have to work out
    // our scale factor.
    bigint denominator = 1_bi;
    for (size_t i = 2; i < args.size(); ++i) {
        bigint ncoeff, dcoeff;
        if (!args[i]->is_rational(&ncoeff, &dcoeff))
            throw expr_error("expected rational coefficients for 'algebraic'");
        denominator = lcm(denominator, dcoeff);
    }
    // And now use that to get the integer version.
    for (size_t i = 2; i < args.size(); ++i) {
        bigint ncoeff, dcoeff;
        bool is_rat = args[i]->is_rational(&ncoeff, &dcoeff);
        assert(is_rat); // because the previous loop already checked it
        (void)is_rat;   // squash warning when -DNDEBUG removes that assertion
        P.push_back(ncoeff * denominator / dcoeff);
    }

    // Eliminate trailing zeroes in the array P, so that the length of
    // the array reflects the actual degree of the polynomial.
    while (!P.empty() && P.back() == 0)
        P.pop_back();

    // Deal with trivial polynomials, i.e. constant ones. These can't
    // possibly have a unique root in any non-empty open interval:
    // either they're zero nowhere in it, or everywhere!
    if (P.size() <= 1)
        throw domain_error("degenerate polynomial");

    AlgebraicPrep prep(nlo, dlo, nhi, dhi, P);
    return prep.spigot();
}

static const VariadicFnWrapper expr_algebraic(
    "algebraic", spigot_algebraic_wrapper);
