#ifndef DIM
 #error DIM undefined
#endif

#if DIM != 3 && DIM != 2
 #error DIM out of bounds
#endif

#ifndef SIZE
 #error SIZE undefined
#endif

#ifndef CST_COEFF
 #error CST_COEFF undefined
#endif

#if CST_COEFF != 0 && CST_COEFF != 1
 #error CST_COEFF out of bounds
#endif

#ifndef RADIUS
 #error RADIUS undefined
#endif

#if RADIUS != 1 && RADIUS != 2
 #error RADIUS out of bounds
#endif

#ifndef DENSE
 #error DENSE undefined
#endif

#if DENSE != 0 && DENSE != 1
 #error DENSE out of bounds
#endif

#ifndef TMP_BLOCK
 #error TMP_BLOCK undefined
#endif

#ifndef REPEAT
 #error REPEAT undefined
#endif


// GHOST  = TMP_BLOCK * RADIUS
// GHOST2 = (TMP_BLOCK - 1) * RADIUS
#if   TMP_BLOCK == 7
 #if RADIUS == 1
  #define GHOST 7
  #define GHOST2 6
 #elif RADIUS == 2
  #define GHOST 14
  #define GHOST2 12
 #endif // RADIUS
#elif TMP_BLOCK == 6
 #if RADIUS == 1
  #define GHOST 6
  #define GHOST2 5
 #elif RADIUS == 2
  #define GHOST 12
  #define GHOST2 10
 #endif // RADIUS
#elif TMP_BLOCK == 5
 #if RADIUS == 1
  #define GHOST 5
  #define GHOST2 4
 #elif RADIUS == 2
  #define GHOST 10
  #define GHOST2 8
 #endif // RADIUS
#elif TMP_BLOCK == 4
 #if RADIUS == 1
  #define GHOST 4
  #define GHOST2 3
 #elif RADIUS == 2
  #define GHOST 8
  #define GHOST2 6
 #endif // RADIUS
#elif TMP_BLOCK == 3
 #if RADIUS == 1
  #define GHOST 3
  #define GHOST2 2
 #elif RADIUS == 2
  #define GHOST 6
  #define GHOST2 4
 #endif // RADIUS
#elif TMP_BLOCK == 2
 #if RADIUS == 1
  #define GHOST 2
  #define GHOST2 1
 #elif RADIUS == 2
  #define GHOST 4
  #define GHOST2 2
 #endif // RADIUS
#elif TMP_BLOCK == 1
 #if RADIUS == 1
  #define GHOST 1
  #define GHOST2 0
 #elif RADIUS == 2
  #define GHOST 2
  #define GHOST2 0
 #endif // RADIUS
#endif // TMP_BLOCK

#ifndef GHOST
 #error GHOST undefined (either TMP_BLOCK or RADIUS out of bounds)
#endif // GHOST


#if DIM == 3 //////////////////////////////////////////////////////////////////

Domain global< [ 0, 0, 0 ] to [ 1, 1, 1 ] >
Layout LayoutSol< Real, Node >@finest {
    innerPoints = [ SIZE, SIZE, SIZE ]
    ghostLayers = [ GHOST, GHOST, GHOST ]
    duplicateLayers = [ 1, 1, 1 ]
}
Layout LayoutRHS< Real, Node >@finest {
    innerPoints = [ SIZE, SIZE, SIZE ]
    ghostLayers = [ GHOST2, GHOST2, GHOST2 ]
    duplicateLayers = [ 1, 1, 1 ]
}

#if CST_COEFF == 1 // ---------------------------------------------------------

Stencil Laplace@finest {
    [ 0,  0,  0] =>  4.8
    [ 1,  0,  0] => -0.8
    [-1,  0,  0] => -0.8
    [ 0,  1,  0] => -0.8
    [ 0, -1,  0] => -0.8
    [ 0,  0,  1] => -0.8
    [ 0,  0, -1] => -0.8
#if RADIUS == 2
    [ 2,  0,  0] => -0.3
    [-2,  0,  0] =>  0.3
    [ 0,  2,  0] => -0.3
    [ 0, -2,  0] =>  0.3
    [ 0,  0,  2] => -0.3
    [ 0,  0, -2] =>  0.3
#endif
#if DENSE == 1
    [-1, -1, -1] => -0.1
    [-1, -1,  0] => -0.2
    [-1, -1,  1] => -0.1
    [-1,  0, -1] => -0.2
    [-1,  0,  1] => -0.2
    [-1,  1, -1] => -0.1
    [-1,  1,  0] => -0.2
    [-1,  1,  1] =>  0.1
    [ 0, -1, -1] => -0.2
    [ 0, -1,  1] =>  0.2
    [ 0,  1, -1] =>  0.2
    [ 0,  1,  1] =>  0.2
    [ 1, -1, -1] => -0.1
    [ 1, -1,  0] =>  0.2
    [ 1, -1,  1] =>  0.1
    [ 1,  0, -1] =>  0.2
    [ 1,  0,  1] =>  0.2
    [ 1,  1, -1] =>  0.1
    [ 1,  1,  0] =>  0.2
    [ 1,  1,  1] =>  0.1
#endif // DENSE
}

Globals {
}

#elif CST_COEFF == 0 // -------------------------------------------------------

#if DENSE != 0 && RADIUS != 1
 #error variable coefficients selected, so DENSE must be 0 and RADIUS must be 1
#endif

Layout LayoutSF< ColumnVector<Real,7>, Node >@finest {
    innerPoints = [ SIZE, SIZE, SIZE ]
    ghostLayers = [ GHOST, GHOST, GHOST ]
    duplicateLayers = [ 1, 1, 1 ]
}

Field LaplaceCoeff< global, LayoutSF, None >@finest
Stencil LaplaceStencil@all { // unused...
    [ 0,  0,  0] => 1
    [ 1,  0,  0] => 1
    [-1,  0,  0] => 1
    [ 0,  1,  0] => 1
    [ 0, -1,  0] => 1
    [ 0,  0,  1] => 1
    [ 0,  0, -1] => 1
}
StencilField Laplace< LaplaceCoeff => LaplaceStencil >@finest

Globals {
    Val kappa : Real = 10.0
}

Function getCoefficient(xx : Real, yy : Real, zz : Real) : Real {
    return exp(kappa * ((xx - xx ** 2) * (yy - yy ** 2) * (zz - zz ** 2)))
}

Function InitLaplace() : Unit {
	loop over LaplaceCoeff@finest sequentially {
        Laplace@finest:[ 0,  0,  0] = (getCoefficient(vf_nodePosition_0@finest + 0.5 * vf_cellWidth_0@finest, vf_nodePosition_1@finest, vf_nodePosition_2@finest) + getCoefficient(vf_nodePosition_0@finest - 0.5 * vf_cellWidth_0@finest, vf_nodePosition_1@finest,vf_nodePosition_2@finest)) / (vf_cellWidth_0@finest * vf_cellWidth_0@finest) + (getCoefficient(vf_nodePosition_0@finest, vf_nodePosition_1@finest + 0.5 * vf_cellWidth_1@finest, vf_nodePosition_2@finest) + getCoefficient(vf_nodePosition_0@finest, vf_nodePosition_1@finest - 0.5 * vf_cellWidth_1@finest, vf_nodePosition_2@finest)) / (vf_cellWidth_1@finest * vf_cellWidth_1@finest) + (getCoefficient(vf_nodePosition_0@finest, vf_nodePosition_1@finest, vf_nodePosition_2@finest + 0.5 * vf_cellWidth_2@finest) + getCoefficient(vf_nodePosition_0@finest, vf_nodePosition_1@finest, vf_nodePosition_2@finest - 0.5 * vf_cellWidth_2@finest)) / (vf_cellWidth_2@finest * vf_cellWidth_2@finest)
        Laplace@finest:[ 1,  0,  0] = -1.0 * getCoefficient(vf_nodePosition_0@finest + 0.5 * vf_cellWidth_0@finest, vf_nodePosition_1@finest, vf_nodePosition_2@finest) / (vf_cellWidth_0@finest * vf_cellWidth_0@finest)
        Laplace@finest:[-1,  0,  0] = -1.0 * getCoefficient(vf_nodePosition_0@finest - 0.5 * vf_cellWidth_0@finest, vf_nodePosition_1@finest, vf_nodePosition_2@finest) / (vf_cellWidth_0@finest * vf_cellWidth_0@finest)
        Laplace@finest:[ 0,  1,  0] = -1.0 * getCoefficient(vf_nodePosition_0@finest, vf_nodePosition_1@finest + 0.5 * vf_cellWidth_1@finest, vf_nodePosition_2@finest) / (vf_cellWidth_1@finest * vf_cellWidth_1@finest)
        Laplace@finest:[ 0, -1,  0] = -1.0 * getCoefficient(vf_nodePosition_0@finest, vf_nodePosition_1@finest - 0.5 * vf_cellWidth_1@finest, vf_nodePosition_2@finest) / (vf_cellWidth_1@finest * vf_cellWidth_1@finest)
        Laplace@finest:[ 0,  0,  1] = -1.0 * getCoefficient(vf_nodePosition_0@finest, vf_nodePosition_1@finest, vf_nodePosition_2@finest + 0.5 * vf_cellWidth_2@finest) / (vf_cellWidth_2@finest * vf_cellWidth_2@finest)
        Laplace@finest:[ 0,  0, -1] = -1.0 * getCoefficient(vf_nodePosition_0@finest, vf_nodePosition_1@finest, vf_nodePosition_2@finest - 0.5 * vf_cellWidth_2@finest) / (vf_cellWidth_2@finest * vf_cellWidth_2@finest)
    }
}


#endif // CST_COEFF // --------------------------------------------------------

#elif DIM == 2 ////////////////////////////////////////////////////////////////

Domain global< [ 0, 0 ] to [ 1, 1 ] >
Layout LayoutSol< Real, Node >@finest {
    innerPoints = [ SIZE, SIZE ]
    ghostLayers = [ GHOST, GHOST ]
    duplicateLayers = [ 1, 1 ]
}
Layout LayoutRHS< Real, Node >@finest {
    innerPoints = [ SIZE, SIZE ]
    ghostLayers = [ GHOST2, GHOST2 ]
    duplicateLayers = [ 1, 1 ]
}

#if CST_COEFF == 1 // ---------------------------------------------------------

Stencil Laplace@finest {
    [ 0,  0] =>  3.2
    [ 1,  0] => -0.8
    [-1,  0] => -0.8
    [ 0,  1] => -0.8
    [ 0, -1] => -0.8
#if RADIUS == 2
    [ 2,  0] => -0.3
    [-2,  0] =>  0.3
    [ 0,  2] => -0.3
    [ 0, -2] =>  0.3
#endif
#if DENSE == 1
    [-1, -1] => -0.1
    [-1,  0] => -0.2
    [-1,  1] => -0.1
    [ 0, -1] => -0.2
    [ 0,  1] =>  0.2
    [ 1, -1] =>  0.1
    [ 1,  0] =>  0.2
    [ 1,  1] =>  0.1
#endif // DENSE
}

Globals {
}

#elif CST_COEFF == 0 // -------------------------------------------------------

#if DENSE != 0 && RADIUS != 1
 #error variable coefficients selected, so DENSE must be 0 and RADIUS must be 1
#endif

Layout LayoutSF< ColumnVector<Real,5>, Node >@all {
    innerPoints = [ SIZE, SIZE ]
    ghostLayers = [ GHOST, GHOST ]
    duplicateLayers = [ 1, 1 ]
}

Field LaplaceCoeff< global, LayoutSF, None >@all
Stencil LaplaceStencil@all { // unused...
    [ 0,  0] => 1
    [ 1,  0] => 1
    [-1,  0] => 1
    [ 0,  1] => 1
    [ 0, -1] => 1
}
StencilField Laplace< LaplaceCoeff => LaplaceStencil >@all

Globals {
    Val kappa : Real = 10.0
}

Function getCoefficient (xx : Real, yy : Real) : Real {
    return exp(kappa * ((xx - xx ** 2) * (yy - yy ** 2)))
}

Function InitLaplace() : Unit {
	loop over LaplaceCoeff@finest sequentially {
        Laplace@finest:[ 0,  0] = (getCoefficient(vf_nodePosition_0@finest + 0.5 * vf_cellWidth_0@finest, vf_nodePosition_1@finest) + getCoefficient(vf_nodePosition_0@finest - 0.5 * vf_cellWidth_0@finest, vf_nodePosition_1@finest)) / (vf_cellWidth_0@finest * vf_cellWidth_0@finest) + (getCoefficient(vf_nodePosition_0@finest, vf_nodePosition_1@finest + 0.5 * vf_cellWidth_1@finest) + getCoefficient(vf_nodePosition_0@finest, vf_nodePosition_1@finest - 0.5 * vf_cellWidth_1@finest)) / (vf_cellWidth_1@finest * vf_cellWidth_1@finest)
        Laplace@finest:[ 1,  0] = -1.0 * getCoefficient(vf_nodePosition_0@finest + 0.5 * vf_cellWidth_0@finest, vf_nodePosition_1@finest) / (vf_cellWidth_0@finest * vf_cellWidth_0@finest)
        Laplace@finest:[-1,  0] = -1.0 * getCoefficient(vf_nodePosition_0@finest - 0.5 * vf_cellWidth_0@finest, vf_nodePosition_1@finest) / (vf_cellWidth_0@finest * vf_cellWidth_0@finest)
        Laplace@finest:[ 0,  1] = -1.0 * getCoefficient(vf_nodePosition_0@finest, vf_nodePosition_1@finest + 0.5 * vf_cellWidth_1@finest) / (vf_cellWidth_1@finest * vf_cellWidth_1@finest)
        Laplace@finest:[ 0, -1] = -1.0 * getCoefficient(vf_nodePosition_0@finest, vf_nodePosition_1@finest - 0.5 * vf_cellWidth_1@finest) / (vf_cellWidth_1@finest * vf_cellWidth_1@finest)
    }
}

#endif // CST_COEFF // --------------------------------------------------------

#endif // DIM /////////////////////////////////////////////////////////////////


#ifdef RBGS
 #define SLOTS 1
#else // ifndef RBGS
 #define SLOTS 2
#endif // RBGS

Field Solution< global, LayoutSol, 0.0 >[SLOTS]@finest
Field RHS< global, LayoutRHS, None >@finest

#if DIM == 3
 #define RADIUS_DIM , RADIUS
 #define PLUS_Z + z
#elif DIM == 2
 #define RADIUS_DIM
 #define PLUS_Z
#endif // DIM

Function Smoother() : Unit {
    loop over fragments {
        repeat TMP_BLOCK times with contraction [RADIUS, RADIUS RADIUS_DIM] {
#ifdef RBGS
            loop over Solution@finest where 0 == ((x + y PLUS_Z) % 2) {
                Solution[0]@finest = Solution[0]@finest + ((0.8 / diag(Laplace@finest)) * (RHS@finest - Laplace@finest * Solution[0]@finest))
            }
            loop over Solution@finest where 1 == ((x + y PLUS_Z) % 2) {
                Solution[0]@finest = Solution[0]@finest + ((0.8 / diag(Laplace@finest)) * (RHS@finest - Laplace@finest * Solution[0]@finest))
            }
#else // ifndef RBGS
            loop over Solution@finest {
                Solution[nextSlot]@finest = Solution[active]@finest + (0.8 / diag(Laplace@finest) * (RHS@finest - Laplace@finest * Solution[active]@finest))
            }
            advance Solution@finest
#endif // RBGS
        }
    }
}

Function InitFields() : Unit {
    native('std::srand(42)')
    loop over Solution@finest sequentially {
        Solution[active]@finest = native('((double)std::rand()/RAND_MAX)')
        Solution[nextSlot]@finest = native('((double)std::rand()/RAND_MAX)')
    }
    loop over RHS@finest sequentially {
        RHS@finest = native('((double)std::rand()/RAND_MAX)')
    }
}

Function Benchmark() : Unit {
    print('Cache warmup')
    repeat 1 times {
        Smoother()
    }
    print('Starting benchmark (temporal blocking)')
    startTimer(benchTimer)
    repeat REPEAT times {
        Smoother()
    }
    stopTimer(benchTimer)
    Var time : Real = getTotalFromTimer(benchTimer)
    print('Runtime: ', time)
#if DIM == 3
    print('MLUPs: ', (1.0 * SIZE * SIZE * SIZE * TMP_BLOCK * REPEAT) / time / 1e3)
#elif DIM == 2
    print('MLUPs: ', (1.0 * SIZE * SIZE * TMP_BLOCK * REPEAT) / time / 1e3)
#endif // DIM
}

Function CheckSolution() : Unit {
    loop over Solution@finest sequentially {
        Var s : Real = Solution[active]@finest * Solution[nextSlot]@finest
        if (s == 0.0 || s == 1./0. || s == -1./0. || s != s) {
            print('ERROR: value out of bounds found: ', s)
        }
    }
}

Function Application() : Unit {
    startTimer(setupWatch)
    initGlobals()
    initDomain()
    InitFields()
#if CST_COEFF == 0
    InitLaplace()
#endif
    stopTimer(setupWatch)
    print('Total time to setup: ', getTotalFromTimer(setupWatch))
    Benchmark()
    CheckSolution()
    destroyGlobals()
}
