#ifndef HALIDE_BOUNDARY_CONDITIONS_H
#define HALIDE_BOUNDARY_CONDITIONS_H

/** \file
 * Support for imposing boundary conditions on Halide::Funcs.
 */

#include <vector>

#include "Expr.h"
#include "Func.h"
#include "Lambda.h"

namespace Halide {

/** namespace to hold functions for imposing boundary conditions on
 *  Halide Funcs.
 *
 *  All functions in this namespace transform a source Func to a
 *  result Func where the result produces the values of the source
 *  within a given region and a different set of values outside the
 *  given region. A region is an N dimensional box specified by
 *  mins and extents.
 *
 *  Three areas are defined:
 *      The image is the entire set of values in the region.
 *      The edge is the set of pixels in the image but adjacent
 *          to coordinates that are not
 *      The interior is the image minus the edge (and is undefined
 *          if the extent of any region is 1 or less).
 *
 *  If the source Func has more dimensions than are specified, the extra ones
 *  are unmodified. Additionally, passing an undefined (default constructed)
 *  'Expr' for the min and extent of a dimension will keep that dimension
 *  unmodified.
 *
 *  Numerous options for specifing the outside area are provided,
 *  including replacement with an expression, repeating the edge
 *  samples, mirroring over the edge, and repeating or mirroring the
 *  entire image.
 *
 *  Using these functions to express your boundary conditions is highly
 *  recommended for correctness and performance. Some of these are hard
 *  to get right. The versions here are both understood by bounds
 *  inference, and also judiciously use the 'likely' intrinsic to minimize
 *  runtime overhead.
 *
 */
namespace BoundaryConditions {

namespace Internal {

inline HALIDE_NO_USER_CODE_INLINE void collect_region(Region &collected_args,
                                                      const Expr &a1, const Expr &a2) {
    collected_args.push_back(Range(a1, a2));
}

template<typename... Args>
inline HALIDE_NO_USER_CODE_INLINE void collect_region(Region &collected_args,
                                                      const Expr &a1, const Expr &a2, Args &&...args) {
    collected_args.push_back(Range(a1, a2));
    collect_region(collected_args, std::forward<Args>(args)...);
}

inline const Func &func_like_to_func(const Func &func) {
    return func;
}

template<typename T>
inline HALIDE_NO_USER_CODE_INLINE Func func_like_to_func(const T &func_like) {
    return lambda(_, func_like(_));
}

}  // namespace Internal

/** Impose a boundary condition such that a given expression is returned
 *  everywhere outside the boundary. Generally the expression will be a
 *  constant, though the code currently allows accessing the arguments
 *  of source.
 *
 *  An ImageParam, Buffer<T>, or similar can be passed instead of a
 *  Func. If this is done and no bounds are given, the boundaries will
 *  be taken from the min and extent methods of the passed
 *  object. Note that objects are taken by mutable ref. Pipelines
 *  capture Buffers via mutable refs, because running a pipeline might
 *  alter the Buffer metadata (e.g. device allocation state).
 *
 *  (This is similar to setting GL_TEXTURE_WRAP_* to GL_CLAMP_TO_BORDER
 *   and putting value in the border of the texture.)
 *
 *  You may pass undefined Exprs for dimensions that you do not wish
 *  to bound.
 */
// @{
Func constant_exterior(const Func &source, const Tuple &value,
                       const Region &bounds);
Func constant_exterior(const Func &source, const Expr &value,
                       const Region &bounds);

template<typename T>
HALIDE_NO_USER_CODE_INLINE Func constant_exterior(const T &func_like, const Tuple &value, const Region &bounds) {
    return constant_exterior(Internal::func_like_to_func(func_like), value, bounds);
}

template<typename T>
HALIDE_NO_USER_CODE_INLINE Func constant_exterior(const T &func_like, const Expr &value, const Region &bounds) {
    return constant_exterior(Internal::func_like_to_func(func_like), value, bounds);
}

template<typename T>
HALIDE_NO_USER_CODE_INLINE Func constant_exterior(const T &func_like, const Tuple &value) {
    Region object_bounds;
    for (int i = 0; i < func_like.dimensions(); i++) {
        object_bounds.push_back({Expr(func_like.dim(i).min()), Expr(func_like.dim(i).extent())});
    }

    return constant_exterior(Internal::func_like_to_func(func_like), value, object_bounds);
}
template<typename T>
HALIDE_NO_USER_CODE_INLINE Func constant_exterior(const T &func_like, const Expr &value) {
    return constant_exterior(func_like, Tuple(value));
}

template<typename T, typename... Bounds,
         typename std::enable_if<Halide::Internal::all_are_convertible<Expr, Bounds...>::value>::type * = nullptr>
HALIDE_NO_USER_CODE_INLINE Func constant_exterior(const T &func_like, const Tuple &value,
                                                  Bounds &&...bounds) {
    Region collected_bounds;
    Internal::collect_region(collected_bounds, std::forward<Bounds>(bounds)...);
    return constant_exterior(Internal::func_like_to_func(func_like), value, collected_bounds);
}
template<typename T, typename... Bounds,
         typename std::enable_if<Halide::Internal::all_are_convertible<Expr, Bounds...>::value>::type * = nullptr>
HALIDE_NO_USER_CODE_INLINE Func constant_exterior(const T &func_like, const Expr &value,
                                                  Bounds &&...bounds) {
    return constant_exterior(func_like, Tuple(value), std::forward<Bounds>(bounds)...);
}
// @}

/** Impose a boundary condition such that the nearest edge sample is returned
 *  everywhere outside the given region.
 *
 *  An ImageParam, Buffer<T>, or similar can be passed instead of a Func. If this
 *  is done and no bounds are given, the boundaries will be taken from the
 *  min and extent methods of the passed object.
 *
 *  (This is similar to setting GL_TEXTURE_WRAP_* to GL_CLAMP_TO_EDGE.)
 *
 *  You may pass undefined Exprs for dimensions that you do not wish
 *  to bound.
 */
// @{
Func repeat_edge(const Func &source, const Region &bounds);

template<typename T>
HALIDE_NO_USER_CODE_INLINE Func repeat_edge(const T &func_like, const Region &bounds) {
    return repeat_edge(Internal::func_like_to_func(func_like), bounds);
}

template<typename T>
HALIDE_NO_USER_CODE_INLINE Func repeat_edge(const T &func_like) {
    Region object_bounds;
    for (int i = 0; i < func_like.dimensions(); i++) {
        object_bounds.push_back({Expr(func_like.dim(i).min()), Expr(func_like.dim(i).extent())});
    }

    return repeat_edge(Internal::func_like_to_func(func_like), object_bounds);
}
// @}

/** Impose a boundary condition such that the entire coordinate space is
 *  tiled with copies of the image abutted against each other.
 *
 *  An ImageParam, Buffer<T>, or similar can be passed instead of a Func. If this
 *  is done and no bounds are given, the boundaries will be taken from the
 *  min and extent methods of the passed object.
 *
 *  (This is similar to setting GL_TEXTURE_WRAP_* to GL_REPEAT.)
 *
 *  You may pass undefined Exprs for dimensions that you do not wish
 *  to bound.
 */
// @{
Func repeat_image(const Func &source, const Region &bounds);

template<typename T>
HALIDE_NO_USER_CODE_INLINE Func repeat_image(const T &func_like, const Region &bounds) {
    return repeat_image(Internal::func_like_to_func(func_like), bounds);
}

template<typename T>
HALIDE_NO_USER_CODE_INLINE Func repeat_image(const T &func_like) {
    Region object_bounds;
    for (int i = 0; i < func_like.dimensions(); i++) {
        object_bounds.push_back({Expr(func_like.dim(i).min()), Expr(func_like.dim(i).extent())});
    }

    return repeat_image(Internal::func_like_to_func(func_like), object_bounds);
}

/** Impose a boundary condition such that the entire coordinate space is
 *  tiled with copies of the image abutted against each other, but mirror
 *  them such that adjacent edges are the same.
 *
 *  An ImageParam, Buffer<T>, or similar can be passed instead of a Func. If this
 *  is done and no bounds are given, the boundaries will be taken from the
 *  min and extent methods of the passed object.
 *
 *  (This is similar to setting GL_TEXTURE_WRAP_* to GL_MIRRORED_REPEAT.)
 *
 *  You may pass undefined Exprs for dimensions that you do not wish
 *  to bound.
 */
// @{
Func mirror_image(const Func &source, const Region &bounds);

template<typename T>
HALIDE_NO_USER_CODE_INLINE Func mirror_image(const T &func_like, const Region &bounds) {
    return mirror_image(Internal::func_like_to_func(func_like), bounds);
}

template<typename T>
HALIDE_NO_USER_CODE_INLINE Func mirror_image(const T &func_like) {
    Region object_bounds;
    for (int i = 0; i < func_like.dimensions(); i++) {
        object_bounds.push_back({Expr(func_like.dim(i).min()), Expr(func_like.dim(i).extent())});
    }

    return mirror_image(Internal::func_like_to_func(func_like), object_bounds);
}

// @}

/** Impose a boundary condition such that the entire coordinate space is
 *  tiled with copies of the image abutted against each other, but mirror
 *  them such that adjacent edges are the same and then overlap the edges.
 *
 *  This produces an error if any extent is 1 or less. (TODO: check this.)
 *
 *  An ImageParam, Buffer<T>, or similar can be passed instead of a Func. If this
 *  is done and no bounds are given, the boundaries will be taken from the
 *  min and extent methods of the passed object.
 *
 *  (I do not believe there is a direct GL_TEXTURE_WRAP_* equivalent for this.)
 *
 *  You may pass undefined Exprs for dimensions that you do not wish
 *  to bound.
 */
// @{
Func mirror_interior(const Func &source, const Region &bounds);

template<typename T>
HALIDE_NO_USER_CODE_INLINE Func mirror_interior(const T &func_like, const Region &bounds) {
    return mirror_interior(Internal::func_like_to_func(func_like), bounds);
}

template<typename T>
HALIDE_NO_USER_CODE_INLINE Func mirror_interior(const T &func_like) {
    Region object_bounds;
    for (int i = 0; i < func_like.dimensions(); i++) {
        object_bounds.push_back({Expr(func_like.dim(i).min()), Expr(func_like.dim(i).extent())});
    }

    return mirror_interior(Internal::func_like_to_func(func_like), object_bounds);
}

// @}

}  // namespace BoundaryConditions

}  // namespace Halide

#endif
