
/* -*- mode: C -*-  */
/*
   IGraph library.
   Copyright (C) 2006-2012  Gabor Csardi <csardi.gabor@gmail.com>
   334 Harvard st, Cambridge MA, 02139 USA

   This program is free software; you can redistribute it and/or modify
   it under the terms of the GNU General Public License as published by
   the Free Software Foundation; either version 2 of the License, or
   (at your option) any later version.

   This program is distributed in the hope that it will be useful,
   but WITHOUT ANY WARRANTY; without even the implied warranty of
   MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
   GNU General Public License for more details.

   You should have received a copy of the GNU General Public License
   along with this program; if not, write to the Free Software
   Foundation, Inc.,  51 Franklin Street, Fifth Floor, Boston, MA
   02110-1301 USA

*/

#include <igraph.h>
#include <stdlib.h>

int print_vector(igraph_vector_t *v) {
    long int i, n = igraph_vector_size(v);
    for (i = 0; i < n; i++) {
        printf("%li ", (long int) VECTOR(*v)[i]);
    }
    printf("\n");
    return 0;
}

int main() {
    igraph_psumtree_t tree;
    igraph_vector_t vec;
    long int i;
    igraph_real_t sum;

    /* Uniform random numbers */
    igraph_vector_init(&vec, 16);
    igraph_psumtree_init(&tree, 16);
    sum = igraph_psumtree_sum(&tree);
    if (sum != 0) {
        printf("Sum: %f instead of 0.\n", sum);
        return 1;
    }

    for (i = 0; i < 16; i++) {
        igraph_psumtree_update(&tree, i, 1);
    }
    if ((sum = igraph_psumtree_sum(&tree)) != 16) {
        printf("Sum: %f instead of 16.\n", sum);
        return 2;
    }

    for (i = 0; i < 16000; i++) {
        igraph_real_t r = ((double)rand()) / RAND_MAX * sum;
        long int idx;
        igraph_psumtree_search(&tree, &idx, r);
        VECTOR(vec)[idx] += 1;
    }
    for (i = 0; i < 16; i++) {
        if (VECTOR(vec)[i] < 800 || VECTOR(vec)[i] > 1200) {
            return 3;
        }
    }

    /* Nonuniform, even indices have twice as much chance */
    for (i = 0; i < 16; i += 2) {
        igraph_psumtree_update(&tree, i, 2);
    }
    if ((sum = igraph_psumtree_sum(&tree)) != 24) {
        printf("Sum: %f instead of 24.\n", sum);
        return 4;
    }

    igraph_vector_null(&vec);
    for (i = 0; i < 24000; i++) {
        igraph_real_t r = ((double)rand()) / RAND_MAX * sum;
        long int idx;
        igraph_psumtree_search(&tree, &idx, r);
        VECTOR(vec)[idx] += 1;
    }
    for (i = 0; i < 16; i++) {
        if (i % 2 == 0 && (VECTOR(vec)[i] < 1800 || VECTOR(vec)[i] > 2200)) {
            return 5;
        }
        if (i % 2 != 0 && (VECTOR(vec)[i] < 800 || VECTOR(vec)[i] > 1200)) {
            return 6;
        }
    }

    /* Test zero probabilities */
    igraph_psumtree_update(&tree, 0, 0);
    igraph_psumtree_update(&tree, 5, 0);
    igraph_psumtree_update(&tree, 15, 0);
    sum = igraph_psumtree_sum(&tree);

    igraph_vector_null(&vec);
    for (i = 0; i < 20000; i++) {
        igraph_real_t r = ((double)rand()) / RAND_MAX * sum;
        long int idx;
        igraph_psumtree_search(&tree, &idx, r);
        VECTOR(vec)[idx] += 1;
    }
    if (VECTOR(vec)[0] != 0 || VECTOR(vec)[5] != 0 || VECTOR(vec)[15] != 0) {
        return 7;
    }

    igraph_vector_destroy(&vec);
    igraph_psumtree_destroy(&tree);

    /****************************************************/
    /* Non power-of-two vector size                     */
    /****************************************************/

    igraph_vector_init(&vec, 9);
    igraph_psumtree_init(&tree, 9);

    for (i = 0; i < 9; i++) {
        igraph_psumtree_update(&tree, i, 1);
    }
    sum = igraph_psumtree_sum(&tree);

    for (i = 0; i < 9000; i++) {
        igraph_real_t r = ((double)rand()) / RAND_MAX * sum;
        long int idx;
        igraph_psumtree_search(&tree, &idx, r);
        VECTOR(vec)[idx] += 1;
    }
    for (i = 0; i < 9; i++) {
        if (VECTOR(vec)[i] < 800 || VECTOR(vec)[i] > 1200) {
            return 8;
        }
    }

    /* Nonuniform, even indices have twice as much chance */
    for (i = 0; i < 9; i += 2) {
        igraph_psumtree_update(&tree, i, 2);
    }
    sum = igraph_psumtree_sum(&tree);

    igraph_vector_null(&vec);
    for (i = 0; i < 14000; i++) {
        igraph_real_t r = ((double)rand()) / RAND_MAX * sum;
        long int idx;
        igraph_psumtree_search(&tree, &idx, r);
        VECTOR(vec)[idx] += 1;
    }
    for (i = 0; i < 9; i++) {
        if (i % 2 == 0 && (VECTOR(vec)[i] < 1800 || VECTOR(vec)[i] > 2200)) {
            return 9;
        }
        if (i % 2 != 0 && (VECTOR(vec)[i] < 800 || VECTOR(vec)[i] > 1200)) {
            return 10;
        }
    }

    /* Test query */
    for (i = 0; i < igraph_psumtree_size(&tree); i++) {
        if (i % 2 == 0 && igraph_psumtree_get(&tree, i) != 2) {
            return 11;
        }
        if (i % 2 != 0 && igraph_psumtree_get(&tree, i) != 1) {
            return 12;
        }
    }

    /* Test zero probabilities */
    igraph_psumtree_update(&tree, 0, 0);
    igraph_psumtree_update(&tree, 5, 0);
    igraph_psumtree_update(&tree, 8, 0);
    sum = igraph_psumtree_sum(&tree);

    igraph_vector_null(&vec);
    for (i = 0; i < 9000; i++) {
        igraph_real_t r = ((double)rand()) / RAND_MAX * sum;
        long int idx;
        igraph_psumtree_search(&tree, &idx, r);
        VECTOR(vec)[idx] += 1;
    }
    if (VECTOR(vec)[0] != 0 || VECTOR(vec)[5] != 0 || VECTOR(vec)[8] != 0) {
        return 11;
    }

    igraph_vector_destroy(&vec);
    igraph_psumtree_destroy(&tree);

    if (!IGRAPH_FINALLY_STACK_EMPTY) {
        return 13;
    }

    return 0;
}
