/*
 * erf.cpp: Error function and its relatives.
 */

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

#include <utility>

using std::swap;

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

/*
 * ErfBaseRational class: compute the simplest erf-like function, i.e.
 * _just_ the integral of exp(-x^2), without any scaling or factors of
 * sqrt(pi) anywhere. We put those on later.
 *
 * We compute this for rationals by translating its obvious power
 * series into a spigot description. This converges like an absolute
 * dog for large |x|, but I don't know of any more efficient way to
 * compute erf at those values. (There isn't any sensible range
 * reduction method, for instance.)
 *
 * Let x = n/d. Then we have
 *
 *                  n      n^3        n^5        n^7
 *   erfbase(n/d) = - - -------- + -------- - -------- + ...
 *                  d   1! 3 d^3   2! 5 d^5   3! 7 d^7
 *
 *              n   1     n^2   1     n^2   1     n^2
 *            = - ( - - ----- ( - - ----- ( - - ----- ( ... ) ) )
 *              d   1   1.d^2   3   2.d^2   5   3.d^2
 *
 * so our matrices go
 *
 *   ( n 0 ) ( -1.n^2   1.d^2 ) ( -3.n^2   2.d^2 ) ( -5.n^2   3.d^2 ) ...
 *   ( 0 d ) (    0   1.1.d^2 ) (    0   2.3.d^2 ) (    0   3.5.d^2 )
 *
 * For convenience of the normal-distribution-style Phi functions
 * which are also clients of this class, we also let it compute
 * e^-fx^2 for some rational factor f. (Because multiplying a nice
 * rational 1/2 into each matrix here is much faster than scaling the
 * input by a separately computed copy of _root_ 2.)
 */
class ErfBaseRational : public Source {
    bigint n, d, fn, fd, n2, d2, k, kodd;
    bool still_forcing { true };
    int crState { -1 };

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

  public:
    ErfBaseRational(const bigint &an, const bigint &ad,
                    const bigint &afn, const bigint &afd);
};

MAKE_CLASS_GREETER(ErfBaseRational);

ErfBaseRational::ErfBaseRational(const bigint &an, const bigint &ad,
                    const bigint &afn, const bigint &afd)
        : n(an), d(ad), fn(afn), fd(afd)
{
    dgreet(n, "/", d, " f=", fn, "/", fd);
}

Spigot ErfBaseRational::clone()
{
    return spigot_clone(this, n, d, fn, fd);
}

bool ErfBaseRational::gen_interval(bigint *low, bigint_or_inf *high)
{
    *low = -1;
    *high = +1;
    return true;
}

bool ErfBaseRational::gen_matrix(Matrix &matrix)
{
    crBegin;

    /*
     * The initial anomalous matrix.
     */
    matrix = { n, 0, 0, d };
    crReturn(still_forcing);

    /*
     * Then the regular series.
     *
     * These matrices represent maps of the form x |-> 1/(2k+1) - mx,
     * for increasing integers k and decreasing absolute values of m.
     * We keep returning force=true until 1/(2k+1) + |m| < 1, after
     * which we can be sure that all further matrices will be
     * refining.
     */
    k = 1;
    kodd = 1;
    n2 = fn*n*n;
    d2 = fd*d*d;
    while (1) {
        matrix = { -kodd*n2, k*d2, 0, k*kodd*d2 };
        if (still_forcing && kodd*bigint_abs(matrix[0]) < (kodd-1)*matrix[3])
            still_forcing = false;
        crReturn(still_forcing);
        ++k;
        kodd += 2;
    }

    crEnd;
}

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

Spigot ErfBaseConstructor::construct(const bigint &n, const bigint &d)
{
    return make_unique<ErfBaseRational>(n, d, 1, 1);
}

static const MonotoneConstructorDebugWrapper<ErfBaseConstructor>
expr_ErfBase("ErfBase");

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

Spigot PhiBaseConstructor::construct(const bigint &n, const bigint &d)
{
    return make_unique<ErfBaseRational>(n, d, 1, 2);
}

static const MonotoneConstructorDebugWrapper<PhiBaseConstructor>
expr_PhiBase("PhiBase");

static Spigot erfscale(void)
{
    return spigot_sqrt(spigot_rational_mul(spigot_pi(), 1, 4));
}

static Spigot Phiscale(void)
{
    return spigot_sqrt(spigot_rational_mul(spigot_pi(), 2, 1));
}

static Spigot spigot_erf(Spigot a)
{
    bigint n, d;
    if (a->is_rational(&n, &d) && n == 0)
        return spigot_integer(0);
    return spigot_div(spigot_monotone(make_shared<ErfBaseConstructor>(),
                                      move(a)),
                      erfscale());
}

static const UnaryFnWrapper expr_erf("erf", spigot_erf);

static Spigot spigot_erfc(Spigot a)
{
    bigint n, d;
    if (a->is_rational(&n, &d) && n == 0)
        return spigot_integer(1);
    return spigot_sub(spigot_integer(1),
                      spigot_div(spigot_monotone(
                                     make_shared<ErfBaseConstructor>(),
                                     move(a)),
                                 erfscale()));
}

static const UnaryFnWrapper expr_erfc("erfc", spigot_erfc);

static Spigot spigot_Phi(Spigot a)
{
    bigint n, d;
    if (a->is_rational(&n, &d) && n == 0)
        return spigot_rational(1, 2);
    return spigot_add(spigot_rational(1,2),
                      spigot_div(spigot_monotone(
                                     make_shared<PhiBaseConstructor>(),
                                     move(a)),
                                 Phiscale()));
}

static const UnaryFnWrapper expr_Phi("Phi", spigot_Phi);
static const UnaryFnWrapper expr_norm("norm", spigot_Phi);

static Spigot erfbase_integrand(Spigot x)
{
    return spigot_exp(spigot_neg(spigot_square(move(x))));
}

class ErfInverseConstructor: public InverseConstructor, public Debuggable {
    virtual Spigot f(bigint n, unsigned dbits) override;
    virtual Spigot fprime_max(bigint nlo, bigint nhi, unsigned dbits) override;
  public:
    ErfInverseConstructor();
};

MAKE_CLASS_GREETER(ErfInverseConstructor);

ErfInverseConstructor::ErfInverseConstructor()
{
    dgreet();
}

Spigot ErfInverseConstructor::f(bigint n, unsigned dbits)
{
    dprint("f(", n, " / 2^", dbits, ")");
    Spigot toret = make_unique<ErfBaseRational>(
        n, bigint_power(2, dbits), 1, 1);
    dprint("f done");
    return toret;
}

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

    dprint("fprime_max([", nlo, ", ", nhi, "] / 2^", dbits, ")");
    if ((nhi >> dbits) < 0) {
        dprint("left of inflection point");
        toret = erfbase_integrand(
            spigot_rational(nhi, bigint_power(2, dbits)));
    } else if ((nlo >> dbits) >= 0) {
        dprint("right of inflection point");
        toret = erfbase_integrand(
            spigot_rational(nlo, bigint_power(2, dbits)));
    } else {
        dprint("interval crosses inflection point");
        toret = spigot_integer(1);     // e^{-x^2} is easy when x = 0
    }
    dprint("fprime_max done");
    return toret;
}

class ErfInverseConstructorPos: public ErfInverseConstructor {
  public:
    ErfInverseConstructorPos();
};

class ErfInverseSetup : public Debuggable {
    Spigot t;
  public:
    ErfInverseSetup(Spigot at);
    Spigot setup();
};

MAKE_CLASS_GREETER(ErfInverseSetup);

ErfInverseSetup::ErfInverseSetup(Spigot at) : t(move(at))
{
    dgreet(t->dbg_id());
}

Spigot ErfInverseSetup::setup()
{
    /*
     * Here we compute the inverse of the function computed by
     * ErfBaseRational. I'll call it EB for the purposes of this
     * comment. So EB'(x) = exp(-x^2), with EB(0) = 0.
     *
     * Denote by L the limit of EB(x) as x -> +infinity, which is
     * sqrt(pi)/2.
     *
     * We start by finding an interval of numbers which reliably
     * bracket the right answer. That is, we need a number with EB(x)
     * < a, and one - ideally not much bigger - with EB(x) > a.
     *
     * Lemma 1: for all x,t, exp(-t^2) <= exp(x^2 - 2xt), with
     * equality only at x=t.
     *
     * Proof: exp is strictly increasing, so this is true iff it's
     * still true with the exps stripped off, i.e. we need to show
     * -t^2 <= x^2 - 2xt, which rearranges to 0 <= (x-t)^2, which is
     * clearly non-negative everywhere, and zero iff x-t = 0. []
     *
     * Lemma 2: for x > 0, L-EB(x) < exp(-x^2)/(2x).
     *
     * Proof: L-EB(x) = int_x^inf exp(-t^2) dt          (by definition)
     *               <= int_x^inf exp(x^2-2xt) dt       (by Lemma 1)
     *                = exp(x^2) int_x^inf exp(-2xt) dt
     *                = exp(x^2) [0 - exp(-2x^2)/(-2x)]
     *                = exp(-x^2)/(2x).
     *
     * And the <= is easily seen to be an < by further observing that
     * equality in the Lemma 1 inequality only holds at one single
     * point, so we cannot still have equality once we integrate. []
     *
     * Lemma 3: for all x>0, L-EB(x) < exp(-x^2).
     *
     * Proof: if x >= 1/2, then 1/(2x) <= 1, so by Lemma 2,
     *     L-EB(x) < exp(-x^2)/(2x) <= exp(-x^2).
     *
     * And on the remaining interval [0, 1/2] it's easy to see that
     * the inequality holds, just by plotting the two graphs and
     * looking at them, or else by observing that exp(-x^2) has a
     * negative second derivative throughout that interval (its 2nd
     * derivative doesn't go positive again until x reaches 1/sqrt(2))
     * while 1-EB(x) has a positive one, hence they curve away from
     * each other and can't cross over. []
     *
     * Theorem: for 0 < t < 1, EB(sqrt(-log(L-t))) > t.
     *
     * Proof: let x = sqrt(-log(L-t)). Then we have
     *
     *    EB(x) > L - exp(-x^2)                    (by Lemma 3)
     *          = L (1 - exp(-x^2 - log L))
     *          = L (1 - exp(log(L-t) - log L))    (by defn of x)
     *          = L (1 - (1-t/L))
     *          = L - (L-t)
     *          = t
     *
     * Corollary: if we're looking for inverse-EB of some t > 0, then
     * sqrt(-log(L-t)) is an upper bound on the answer. (For t < 0,
     * just flip all the signs, of course.)
     *
     * To find a lower bound, we just do one iteration of
     * Newton-Raphson. Since EB has a negative 2nd derivative (for
     * x>0), this should always give us something on the far side of
     * the root.
     */

    /*
     * Start by checking the input number's sign. To avoid an
     * exactness hazard at EB^{-1}(0), we do this only approximately,
     * leaving a small interval around 0 where we aren't sure of the
     * sign. In that interval it's safe to choose very simple upper
     * and lower bounds anyway.
     */
    int sign;
    {
        StaticGenerator test(t->clone());
        bigint approx = test.get_approximate_approximant(64);
        if (approx < -2)
            sign = -1;
        else if (approx > +2)
            sign = 1;
        else
            sign = 0;
    }

    bigint nlo, nhi;
    unsigned dbits;

    if (sign == 0) {
        /*
         * The input number could be in the range [-1/16, +1/16].
         * (get_approximate_approximant returned at least -2 and at
         * most +2, and guarantees to be within 2 of the real answer,
         * so the real answer is in the range [-4/64, +4/64].)
         *
         * Since EB^{-1}(1/16) = 0x0.100559..., we'll pick 0x0.1006.
         */
        nlo = -0x803;
        nhi = +0x803;
        dbits = 15;
    } else {
        unsigned dbitslo, dbitshi;

        // Normalise t to positive sign to make this less confusing
        Spigot T = t->clone();
        if (sign < 0)
            T = spigot_neg(move(T));

        // Find an upper bound.
        Spigot upperbound, erfbase_of_upperbound;
        {
            bigint d;
            Spigot L = spigot_rational_mul(spigot_sqrt(spigot_pi()), 1, 2);
            BracketingGenerator boundgen(
                spigot_sqrt(spigot_neg(spigot_log(spigot_sub(
                                                      move(L), T->clone())))));
            while (1) {
                bigint n1, n2;
                boundgen.get_bracket_shift(&n1, &n2, &dbitshi);

                dprint("upper bound: trying (",
                       n1, ", ", n2, ") / 2^", dbitshi);

                d = bigint_power(2, dbitshi);

                int s = parallel_sign_test(
                    spigot_sub(make_unique<ErfBaseRational>(n1, d, 1, 1),
                               T->clone()),
                    spigot_sub(make_unique<ErfBaseRational>(n2, d, 1, 1),
                               T->clone()));

                dprint("s = ", s);

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

            // Having found a nice rational upper bound, replace the
            // theoretical one with that, for speed and simplicity.
            upperbound = spigot_rational(nhi, d);
            erfbase_of_upperbound = make_unique<ErfBaseRational>(nhi, d, 1, 1);
        }

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

        // Now do a Newton-Raphson iteration to find a lower bound.
        {
            Spigot fprime = erfbase_integrand(upperbound->clone());

            BracketingGenerator boundgen(
                spigot_sub(move(upperbound),
                           spigot_div(spigot_sub(
                                          move(erfbase_of_upperbound),
                                          T->clone()), move(fprime))));

            while (1) {
                bigint n1, n2;
                boundgen.get_bracket_shift(&n1, &n2, &dbitslo);

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

                bigint d = bigint_power(2, dbitslo);

                int s = parallel_sign_test(
                    spigot_sub(make_unique<ErfBaseRational>(n1, d, 1, 1),
                               T->clone()),
                    spigot_sub(make_unique<ErfBaseRational>(n2, d, 1, 1),
                               T->clone()));

                dprint("s = ", s);

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

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

    dprint("bounds (", nlo, ", ", nhi, ") / 2^", dbits);
    if (sign < 0) {
        // We sign-flipped everything above to keep the logic simple.
        // Now sign-flip it back, so that we're returning the right
        // answer directly.
        swap(nlo, nhi);
        nlo = -nlo;
        nhi = -nhi;
        dprint("flipped bounds (", nlo, ", ", nhi, ") / 2^", dbits);
    }

    return spigot_invert(
        make_shared<ErfInverseConstructor>(), nlo, nhi, dbits, move(t));
}

static Spigot spigot_erfbaseinv(Spigot t)
{
    return ErfInverseSetup(move(t)).setup();
}
static const UnaryDebugFnWrapper expr_erfbaseinv(
    "ErfBaseInv", spigot_erfbaseinv);

static Spigot spigot_erfinv(Spigot t)
{
    return ErfInverseSetup(spigot_mul(move(t), erfscale())).setup();
}

static const UnaryFnWrapper expr_erfinv("erfinv", spigot_erfinv);
static const UnaryFnWrapper expr_inverf("inverf", spigot_erfinv);

static Spigot spigot_erfcinv(Spigot a)
{
    return spigot_erfinv(spigot_sub(spigot_integer(1), move(a)));
}

static const UnaryFnWrapper expr_erfcinv("erfcinv", spigot_erfcinv);
static const UnaryFnWrapper expr_inverfc("inverfc", spigot_erfcinv);

static Spigot spigot_Phiinv(Spigot a)
{
    Spigot atrans = spigot_mobius(move(a), 2, -1, 0, 2);
    Spigot rootpi = spigot_sqrt(spigot_pi());
    Spigot ebinv = ErfInverseSetup(spigot_mul(move(atrans),
                                              move(rootpi))).setup();
    return spigot_mul(spigot_sqrt(spigot_integer(2)), move(ebinv));
}

static const UnaryFnWrapper expr_Phiinv("Phiinv", spigot_Phiinv);
static const UnaryFnWrapper expr_invPhi("invPhi", spigot_Phiinv);
static const UnaryFnWrapper expr_norminv("norminv", spigot_Phiinv);
static const UnaryFnWrapper expr_invnorm("invnorm", spigot_Phiinv);
static const UnaryFnWrapper expr_probit("probit", spigot_Phiinv);
