#include "FunctionDAG.h"
#include "Halide.h"
#include <sstream>

using namespace Halide;

extern "C" int mul_by_two(
    halide_buffer_t *input,
    halide_buffer_t *output) {
    if (input->is_bounds_query()) {
        // Bounds query: infer the input dimensions from the output dimensions. In
        // this example, the dimensions are exactly the same
        for (int i = 0; i < 2; ++i) {
            input->dim[i] = output->dim[i];
        }
        return 0;
    }

    // Actual computation: return 2 times x as an example. The first dimension is
    // the innermost, so iterate over it last to avoid inefficient memory access
    // patterns.
    for (int j = 0; j < input->dim[1].extent; ++j) {
        for (int i = 0; i < input->dim[0].extent; ++i) {
            float *out = (float *)output->host + i * output->dim[0].stride +
                         j * output->dim[1].stride;
            float *in = (float *)input->host + i * input->dim[0].stride +
                        j * input->dim[1].stride;
            (*out) = 2 * (*in);
        }
    }
    return 0;
}

void test_coeff_wise(const Target &target) {
    Var x("x"), y("y");

    std::ostringstream with_extern;
    {
        Func f("f"), g("g"), h("h");
        f(x, y) = (x + y) * (x + y);

        Halide::ExternFuncArgument arg = f;
        std::vector<Var> vars = {x, y};
        Halide::Type input_type = Halide::Float(32);
        g.define_extern(
            "mul_by_two",
            {arg},
            input_type,
            vars,
            Halide::NameMangling::C);
        g.function().extern_definition_proxy_expr() = f(x, y) * 2.0f;

        h(x, y) = g(x, y) * 2 + 1;

        h.set_estimate(x, 0, 1000).set_estimate(y, 0, 1000);
        std::vector<Halide::Internal::Function> v;
        v.push_back(h.function());
        Halide::Internal::Autoscheduler::FunctionDAG d(v, target);

        d.dump(with_extern);
    }

    std::ostringstream without_extern;
    {
        Func f("f"), g("g"), h("h");
        f(x, y) = (x + y) * (x + y);
        g(x, y) = f(x, y) * 2.0f;
        h(x, y) = g(x, y) * 2 + 1;

        h.set_estimate(x, 0, 1000).set_estimate(y, 0, 1000);
        std::vector<Halide::Internal::Function> v;
        v.push_back(h.function());
        Halide::Internal::Autoscheduler::FunctionDAG d(v, target);

        d.dump(without_extern);
    }

    // Disabled for now: there is still work to do to populate the jacobian
    // assert(with_extern.str() == without_extern.str());
}

extern "C" int matmul(
    halide_buffer_t *input1,
    halide_buffer_t *input2,
    halide_buffer_t *output) {
    if (input1->is_bounds_query() || input2->is_bounds_query()) {
        // Bounds query: infer the input dimensions from the output dimensions.
        // We leave the k dimension alone since we can't infer it from the output dimensions.
        input1->dim[0].min = output->dim[0].min;
        input1->dim[0].extent = output->dim[0].extent;
        input2->dim[1].min = output->dim[1].min;
        input2->dim[1].extent = output->dim[1].extent;
        return 0;
    }

    // Actual computation: return input1 * input2.
    const int max_i = output->dim[0].min + output->dim[0].extent;
    const int max_j = output->dim[1].min + output->dim[1].extent;
    for (int i = output->dim[0].min; i < max_i; ++i) {
        for (int j = output->dim[1].min; j < max_j; ++j) {
            int pos[2] = {i, j};
            float *out = (float *)output->address_of(pos);
            *out = 0.0f;
            for (int k = 0; k < input1->dim[1].extent; ++k) {
                int pos1[2] = {i, k};
                float *in1 = (float *)input1->address_of(pos1);
                int pos2[2] = {k, j};
                float *in2 = (float *)input2->address_of(pos2);
                (*out) += (*in1) * (*in2);
            }
        }
    }
    return 0;
}

void test_matmul(const Target &target) {
    Var x("x"), y("y"), k("k");
    RDom r(0, 200);
    Halide::Buffer<float> input1(200, 200);
    Halide::Buffer<float> input2(200, 200);

    std::ostringstream with_extern;
    {
        Func mm("mm"), h("h");

        Halide::ExternFuncArgument arg1 = input1;
        Halide::ExternFuncArgument arg2 = input2;
        std::vector<Var> vars = {x, y};
        Halide::Type input_type = Halide::Float(32);
        mm.define_extern(
            "matmul",
            {arg1, arg2},
            {input_type, input_type},
            vars,
            Halide::NameMangling::C);
        mm.function().extern_definition_proxy_expr() = Halide::sum(input1(x, r) * input2(r, y));

        h(x, y) = mm(x, y);

        h.set_estimate(x, 0, 200).set_estimate(y, 0, 200);
        std::vector<Halide::Internal::Function> v;
        v.push_back(h.function());
        Halide::Internal::Autoscheduler::FunctionDAG d(v, target);

        d.dump(with_extern);
    }
    std::ostringstream without_extern;
    {
        Func mm("mm"), h("h");
        mm(x, y) = Halide::sum(input1(x, r) * input2(r, y));
        h(x, y) = mm(x, y);

        h.set_estimate(x, 0, 200).set_estimate(y, 0, 200);
        std::vector<Halide::Internal::Function> v;
        v.push_back(h.function());
        Halide::Internal::Autoscheduler::FunctionDAG d(v, target);

        d.dump(without_extern);
    }

    std::cout << "with_extern:\n " << with_extern.str()
              << "\n\nwithout_extern:\n " << without_extern.str() << "\n";
}

int main(int argc, char **argv) {
    // Use a fixed target for the analysis to get consistent results from this test.
    Target target("x86-64-linux-sse41-avx-avx2");

    test_coeff_wise(target);
    test_matmul(target);

    return 0;
}