#include "bitree.h"
#include "log2.h"
#include "helper.h"

void ec_bitree_to_counts(unsigned long *_this, long int _sz, long int _split)
{
    long int p;
    long int q;
    long int s;

    for (p = _split; p > 1; p = q) {
        q = p >> 1;
        for (s = p - 1; s < _sz; s += p)
            _this[s] -= _this[s - q];
    }
}

void ec_bitree_from_counts(unsigned long *_this, long int _sz)
{
    long int p;
    long int q;
    long int s;

    for (q = 1, p = 2; p <= _sz; q = p, p = q << 1) {
        for (s = p - 1; s < _sz; s += p)
            _this[s] += _this[s - q];
    }
}

unsigned long ec_bitree_get_cumul(const unsigned long *_this, long int _sym)
{
    unsigned long ret = 0;

    while (_sym > 0) {
        ret += _this[_sym - 1];
        _sym &= _sym - 1;
    }
    return ret;
}

unsigned long ec_bitree_get_freq(const unsigned long *_this, long int _sym)
{
    unsigned long ret;
    int p;

    ret = _this[_sym];
    p = _sym & (_sym + 1);
    while (_sym != p) {
        ret -= _this[_sym - 1];
        _sym &= _sym - 1;
    }
    return ret;
}

#if 0
/*Fenwick's approach to re-scaling the counts.
  This tests to be twice as slow or more than the one below, even with inline
   functions enabled, and without loop vectorization (which would make Moffat's
   approach even faster).*/
void ec_bitree_halve(unsigned *_this, int _sz, int _split)
{
    int i;
    for (i = _sz; i-- > 0;) {
        ec_bitree_update(_this, _sz, i,
                         -(int) (ec_bitree_get_freq(_this, i) >> 1));
    }
}
#else
/*Fenwick mentions that this approach is also possible, and this is what
   Moffat uses.
  Simply convert the tree into a simple array of counts, perform the halving,
   and then convert it back.*/
void ec_bitree_halve(unsigned long *_this, long int _sz, long int _split)
{
    int i;

    ec_bitree_to_counts(_this, _sz, _split);
    /*LOOP VECTORIZES. */
    for (i = 0; i < _sz; i++)
        _this[i] -= _this[i] >> 1;
    ec_bitree_from_counts(_this, _sz);
}
#endif

void ec_bitree_update(unsigned long *_this, long int _sz, long int _sym, long int _val)
{
    do {
        _this[_sym] += _val;
        _sym += (_sym + 1) & -(_sym + 1);
    } while (_sym < _sz);
}

long int ec_bitree_find_and_update(unsigned long *_this, long int _sz, long int _split,
                                   unsigned long _freq, unsigned long *_fl, long int _val)
{
    long int base = -1;
    long int test;
    long int fl = 0;

    while (_split > 0) {
        test = base + _split;
        if (test < _sz) {
            if (_freq >= _this[test]) {
                _freq -= _this[test];
                fl += _this[test];
                base = test;
            } else {
                _this[test] += _val;
            }
        }
        _split >>= 1;
    }
    *_fl = fl;
    return base + 1;
}

#ifdef BITREE_TEST
#include <stdio.h>
#include <unistd.h>

/* Simple regression test code. */
static void ec_bitree_print(unsigned long *_this, long int _sz)
{
    int i;

    for (i = 0; i < _sz; i++) {
        printf("%3lu%c", _this[i], i + 1 < _sz ? ' ' : '\n');
    }
}

int main(void)
{
    unsigned long int t[4096];
    long int fl;
    long int s;
    long int i;
    long int s1 = -1;
    long int fl1 = -1;
    long int flprev = -1;

    read(0, t, sizeof(t));
    for (i = 0; i < ARRAY_SIZE(t); i++)
        t[i] %= 4096UL;
    ec_bitree_print(t, ARRAY_SIZE(t));
    ec_bitree_from_counts(t, ARRAY_SIZE(t));
    ec_bitree_print(t, ARRAY_SIZE(t));

    for (i = 0; i <= ARRAY_SIZE(t); i++)
        printf("%3lu%c", ec_bitree_get_cumul(t, i), i < ARRAY_SIZE(t) ? ' ' : '\n');
    printf("\n");

    for (i = 0; i < t[ARRAY_SIZE(t)-1]; i++) {
        s = ec_bitree_find_and_update(t, ARRAY_SIZE(t), __roundup_pow_of_two(ARRAY_SIZE(t)), i, &fl, 0);
        if (s != s1 || fl != fl1 || i == t[ARRAY_SIZE(t)-1]-1) {
            printf("%7ld: %4ld %7ld %4ld\n", i, s, fl, fl-flprev);
            flprev = fl;
            s1 = s;
            fl1 = fl;
        }
    }
    printf("\n");

    for (i = 0; i < ARRAY_SIZE(t); i++) {
        s = ec_bitree_find_and_update(t, ARRAY_SIZE(t), __roundup_pow_of_two(ARRAY_SIZE(t)),
                                      ec_bitree_get_cumul(t, i), &fl, 100);
        ec_bitree_to_counts(t, ARRAY_SIZE(t), __roundup_pow_of_two(ARRAY_SIZE(t)));
        ec_bitree_print(t, ARRAY_SIZE(t));
        ec_bitree_from_counts(t, ARRAY_SIZE(t));
        ec_bitree_update(t, ARRAY_SIZE(t), s, -100);
    }
    fflush(stdout);
    return 0;
}
#endif