#include "Halide.h"
#include <stdio.h>

using namespace Halide;

int call_count[4];
extern "C" HALIDE_EXPORT_SYMBOL int call_counter(int x, int idx) {
    call_count[idx]++;
    return x;
}
HalideExtern_2(int, call_counter, int, int);

void reset_counts() {
    for (int i = 0; i < 4; i++) {
        call_count[i] = 0;
    }
}

void check_counts(int a = 0, int b = 0, int c = 0, int d = 0) {
    int correct[] = {a, b, c, d};
    for (int i = 0; i < 4; i++) {
        if (correct[i] != call_count[i]) {
            printf("call_count[%d] was supposed to be %d but instead is %d\n", i, correct[i], call_count[i]);
            exit(1);
        }
    }
}

int main(int argc, char **argv) {
    Var x;
    Param<bool> toggle1, toggle2;

    {
        // Make a diamond-shaped graph where only one of the two
        // side-lobes is used.
        Func f1, f2, f3, f4;
        f1(x) = x;
        f2(x) = call_counter(f1(x) + 1, 0);
        f3(x) = call_counter(f1(x) + 2, 1);
        f4(x) = select(toggle1, f2(x), f3(x));

        f1.compute_root();
        f2.compute_root();
        f3.compute_root();

        f4.compile_jit();

        reset_counts();
        toggle1.set(true);
        f4.realize({10});
        check_counts(10, 0);

        reset_counts();
        toggle1.set(false);
        f4.realize({10});
        check_counts(0, 10);
    }

    {
        // Make a diamond-shaped graph where the first node can be
        // used in one of two ways.
        Func f1, f2, f3, f4;

        f1(x) = call_counter(x, 0);
        f2(x) = call_counter(f1(x) + 1, 1);
        f3(x) = call_counter(f1(x) + 1, 2);
        f4(x) = select(toggle1, f2(x), 0) + select(toggle2, f3(x), 0);

        f1.compute_root();
        f2.compute_root();
        f3.compute_root();

        f4.compile_jit();

        reset_counts();
        toggle1.set(true);
        toggle2.set(true);
        f4.realize({10});
        check_counts(10, 10, 10);

        reset_counts();
        toggle1.set(false);
        toggle2.set(true);
        f4.realize({10});
        check_counts(10, 0, 10);

        reset_counts();
        toggle1.set(true);
        toggle2.set(false);
        f4.realize({10});
        check_counts(10, 10, 0);

        reset_counts();
        toggle1.set(false);
        toggle2.set(false);
        f4.realize({10});
        check_counts(0, 0, 0);
    }

    {
        // Make a tuple-valued func where one value is used but the
        // other isn't. Currently we need to evaluate both, because we
        // have no way to turn only one of them off, and there might
        // be a recursive dependence of one on the other in an update
        // step.
        Func f1, f2;
        f1(x) = Tuple(call_counter(x, 0), call_counter(x + 1, 1));
        f2(x) = select(toggle1, f1(x)[0], 0) + f1(x)[1];
        f1.compute_root();

        f2.compile_jit();

        reset_counts();
        toggle1.set(true);
        f2.realize({10});
        check_counts(10, 10);

        reset_counts();
        toggle1.set(false);
        f2.realize({10});
        check_counts(10, 10);
    }

    {
        // Make a tuple-valued func where neither value is used when
        // the toggle is false.
        Func f1, f2;
        f1(x) = Tuple(call_counter(x, 0), call_counter(x + 1, 1));
        f2(x) = select(toggle1, f1(x)[0], 0);
        f1.compute_root();
        f2.realize({10});

        f2.compile_jit();

        reset_counts();
        toggle1.set(true);
        f2.realize({10});
        check_counts(10, 10);

        reset_counts();
        toggle1.set(false);
        f2.realize({10});
        check_counts(0, 0);
    }

    {
        // Make our two-toggle diamond-shaped graph again, but use a more complex schedule.
        Func f1, f2, f3, f4;

        f1(x) = call_counter(x, 0);
        f2(x) = call_counter(f1(x) + 1, 1);
        f3(x) = call_counter(f1(x) + 1, 2);
        f4(x) = select(toggle1, f2(x), 0) + select(toggle2, f3(x), 0);

        Var xo, xi;
        f4.split(x, xo, xi, 5);
        f1.compute_at(f4, xo);
        f2.store_root().compute_at(f4, xo);
        f3.store_at(f4, xo).compute_at(f4, xi);

        f4.compile_jit();

        reset_counts();
        toggle1.set(true);
        toggle2.set(true);
        f4.realize({10});
        check_counts(10, 10, 10);

        reset_counts();
        toggle1.set(false);
        toggle2.set(true);
        f4.realize({10});
        check_counts(10, 0, 10);

        reset_counts();
        toggle1.set(true);
        toggle2.set(false);
        f4.realize({10});
        check_counts(10, 10, 0);

        reset_counts();
        toggle1.set(false);
        toggle2.set(false);
        f4.realize({10});
        check_counts(0, 0, 0);
    }

    {
        // Test the interaction with sliding window. We don't need value of
        // g(5), but we need all values of f which is computed inside the g's
        // loop. Make sure we don't skip the computation of f.
        reset_counts();
        Func f("f"), g("g"), h("h");
        f(x) = call_counter(x, 0);
        g(x) = f(x) + f(x - 1);
        h(x) = select(x == 5, 0, g(x));

        f.store_root().compute_at(g, x);
        g.compute_at(h, x);
        h.realize({10});
        check_counts(11);
    }

    printf("Success!\n");
    return 0;
}