// mini-cfd-2D.cpp

// Copied comments from cfd-mini-c.c and coverted it to 2D. 


//========================

// On linux, you can compile and run it like so:
// g++ mini-cfd-2D.cpp -g -I ../include -L ../bin -lHalide `libpng-config --cflags --ldflags` -lpthread -ldl -o mini-cfd-2D
// LD_LIBRARY_PATH=../bin ./mini-cfd-2D

// On os x:
// g++ mini-cfd-2D.cpp -g -I ../include -L ../bin -lHalide `libpng-config --cflags --ldflags` -o mini-cfd-2D
// DYLD_LIBRARY_PATH=../bin ./mini-cfd-2D

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

using namespace Halide;

// mini-cfd parameters
Halide::Expr dx = cast<double>(0.5f);
Halide::Expr factor1 = cast<double>(1.f/12.f);
Halide::Expr factor2 = cast<double>(2.f);
int boxsize = 12000;

//There are 5 components: p, e, u, v, w (density, energy, velocity (3D))
// Each of these components is represented as a 3D array.
// We are only using 2D arrays in this Halide version. 
int main(int argc, char **argv) {
    // Correspond to x and y in a box.
    Var ix("ix"), iy("iy");

    // p_DATA_new(iz,iy,ix) = dx*(iz+iy+ix);
    Halide::Func p_init, e_init, u_init, v_init;
    p_init(iy,ix) = dx*(iy+ix);
    e_init(iy,ix) = cast<double>(1.0f)*dx*(iy+ix);
    u_init(iy,ix) = cast<double>(2.0f)*dx*(iy+ix);
    v_init(iy,ix) = cast<double>(3.0f)*dx*(iy+ix);

    // g(c,z,y,x) = factor1*
    //              (c[z][y][x-2]+7*(c[z][y][x-1]+c[z][y][x])+c[z][y][x+1])
    Halide::Func p_gx, e_gx, u_gx, v_gx;
    p_gx(iy,ix) = factor1*(p_init(iy,ix-2) + cast<double>(7.0f)*p_init(iy,ix-1)
                          + p_init(iy,ix) + p_init(iy,ix+1));
    e_gx(iy,ix) = factor1*(e_init(iy,ix-2) + cast<double>(7.0f)*e_init(iy,ix-1)
                          + e_init(iy,ix) + e_init(iy,ix+1));
    u_gx(iy,ix) = factor1*(u_init(iy,ix-2) + cast<double>(7.0f)*u_init(iy,ix-1)
                          + u_init(iy,ix) + u_init(iy,ix+1));
    v_gx(iy,ix) = factor1*(v_init(iy,ix-2) + cast<double>(7.0f)*v_init(iy,ix-1)
                          + v_init(iy,ix) + v_init(iy,ix+1));

    // g'(c,z,y,x) = factor1*
    //               (c[z][y-2][x]+7*(c[z][y-1][x]+c[z][y][x])+c[z][y+1][x])
    Halide::Func p_gy, e_gy, u_gy, v_gy;
    p_gy(iy,ix) = factor1*(p_init(iy-2,ix) + cast<double>(7.0f)*p_init(iy-1,ix)
                          + p_init(iy,ix) + p_init(iy+1,ix));
    e_gy(iy,ix) = factor1*(e_init(iy-2,ix) + cast<double>(7.0f)*e_init(iy-1,ix)
                          + e_init(iy,ix) + e_init(iy+1,ix));
    u_gy(iy,ix) = factor1*(u_init(iy-2,ix) + cast<double>(7.0f)*u_init(iy-1,ix)
                          + u_init(iy,ix) + u_init(iy+1,ix));
    v_gy(iy,ix) = factor1*(v_init(iy-2,ix) + cast<double>(7.0f)*v_init(iy-1,ix)
                          + v_init(iy,ix) + v_init(iy+1,ix));

    // h(c,z,y,x)  = factor2*
    //               (g(c,z,y,x+1)*g(u_t,z,y,x+1)-g(c,z,y,x)*g(u_t,z,y,x))
    Halide::Func p_hx, e_hx, u_hx, v_hx;
    p_hx(iy,ix) = factor2*(p_gx(iy,ix+1)*u_gx(iy,ix+1)-p_gx(iy,ix)*u_gx(iy,ix));
    e_hx(iy,ix) = factor2*(e_gx(iy,ix+1)*u_gx(iy,ix+1)-e_gx(iy,ix)*u_gx(iy,ix));
    u_hx(iy,ix) = factor2*(u_gx(iy,ix+1)*u_gx(iy,ix+1)-u_gx(iy,ix)*u_gx(iy,ix));
    v_hx(iy,ix) = factor2*(v_gx(iy,ix+1)*u_gx(iy,ix+1)-v_gx(iy,ix)*u_gx(iy,ix));


    // h'(c,z,y,x) = factor2*
    //               (g'(c,z,y+1,x)*g'(v_t,z,y+1,x)-g'(c,z,y,x)*g'(v_t,z,y,x))
    Halide::Func p_hy, e_hy, u_hy, v_hy;
    p_hy(iy,ix) = factor2*(p_gy(iy+1,ix)*v_gy(iy+1,ix)-p_gy(iy,ix)*v_gy(iy,ix));
    e_hy(iy,ix) = factor2*(e_gy(iy+1,ix)*v_gy(iy+1,ix)-e_gy(iy,ix)*v_gy(iy,ix));
    u_hy(iy,ix) = factor2*(u_gy(iy+1,ix)*v_gy(iy+1,ix)-u_gy(iy,ix)*v_gy(iy,ix));
    v_hy(iy,ix) = factor2*(v_gy(iy+1,ix)*v_gy(iy+1,ix)-v_gy(iy,ix)*v_gy(iy,ix));
    
    // p_{t+1}(z,y,x) = h(p_t,z,y,x) + h'(p_t,z,y,x) + h"(p_t,z,y,x)
    Halide::Func p, e, u, v;
    p(iy,ix) = p_hx(iy,ix) + p_hy(iy,ix);
    e(iy,ix) = e_hx(iy,ix) + e_hy(iy,ix);
    u(iy,ix) = u_hx(iy,ix) + u_hy(iy,ix);
    v(iy,ix) = v_hx(iy,ix) + v_hy(iy,ix);
   
    // PA4 TODO: Put scheduling code here.

    // Generate code and do computation.
    // Somewhat artificial computation here.  In real code
    // keep all 5 variables separate.
    Halide::Func sum;
    sum(iy,ix) = p(iy,ix) + e(iy,ix) + u(iy,ix) + v(iy,ix);

    Halide::Image<double> output = sum.realize(boxsize,boxsize);

    /*
    // Output for a sanity check while debugging.
    for (int j=0; j<output.height(); j++) {
        for (int i=0; i<output.width(); i++) {
            printf("output(%d,%d) = %g\n", i, j, output(i,j) );
        }
    }
    */

    return 0;
}