/*
 * Lambert W function, i.e. the inverse of f(t) = t e^t.
 *
 * This function allows you to solve several related kinds of
 * equation, by composing it with assorted operations:
 *
 * If      t e^t = x      then t = W(x)
 * If    e^t / t = x      then t = -W(-1/x)
 * If    e^t + t = x      then t = log(W(exp(x)))
 * If    e^t - t = x      then t = log(-W(-exp(-x)))
 * If    t log t = x      then t = exp(W(x))
 * If  t / log t = x      then t = exp(-W(-1/x))
 * If  t + log t = x      then t = W(exp(x))
 * If  t - log t = x      then t = -W(-exp(-x))
 * If        t^t = x      then t = exp(W(log x))
 * If        x^t = t      then t = exp(-W(-log(x)))
 *
 * Since W has two branches (both starting at W(-1/e) = -1, but one
 * heads upwards and is valid for all x >= -1/e while the other heads
 * downwards and is only valid for -1/e <= x < 0), most of the above
 * types of equation can have different numbers of solutions depending
 * on whether the value passed to W is less than -1/e, greater than 0,
 * or in between the two.
 */

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

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

static Spigot xey(const bigint &nx, const bigint &ny, const bigint &d)
{
    return spigot_mul(spigot_rational(nx, d),
                      spigot_exp(spigot_rational(ny, d)));
}

static Spigot xex(const bigint &n, const bigint &d)
{
    return xey(n, n, d);
}

static Spigot xexprime(const bigint &n, const bigint &d)
{
    return xey(n+d, n, d);
}

class XexConstructor : public MonotoneConstructor {
    virtual Spigot construct(const bigint &n, const bigint &d) override;
};

Spigot XexConstructor::construct(const bigint &n, const bigint &d)
{
    return xex(n, d);
}

class XexInverseConstructor: public InverseConstructor, public Debuggable {
    virtual Spigot f(bigint n, unsigned dbits) override;
  public:
    XexInverseConstructor() = default;
};

Spigot XexInverseConstructor::f(bigint n, unsigned dbits)
{
    dprint("f(", n, " / 2^", dbits, ")");
    Spigot toret = xex(n, bigint_power(2, dbits));
    dprint("f done");
    return toret;
}

class XexInverseConstructorPos: public XexInverseConstructor {
    virtual Spigot fprime_max(bigint nlo, bigint nhi, unsigned dbits) override;
  public:
    XexInverseConstructorPos();
};

MAKE_CLASS_GREETER(XexInverseConstructorPos);

XexInverseConstructorPos::XexInverseConstructorPos()
{
    dgreet();
}

Spigot XexInverseConstructorPos::fprime_max(bigint nlo, bigint nhi,
                                            unsigned dbits)
{
    dprint("fprime_max([", nlo, ", ", nhi, "] / 2^", dbits, ")");
    Spigot toret = xexprime(nhi, bigint_power(2, dbits));
    dprint("fprime_max done");
    return toret;
}

class XexInverseConstructorNeg: public XexInverseConstructor {
    virtual Spigot fprime_max(bigint nlo, bigint nhi, unsigned dbits) override;
  public:
    XexInverseConstructorNeg();
};

MAKE_CLASS_GREETER(XexInverseConstructorNeg);

XexInverseConstructorNeg::XexInverseConstructorNeg()
{
    dgreet();
}

Spigot XexInverseConstructorNeg::fprime_max(bigint nlo, bigint nhi,
                                            unsigned dbits)
{
    Spigot toret;

    dprint("fprime_max([", nlo, ", ", nhi, "] / 2^", dbits, ")");
    if ((nhi >> dbits) < -2) {
        dprint("left of inflection point");
        toret = xexprime(nhi, bigint_power(2, dbits));
    } else if ((nlo >> dbits) >= -2) {
        dprint("right of inflection point");
        toret = xexprime(nlo, bigint_power(2, dbits));
    } else {
        dprint("interval crosses inflection point");
        toret = spigot_neg(spigot_exp(spigot_integer(-2)));
    }
    dprint("fprime_max done");
    return toret;
}

class LambertWPosSetup : public Debuggable {
    Spigot x;
  public:
    LambertWPosSetup(Spigot ax);
    Spigot setup();
};

MAKE_CLASS_GREETER(LambertWPosSetup);

LambertWPosSetup::LambertWPosSetup(Spigot ax) : x(move(ax))
{
    dgreet(x->dbg_id());
}

Spigot LambertWPosSetup::setup()
{
    /*
     * This is the positive branch of W, i.e. the one that starts from
     * W(-1/e) = -1 and goes upwards through the origin into the
     * positive x,y quadrant.
     *
     * We need to find two numbers such that f(lo) < x < f(hi),
     * reasonably close together so that f' doesn't vary too much in
     * between.
     *
     * The bounds I've come up with are:
     *
     *   log (1+x) - log log (6+x)  <=  W(x)  <=  log (1+x)
     *
     * (Why 6? Because the lower bound would be tight if you replaced
     * the 6 with (e^e+1)/e ~= 5.94; that's the value at which the LHS
     * equals W(x) when x=-1. 6 is the next integer on the right side
     * of the bound, so much easier to compute.)
     *
     * Here's my slightly handwavy proof of those bounds.
     *
     * To show the upper bound W(x) <= log (1+x):
     *
     * Taking f of both sides, this is true iff
     *          x <= (1+x) log (1+x)
     * <=>      0 <= (1+x) log (1+x) - x
     * Now d/dx ((1+x) log (1+x) - x) = (1+x)/(1+x) + log(1+x) - 1
     *                                = log(1+x)
     * For x >= -1/e, this is increasing, with a single zero at x=0.
     * So (1+x) log (1+x) - x is convex, with a single global min at x=0.
     * And at x=0, it evaluates to 0. Hence, it's >= 0 everywhere, done.
     *
     * To show the lower bound log (1+x) - log log (6+x) <= W(x),
     * we'll show that the two sides are never equal (for x >= -1/e),
     * so either LHS < RHS throughout that range, or LHS > RHS
     * throughout, and it's easy to tell it's the former just by
     * evaluating at any one point.
     *
     * So, suppose we were to have equality:
     *
     * If              log (1+x) - log log (6+x) = W(x)                (1)
     * then taking exp,        (1+x) / log (6+x) = exp(W(x))           (2)
     * Multiplying (1) by (2) and using W(x) exp(W(x)) = x gives us
     *        (1+x) / log (6+x) * (log (1+x) - log log (6+x)) = x
     * <=>      log (1+x) - log log (6+x) - x/(1+x) log (6+x) = 0      (3)
     *
     * We'll show that the derivative of the LHS of (3) is never zero,
     * which proves it's monotonic one way or the other. Then a single
     * spot check will show it's decreasing, and since it also starts
     * off negative at the low end of our x range (-1/e), it must be
     * negative everywhere. So if the LHS of (3) is always negative,
     * then it's never zero, and it would have to be zero if we had
     * equality in (1), and therefore, log (1+x) - log log (6+x) and
     * W(x) don't cross over.
     *
     * Differentiating the LHS of (3) conveniently makes the log(1+x)
     * term turn back into a rational function, so the only logs left
     * are log(6+x). In fact the equation can be regarded as a
     * quadratic in log(6+x), with its coefficients being polynomials
     * in x. Solving the quadratic as if log(6+x) was a separate
     * variable[1], we get
     *
     *   log(6+x) = (x+1)/(x+6) (3 +- sqrt(3-x))                       (4)
     *
     * The LHS of (4) is real. For x>3, the RHS is not real, and so
     * the LHS can't possibly equal it. That leaves the interval
     * [-1/e,3]; on that interval, the LHS is bounded below by
     * log(6-1/e) ~= 1.728, and (here's where I run out of motivation
     * to do the rigorous maximum-finding) plotting both signs of the
     * RHS in gnuplot, you can see that the largest value it takes on
     * that interval is less than 1.53. So the LHS and RHS of (4)
     * can't be equal on [-1/e,3] _or_ on [3,inf), so we're done.
     * Phew!
     *
     * [1] Proof by getting Maxima to do the heavy lifting:
     *   df : diff(log (1+x) - log (log (6+x)) - x/(1+x) * log(6+x), x);
     *   quad : subst(log(x+6)=u, df);
     *   factor(solve(quad, u));
     */

    bigint n, d;
    if (x->is_rational(&n, &d) && n == 0)
        return spigot_integer(0);

    /*
     * Check the lower bound.
     *
     * We don't want to do this _up front_ with complete precision,
     * because it would be an exactness hazard. W is actually defined
     * at -1/e itself, so we'd like an expression of the form W(-1/e)
     * + stuff to evaluate the same as (-1) + stuff without hanging at
     * startup due to the bounds check in W.
     *
     * So we do a quick check for really obvious cases using an
     * approximant, and then defer the precise bounds check using
     * spigot_enforce. (Hence it might come up part way through
     * output, if for example the input to W is a stream of digits
     * from an input file which _turns out_ part way along to be just
     * less than -1/e.)
     */
    {
        StaticGenerator test(spigot_sub(xex(-1, 1), x->clone()));
        bigint approx = test.get_approximate_approximant(64);
        dprint("approximate bounds check -> ", approx, " / 64");
        if (approx > 1)                // immediately obviously out of range
            throw domain_error("W of less than -1/e");
        else if (approx > -2)          // _might_ be out of range
            x = spigot_enforce(move(x), ENFORCE_GE, xex(-1, 1),
                               domain_error("W of less than -1/e"));
    }

    bigint nhi;
    unsigned dbitshi;
    {
        BracketingGenerator boundgen(spigot_log1p(x->clone()));
        while (1) {
            bigint n1, n2;
            boundgen.get_bracket_shift(&n1, &n2, &dbitshi);
            dprint("upper bound: trying (", n1, ", ", n2, ") / 2^", dbitshi);

            bigint d = bigint_power(2, dbitshi);

            int s = parallel_sign_test
                (spigot_sub(xex(n1, d), x->clone()),
                 spigot_sub(xex(n2, d), x->clone()));

            dprint("s = ", s);

            if (s == +1) {
                nhi = n1;
                break;
            } else if (s == +2) {
                nhi = n2;
                break;
            }
        }
    }

    dprint("got upper bound ", nhi, " / 2^", dbitshi);

    bigint nlo;
    unsigned dbitslo;
    {
        BracketingGenerator boundgen(
            spigot_sub(spigot_log1p(x->clone()),
                       spigot_log(spigot_log(spigot_mobius(x->clone(),
                                                           1, 6, 0, 1)))));
        while (1) {
            bigint n1, n2;
            boundgen.get_bracket_shift(&n1, &n2, &dbitslo);
            dprint("lower bound: trying (", n1, ", ", n2, ") / 2^", dbitslo);

            if (n1 < 0 /* optimise to avoid shift */ &&
                n1 < -(1_bi << dbitslo)) {
                nlo = -(1_bi << dbitslo);
                break;
            }

            bigint d = bigint_power(2, dbitslo);

            int s = parallel_sign_test
                (spigot_sub(xex(n1, d), x->clone()),
                 spigot_sub(xex(n2, d), x->clone()));

            dprint("s = ", s);

            if (s == -1) {
                nlo = n1;
                break;
            } else if (s == -2) {
                nlo = n2;
                break;
            }
        }
    }

    dprint("got lower bound ", nlo, " / 2^", dbitslo);

    unsigned dbits = match_dyadic_rationals(nlo, dbitslo, nhi, dbitshi);

    dprint("got overall bounds (", nlo, ", ", nhi, ") / 2^", dbits);

    return spigot_invert(make_shared<XexInverseConstructorPos>(),
                         nlo, nhi, dbits, move(x));
}

class LambertWNegSetup : public Debuggable {
    Spigot x;
  public:
    LambertWNegSetup(Spigot ax);
    Spigot setup();
};

MAKE_CLASS_GREETER(LambertWNegSetup);

LambertWNegSetup::LambertWNegSetup(Spigot ax) : x(move(ax))
{
    dgreet(x->dbg_id());
}

Spigot LambertWNegSetup::setup()
{
    /*
     * This is the negative branch of W, i.e. the one that starts from
     * W(-1/e) = -1, heads downwards, whizzes out to -infinity as x
     * approaches zero, and has no value at all for positive x.
     */

    /*
     * Bounds checking for sensible error messages.
     */
    {
        StaticGenerator test(x->clone());
        if (test.get_sign() >= 0)
            throw domain_error("Wn of a non-negative number");
    }

    /*
     * Lower-bound check, same as the version above.
     */
    {
        StaticGenerator test(spigot_sub(xex(-1, 1), x->clone()));
        bigint approx = test.get_approximate_approximant(64);
        if (approx > 1)                // immediately obviously out of range
            throw domain_error("Wn of less than -1/e");
        else if (approx > -2)          // _might_ be out of range
            x = spigot_enforce(move(x), ENFORCE_GE, xex(-1, 1),
                               domain_error("Wn of less than -1/e"));
    }

    /*
     * Find some bounds.
     *
     * We have -1/e <= x < 0 (probably - unless the deferred bounds
     * check above is going to kick in later). We seek t = Wn(x) < -1,
     * or rather, we seek two values of t such that f(t_0) < x and
     * f(t_1) > x.
     *
     * We have f(log(-x)) = log(-x) e^(log(-x)) = x (-log(-x)) < x
     * (since (-log(-x)) >= 1 given the range of x, and multiplying
     * the negative number x by more than 1 makes it smaller still).
     * So t_0 = log(-x) will do, although we'll just use t_0 = -1 if
     * it looks dangerously close to that.
     *
     * Also have f(2 log(-x)) = 2 log(-x) e^(2 log(-x)) = 2x^2 log(-x)
     * = x (2x log -x), and since 0 < 2x log -x < 2/e given the range
     * of x, this is multiplying the negative x by a positive number
     * _less_ than 1, which makes it larger. So t_1 = 2 log(-x) will
     * also do, and we've got only a factor of two separating our
     * bounds.
     */

    bigint nhi;
    unsigned dbitshi;
    {
        BracketingGenerator boundgen
            (spigot_rational_mul(spigot_log(spigot_neg(x->clone())), 2, 1));
        while (1) {
            bigint n1, n2;
            boundgen.get_bracket_shift(&n1, &n2, &dbitshi);
            dprint("upper bound: trying (", n1, ", ", n2, ") / 2^", dbitshi);

            bigint d = bigint_power(2, dbitshi);

            int s = parallel_sign_test
                (spigot_sub(xex(n1, d), x->clone()),
                 spigot_sub(xex(n2, d), x->clone()));

            dprint("s = ", s);

            if (s == +1) {
                nhi = n1;
                break;
            } else if (s == +2) {
                nhi = n2;
                break;
            }
        }
    }

    bigint nlo;
    unsigned dbitslo;
    {
        BracketingGenerator boundgen
            (spigot_log(spigot_neg(x->clone())));
        while (1) {
            bigint n1, n2;
            boundgen.get_bracket_shift(&n1, &n2, &dbitslo);
            if ((n1 >> dbitslo) >= -2) {
                dprint("lower bound: got (", n1, ", ", n2, ") / 2^", dbitslo,
                       "; just use -1");
                nlo = -1;
                dbitslo = 0;
                break;
            }

            dprint("lower bound: trying (", n1, ", ", n2, ") / 2^", dbitslo);

            bigint d = bigint_power(2, dbitslo);

            int s = parallel_sign_test
                (spigot_sub(xex(n1, d), x->clone()),
                 spigot_sub(xex(n2, d), x->clone()));

            dprint("s = ", s);

            if (s == -1) {
                nlo = n1;
                break;
            } else if (s == -2) {
                nlo = n2;
                break;
            }
        }
    }

    dprint("bounds (", nhi, "/2^", dbitshi, ", ", nlo, "/2^", dbitslo, ")");
    unsigned dbits = match_dyadic_rationals(nlo, dbitslo, nhi, dbitshi);
    dprint("evened bounds (", nhi, ", ", nlo, ") / 2^", dbits, ")");

    return spigot_invert(make_shared<XexInverseConstructorNeg>(),
                         nhi, nlo, dbits, move(x));
}

static Spigot spigot_lambertw_pos(Spigot x)
{
    return LambertWPosSetup(move(x)).setup();
}

static Spigot spigot_lambertw_neg(Spigot x)
{
    return LambertWNegSetup(move(x)).setup();
}

static const UnaryFnWrapper expr_W("W", spigot_lambertw_pos);
static const UnaryFnWrapper expr_Wn("Wn", spigot_lambertw_neg);
