/*
    Copyright (C) 2012 Fredrik Johansson

    This file is part of Arb.

    Arb is free software: you can redistribute it and/or modify it under
    the terms of the GNU Lesser General Public License (LGPL) as published
    by the Free Software Foundation; either version 2.1 of the License, or
    (at your option) any later version.  See <http://www.gnu.org/licenses/>.
*/

#include "hypgeom.h"

static __inline__ void
fmpz_poly_evaluate_si(fmpz_t y, const fmpz_poly_t poly, slong x)
{
    fmpz_set_si(y, x);
    fmpz_poly_evaluate_fmpz(y, poly, y);
}

static void
bsplit_recursive_fmpz(fmpz_t P, fmpz_t Q, fmpz_t B, fmpz_t T,
    const hypgeom_t hyp, slong a, slong b, int cont)
{
    if (b - a == 1)
    {
        if (a == 0)
        {
            fmpz_one(P);
            fmpz_one(Q);
        }
        else
        {
            fmpz_poly_evaluate_si(P, hyp->P, a);
            fmpz_poly_evaluate_si(Q, hyp->Q, a);
        }

        fmpz_poly_evaluate_si(B, hyp->B, a);
        fmpz_poly_evaluate_si(T, hyp->A, a);
        fmpz_mul(T, T, P);
    }
    else
    {
        slong m;
        fmpz_t P2, Q2, B2, T2;

        m = (a + b) / 2;

        fmpz_init(P2);
        fmpz_init(Q2);
        fmpz_init(B2);
        fmpz_init(T2);

        bsplit_recursive_fmpz(P, Q, B, T, hyp, a, m, 1);
        bsplit_recursive_fmpz(P2, Q2, B2, T2, hyp, m, b, 1);

        if (fmpz_is_one(B) && fmpz_is_one(B2))
        {
            fmpz_mul(T, T, Q2);
            fmpz_addmul(T, P, T2);
        }
        else
        {
            fmpz_mul(T, T, B2);
            fmpz_mul(T, T, Q2);
            fmpz_mul(T2, T2, B);
            fmpz_addmul(T, P, T2);
        }

        fmpz_mul(B, B, B2);
        fmpz_mul(Q, Q, Q2);
        if (cont)
            fmpz_mul(P, P, P2);

        fmpz_clear(P2);
        fmpz_clear(Q2);
        fmpz_clear(B2);
        fmpz_clear(T2);
    }
}

static void
bsplit_recursive_arb(arb_t P, arb_t Q, arb_t B, arb_t T,
    const hypgeom_t hyp, slong a, slong b, int cont, slong prec)
{
    if (b - a < 4)
    {
        fmpz_t PP, QQ, BB, TT;

        fmpz_init(PP);
        fmpz_init(QQ);
        fmpz_init(BB);
        fmpz_init(TT);

        bsplit_recursive_fmpz(PP, QQ, BB, TT, hyp, a, b, cont);

        arb_set_fmpz(P, PP);
        arb_set_fmpz(Q, QQ);
        arb_set_fmpz(B, BB);
        arb_set_fmpz(T, TT);

        fmpz_clear(PP);
        fmpz_clear(QQ);
        fmpz_clear(BB);
        fmpz_clear(TT);
    }
    else
    {
        slong m;
        arb_t P2, Q2, B2, T2;

        m = (a + b) / 2;

        arb_init(P2);
        arb_init(Q2);
        arb_init(B2);
        arb_init(T2);

        bsplit_recursive_arb(P, Q, B, T, hyp, a, m, 1, prec);
        bsplit_recursive_arb(P2, Q2, B2, T2, hyp, m, b, 1, prec);

        if (arb_is_one(B) && arb_is_one(B2))
        {
            arb_mul(T, T, Q2, prec);
            arb_addmul(T, P, T2, prec);
        }
        else
        {
            arb_mul(T, T, B2, prec);
            arb_mul(T, T, Q2, prec);
            arb_mul(T2, T2, B, prec);
            arb_addmul(T, P, T2, prec);
        }

        arb_mul(B, B, B2, prec);
        arb_mul(Q, Q, Q2, prec);
        if (cont)
            arb_mul(P, P, P2, prec);

        arb_clear(P2);
        arb_clear(Q2);
        arb_clear(B2);
        arb_clear(T2);
    }
}

void
arb_hypgeom_sum(arb_t P, arb_t Q, const hypgeom_t hyp, slong n, slong prec)
{
    if (n < 1)
    {
        arb_zero(P);
        arb_one(Q);
    }
    else
    {
        arb_t B, T;
        arb_init(B);
        arb_init(T);
        bsplit_recursive_arb(P, Q, B, T, hyp, 0, n, 0, prec);
        if (!arb_is_one(B))
            arb_mul(Q, Q, B, prec);
        arb_swap(P, T);
        arb_clear(B);
        arb_clear(T);
    }
}

void
arb_hypgeom_infsum(arb_t P, arb_t Q, hypgeom_t hyp, slong target_prec, slong prec)
{
    mag_t err, z;
    slong n;

    mag_init(err);
    mag_init(z);

    mag_set_fmpz(z, hyp->P->coeffs + hyp->P->length - 1);
    mag_div_fmpz(z, z, hyp->Q->coeffs + hyp->Q->length - 1);

    if (!hyp->have_precomputed)
    {
        hypgeom_precompute(hyp);
        hyp->have_precomputed = 1;
    }

    n = hypgeom_bound(err, hyp->r, hyp->boundC, hyp->boundD,
        hyp->boundK, hyp->MK, z, target_prec);

    arb_hypgeom_sum(P, Q, hyp, n, prec);

    if (arf_sgn(arb_midref(Q)) < 0)
    {
        arb_neg(P, P);
        arb_neg(Q, Q);
    }

    /* We have p/q = s + err i.e. (p + q*err)/q = s */
    {
        mag_t u;
        mag_init(u);
        arb_get_mag(u, Q);
        mag_mul(u, u, err);
        mag_add(arb_radref(P), arb_radref(P), u);
        mag_clear(u);
    }

    mag_clear(z);
    mag_clear(err);
}