/*
 * unary.cpp: Simple unary operations on continued fractions.
 */

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

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

class Abs : public BinaryIntervalSource {
    Spigot x_orig;
    BracketingGenerator bg_x;

    virtual Spigot clone() override;
    virtual void gen_bin_interval(bigint *ret_lo, bigint *ret_hi,
                                  unsigned *ret_bits) override;

  public:
    Abs(Spigot x);
};

MAKE_CLASS_GREETER(Abs);

Abs::Abs(Spigot x)
    : x_orig(x->clone()), bg_x(move(x))
{
    dgreet();
}

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

void Abs::gen_bin_interval(bigint *ret_lo, bigint *ret_hi,
                           unsigned *ret_bits)
{
    bg_x.get_bracket_shift(ret_lo, ret_hi, ret_bits);
    dprint("got x bracket (", *ret_lo, ",", *ret_hi, ") / 2^", *ret_bits);

    /*
     * Adjust the interval we got to ensure it's all positive.
     */
    if (*ret_hi < 0) {
        /*
         * The whole interval was negative, so just reflect it.
         */
        bigint tmp = -*ret_hi;
        *ret_hi = -*ret_lo;
        *ret_lo = tmp;
    } else if (*ret_lo < 0) {
        /*
         * The interval crosses zero, so replace it with an interval
         * with one end at zero and the other end at the maximum
         * extent.
         */
        bigint tmp = -*ret_lo;
        if (*ret_hi < tmp)
            *ret_hi = tmp;
        *ret_lo = 0;
    }
}

class Prepend : public Source {
    /*
     * This class prepends one extra spigot matrix to the stream
     * provided by another spigot, which permits us to apply any
     * Mobius transformation.
     */
    unique_ptr<Source> x;
    bigint a, b, c, d;
    bool x_force;
    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:
    Prepend(Spigot ax, bigint aa, bigint ab, bigint ac, bigint ad);
};

MAKE_CLASS_GREETER(Prepend);

Prepend::Prepend(Spigot ax, bigint aa, bigint ab, bigint ac, bigint ad)
    : x(ax.release()->toSource()), a(aa), b(ab), c(ac), d(ad)
{
    dgreet(Matrix{aa, ab, ac, ad});
}

Spigot Prepend::clone()
{
    return spigot_clone(this, x->clone(), a, b, c, d);
}

bool Prepend::gen_interval(bigint *low, bigint_or_inf *high)
{
    x_force = x->gen_interval(low, high);
    return true; /* force the absorption of our prefix matrix */
}

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

    matrix = { a, b, c, d };
    crReturn(x_force);

    while (1) {
        crReturn(x->gen_matrix(matrix));
    }

    crEnd;
}

Spigot spigot_identity(Spigot x)
{
    return x;
}

Spigot spigot_reciprocal(Spigot a)
{
    bigint n, d;
    if (a->is_rational(&n, &d)) {
        if (n == 0) {
            throw domain_error("reciprocal of zero");
        }
        return spigot_rational(d, n);
    } else {
        return make_unique<Prepend>(move(a), 0, 1, 1, 0);
    }
}

Spigot spigot_rational_mul(Spigot a, const bigint &n, const bigint &d)
{
    bigint an, ad;
    if (a->is_rational(&an, &ad)) {
        return spigot_rational(an * n, ad * d);
    } else {
        return make_unique<Prepend>(move(a), n, 0, 0, d);
    }
}

Spigot spigot_mobius(Spigot x,
                     const bigint &a, const bigint &b,
                     const bigint &c, const bigint &d)
{
    bigint an, ad;
    if (x->is_rational(&an, &ad)) {
        return spigot_rational(a*an + b*ad, c*an + d*ad);
    } else {
        return make_unique<Prepend>(move(x), a, b, c, d);
    }
}

bigint spigot_to_integer(Spigot x, RoundingMode rmode, bool *known_exact)
{
    StaticGenerator sg(move(x));
    bigint n;

    int sign = 0;
    bool odd = false;

    // Collect preliminary information that some rounding modes will
    // need in order to decide what to do.
    switch (rmode) {
      case ROUND_TOWARD_ZERO:
      case ROUND_TO_NEAREST_TOWARD_ZERO:
      case ROUND_TO_NEAREST_AWAY_FROM_ZERO:
        // These require an _approximate_ sign test to decide which of
        // two strategies we're going to use. It only needs to be
        // approximate, because the whole interval (-1/2,+1/2) rounds
        // to 0 anyway.
        n = sg.get_approximate_approximant(5);
        if (n < -2)
            sign = -1;
        else if (n > +2)
            sign = +1;
        else
            sign = 0;
        break;

      case ROUND_AWAY_FROM_ZERO:
        // This requires an _exact_ sign test, because the only value
        // that can return 0 is 0 exactly. If that causes an
        // exactness-hazard hang, then so be it - there was nothing we
        // could do about that.
        sign = sg.get_sign();
        break;

      case ROUND_TO_NEAREST_EVEN:
      case ROUND_TO_NEAREST_ODD:
        // These require an 'approximate parity test'. Specifically,
        // we must narrow down to a unique _rounding boundary_ - i.e.
        // a unique (integer + 1/2) - which is the only one the number
        // might turn out to be on.
        //
        // We do this by getting an approximate value of the number
        // and then taking its floor. If the number is near enough to
        // a rounding boundary (k + 1/2) that there's a chance it
        // might land on it, this will reliably give us k.
        // Contrapositively, in all cases where this approximation
        // might fail to give the exact floor of x, it's because x is
        // near enough to an integer that the 'round to nearest'
        // operation isn't going to go anywhere near a tie-break case
        // anyway, so the value won't be needed.
        n = fdiv(sg.get_approximate_approximant(16), 16);
        odd = (n % 2U) != 0;

      default:
        // In all other cases, no preprocessing needed here.
        break;
    }

    // Now decide what we're actually going to do.
    //
    // In all cases, we're going to take either floor(x), or
    // -floor(-x) (i.e. ceil(x)), and we'll set 'ceiling' to indicate
    // the latter. We may also add 1/2 immediately before taking floor
    // (but after negation), which we indicate by setting add_half.

    bool ceiling;
    bool add_half;

    switch (rmode) {
      case ROUND_DOWN:
        // This is the bare floor function.
        ceiling = false;
        add_half = false;
        break;
      case ROUND_UP:
        // This is the ceiling function.
        ceiling = true;
        add_half = false;
        break;
      case ROUND_TOWARD_ZERO:
        // Floor or ceiling depending on result of approximate sign
        // test: ceiling if negative, floor if positive (and in the
        // middle case either one will work anyway).
        ceiling = sign < 0;
        add_half = false;
        break;
      case ROUND_AWAY_FROM_ZERO:
        // Floor or ceiling depending on result of exact sign test:
        // floor if negative, ceiling if positive. (If the sign is
        // _exactly_ 0, then again, either will work.)
        ceiling = sign > 0;
        add_half = false;
        break;
      case ROUND_TO_NEAREST_UP:
        // The schoolchild rounding rule: floor(x + 1/2).
        ceiling = false;
        add_half = true;
        break;
      case ROUND_TO_NEAREST_DOWN:
        // Conversely, ceil(x - 1/2) = -floor(-x + 1/2).
        ceiling = true;
        add_half = true;
        break;
      case ROUND_TO_NEAREST_TOWARD_ZERO:
        // Decide between the above two policies using the result of
        // the approximate sign test. As before, the middle case will
        // work no matter which we pick.
        ceiling = sign > 0;
        add_half = true;
        break;
      case ROUND_TO_NEAREST_AWAY_FROM_ZERO:
        // Same again, but reverse the sense of the decision.
        ceiling = sign < 0;
        add_half = true;
        break;
      case ROUND_TO_NEAREST_EVEN:
        // Use the parity of the nearest rounding boundary to decide
        // which rounding rule to apply. If the nearest rounding
        // boundary is (even + 1/2), then the number will need to
        // round _down_ if it turns out to be exactly on that
        // boundary, i.e. we behave as if we're doing nearest-down.
        // Conversely, for (odd + 1/2), we behave like nearest-up, so
        // that the exact boundary value will tie-break to the even
        // integer above it.
        ceiling = !odd;      // because nearest-down sets ceiling=true
        add_half = true;
        break;
      case ROUND_TO_NEAREST_ODD:
        // Similar, but exactly flip the sense of the test.
        ceiling = odd;
        add_half = true;
        break;
      default:
        throw internal_error("bad rounding mode in spigot_to_integer");
    }

    // Now we've finished making decisions, and can calculate and
    // return our answer.

    if (ceiling)
        sg.premultiply({-1, 0, 0, 1});

    if (add_half) {
        sg.premultiply({2, 1, 0, 2});

        // If we're rounding to the nearest integer, we never find out
        // whether the input was _exactly_ an integer, because our
        // focus is on the rounding boundaries half an integer away
        // from there.
        if (known_exact)
            *known_exact = false;
        known_exact = nullptr;
    }

    n = sg.get_floor(known_exact);

    if (ceiling)
        n = -n;

    return n;
}

template<RoundingMode rmode>
static Spigot spigot_round_to_integral_value(Spigot a)
{
    return spigot_integer(spigot_to_integer(a->clone(), rmode));
}

static const UnaryFnWrapper expr_floor(
    "floor", spigot_round_to_integral_value<ROUND_DOWN>);
static const UnaryFnWrapper expr_ceil(
    "ceil", spigot_round_to_integral_value<ROUND_UP>);

template<RoundingMode rmode>
static Spigot spigot_fractional_part(Spigot a)
{
    bool known_exact;
    bigint intpart = spigot_to_integer(a->clone(), rmode, &known_exact);
    if (known_exact)
        return spigot_integer(0);
    else
        return spigot_mobius(move(a), 1, -intpart, 0, 1);
}

static const UnaryFnWrapper expr_frac(
    "frac", spigot_fractional_part<ROUND_DOWN>);

template<RoundingMode rmode>
static Spigot spigot_remainder(Spigot a, Spigot b)
{
    Spigot a_over_b = spigot_div(move(a), b->clone());
    return spigot_mul(move(b), spigot_fractional_part<rmode>(move(a_over_b)));
}

static const BinaryFnWrapper expr_fmod(
    "fmod", spigot_remainder<ROUND_TOWARD_ZERO>);

Spigot spigot_mod(Spigot a, Spigot b)
{
    return spigot_remainder<ROUND_DOWN>(move(a), move(b));
}

Spigot spigot_rem(Spigot a, Spigot b)
{
    return spigot_remainder<ROUND_TO_NEAREST_EVEN>(move(a), move(b));
}

#define GENERAL_ROUND_FNS(opt, mode) \
static const UnaryFnWrapper expr_round_##opt( \
    "round_" #opt, spigot_round_to_integral_value<mode>); \
static const UnaryFnWrapper expr_fracpart_##opt( \
    "fracpart_" #opt, spigot_fractional_part<mode>); \
static const BinaryFnWrapper expr_remainder_##opt( \
    "remainder_" #opt, spigot_remainder<mode>);
RoundingMode_ABBR_LIST(GENERAL_ROUND_FNS)
#undef RMODE_OPTION

Spigot spigot_neg(Spigot a)
{
    return spigot_rational_mul(move(a), -1, 1);
}

Spigot spigot_abs(Spigot a)
{
    bigint n, d;
    if (a->is_rational(&n, &d))
        return spigot_rational(bigint_abs(n), bigint_abs(d));

    return make_unique<Abs>(move(a));
}

static const UnaryFnWrapper expr_abs("abs", spigot_abs);

static Spigot spigot_sign(Spigot a)
{
    bigint n, d;
    int toret;
    if (a->is_rational(&n, &d))
        toret = n > 0 ? +1 : n < 0 ? -1 : 0;
    else
        toret = get_sign(move(a));
    return spigot_integer(toret);
}

static const UnaryFnWrapper expr_sign("sign", spigot_sign);
