nixMark · December 3, 2023 01:05
diff --git a/wfc.hpp b/wfc.hpp
 #pragma once

 #include "util/types.hpp"
 #include "util/std.hpp"
 #include "util/ndarray.hpp"
 #include "util/collections.hpp"
 #include "util/rand.hpp"
 #include "util/hash.hpp"
 #include "util/assert.hpp"
 #include "util/bitset.hpp"
 #include "util/result.hpp"
 #include "util/omp.hpp"

 namespace wfc {
 enum Flags {
    // add rotations of default patterns to pattern set
    FLAGS_ROTATE = 1 << 0,

    // add reflections of default patterns to pattern set
    FLAGS_REFLECT = 1 << 1
 };

 // which pattern selection function to use
 enum class PatternFunction {
    WEIGHTED
 };

 // which cell selection function to use
 enum class NextCellFunction {
    MIN_ENTROPY
 };

 // behavior of borders for pattern calculation
 // EXCLUDE: any pattern which would include a border is not used
 // ZERO: borders have the value T(0)
 // CLAMP: borders are clamped to their nearest defined value
 // WRAP: borders are wrapped around the input data
 enum BorderBehavior {
    EXCLUDE,
    ZERO,
    CLAMP,
    WRAP
 };

 // permute n-dimensional rotations of an array
 template <usize N>
 struct Rotator {};

 template <>
 struct Rotator<2> {
    using V = ivec2;

    // permutes rotations of src
    template <typename T, usize S>
    static inline std::array<std::array<T, S * S>, 4> permute(
        const std::array<T, S * S> &src) {
        std::array<std::array<T, S * S>, 4> dst;
        dst[0] = src;
        dst[1] = rotate_ccw<T, S>(dst[0]);
        dst[2] = rotate_ccw<T, S>(dst[1]);
        dst[3] = rotate_ccw<T, S>(dst[2]);
        return dst;
    }

 private:
    template <typename T, usize S>
    static inline std::array<T, S * S> rotate(
        const std::array<T, S * S> &src) {
        std::array<T, S * S> dst;
        for (usize i = 0; i < S; i++) {
            for (usize j = 0; j < S; j++) {
                ndarray::at(V(S), &dst[0], { i, j }) =
                    ndarray::at(V(S), &src[0], { S - j - 1, i });
            }
        }
        return dst;
    }
 };

 // permutes n-dimensional reflections of an array
 template <usize N>
 struct Reflector {
    using V = math::vec<N, int, math::defaultp>;

    template <typename T, usize S, usize VOL = math::cexp::pow(S, N)>
    static inline auto permute(
        const std::array<T, S * S> &src) {
        std::array<std::array<T, VOL>, math::cexp::pow(2, N)> dst;

        // iterate over possible permutations of axes
        usize i = 0;
        ndarray::each(
            V(2),
            [&](const V &v) {
                const auto which = math::vec<N, bool, math::defaultp>(v);
                dst[i] = src;
                for (usize j = 0; j < N; j++) {
                    if (which[j]) {
                        dst[i] = reflect_axis<T, S>(dst[i], j);
                    }
                }
                i++;
            });

        return dst;
    }

 private:
    template <typename T, usize S, usize VOL = math::cexp::pow(S, N)>
    static inline auto reflect_axis(
        const std::array<T, VOL> &src, usize axis) {
        std::array<T, VOL> dst;
        ndarray::each(
            V(S),
            [&](const V &v) {
                V u = v;
                u[axis] = S - u[axis] - 1;
                ndarray::at(V(S), &dst[0], v) =
                    ndarray::at(V(S), &src[0], u);
            });
        return dst;
    }
 };

 // neighbors in N dimensions
 template <usize N>
 struct Neighbors {};

 template <>
 struct Neighbors<2> {
    static constexpr std::array<ivec2, 4> neighbors =
        {
            ivec2(-1, 0),
            ivec2(1, 0),
            ivec2(0, -1),
            ivec2(0, 1),
        };
 };

 template <>
 struct Neighbors<3> {
    static constexpr std::array<ivec3, 9> neighbors =
        {
            ivec3(-1, 0, 0),
            ivec3(1, 0, 0),
            ivec3(0, -1, 0),
            ivec3(0, 1, 0),
            ivec3(0, 0, -1),
            ivec3(0, 0, 1),
        };
 };

 // implements n-dimensional wave function collapse with the "overlapping model"
 // T: pattern data type
 // N: number of dimension
 // S: size of patterns (MUST BE ODD!)
 // D: bitset size
 // V (defaulted): ivecN
 template <
    typename T,
    usize N,
    usize S,
    usize D,
    typename V = math::vec<N, int, math::defaultp>>
    requires ((S % 2) == 1)
 struct WFC {
    // forward declarations
    struct Pattern;
    struct Element;
    struct Wave;

    // function which takes current wave and returns next cell to collapse
    using NextCellFn = std::function<Element&(Wave&)>;

    // function which chooses the pattern to collapse a cell to
    using PatternFn = std::function<usize(Wave&, Element&)>;

    // function used for optional callbacks during collapse process
    using CallbackFn = std::function<void(const Wave&)>;

    // S ** N (volume of one pattern)
    static constexpr auto VOL = math::cexp::pow(S, N);

    // pattern derived from input data
    struct Pattern {
        // unique pattern ID
        usize id;

        // normalized frequency of this pattern in input data
        f32 frequency;

        // value at the center of this pattern
        T value;

        // pattern data
        std::array<T, VOL> data;

        // valid neighbors on each side of this pattern
        std::array<Bitset<D>, 2 * N> valid;

        explicit Pattern(const std::array<T, VOL> &data)
            : data(data),
              value(ndarray::at(V(S), &data[0], V(S) / 2)) {}

        // NOTE: comparison is on data only
        auto operator<=>(const Pattern &other) const {
            return (*this) == other ?
                std::strong_ordering::equal
                : this->hash() <=> other.hash();
        }

        // NOTE: equality comparison is on data only
        bool operator==(const Pattern &other) const {
            return this->hash() == other.hash()
                && this->data == other.data;
        }

        inline u64 hash() const {
            if (this->_hash) {
                return this->_hash;
            }

            u64 v = 0x12345;
            for (const auto &x : this->data) {
                v ^= x + 0x9e3779b9 + (v << 6) + (v >> 2);
            }
            return this->_hash = v;
        }

 private:
        // stored hash, only calculated once
        mutable u64 _hash = 0;
    };

    // wave element
    struct Element {
        // position in output space
        V pos;

        // coefficient, marked bits are VALID choices
        Bitset<D> c;

        // number of valid bits remaining
        usize popcnt;

        // value post-collapse (std::nullopt if not collapsed)
        std::optional<T> value = std::nullopt;

        // memoized entropy values
        f32 sum_weights = 0.0f;
        f32 sum_weight_log_weights = 0.0f;
        f32 entropy = 0.0f;

        // intialize entropy values, coefficient
        void init(const Wave &w, const Bitset<D> &mask) {
            this->c = mask;
            this->popcnt = this->c.popcnt();

            for (
                auto it = this->c.begin_on();
                it != this->c.end_on();
                it++) {
                const auto weight = w.wfc.patterns[*it].frequency;
                this->sum_weights += weight;
                this->sum_weight_log_weights += weight * std::log(weight);
            }
        }

        // applies a mask to this element's coefficient, updating memoized
        // entropy values
        // returns false on failure (0 popcnt/contradiction)
        bool apply(const Wave &w, const Bitset<D> &mask) {
            // get bits which are on in the current coefficient but off in the
            // new mask. exit early if nothing changes.
            const auto diff = this->c & mask;
            if (diff.popcnt() == 0) {
                return true;
            }

            this->c &= mask;

            for (
                auto it = diff.begin_on();
                it != diff.end_on();
                it++) {
                const auto weight = w.wfc.patterns[*it].frequency;
                this->sum_weights -= weight;
                this->sum_weight_log_weights -= weight * std::log(weight);
            }

            this->entropy =
                std::log(this->sum_weights)
                    - (this->sum_weight_log_weights / this->sum_weights);
            this->popcnt = this->c.popcnt();
            return this->popcnt != 0;
        }

        // collapse this element to pattern n
        bool collapse(usize n, const T &value) {
            if (!this->c[n]) {
                return false;
            }

            this->value = value;
            this->c.reset();
            this->c.set(n);
            this->popcnt = 1;
            this->entropy = 0.0f;
            this->sum_weights = 0.0f;
            this->sum_weight_log_weights = 0.0f;
            return true;
        }

        // returns true if this element is collapsed
        bool collapsed() const {
            return static_cast<bool>(this->value);
        }
    };

    // wave to collapse
    struct Wave {
        const WFC &wfc;

        // size of output wave
        V size_wave;

        // output wave elements
        std::vector<Element> wave;

        // optional preset values
        const std::optional<T> *preset;

        // total number of collapsed elements
        usize num_collapsed = 0;

        Wave(
            const WFC &wfc,
            const V &size_wave,
            const std::optional<T> *preset = nullptr)
            : wfc(wfc),
              size_wave(size_wave),
              preset(preset) {}

        // collapses the specified wave element to the only remaining
        // possibility, or n if specified
        result::Result<void, V> collapse(
            Element &e,
            usize n = std::numeric_limits<usize>::max()) {
            n = (n == std::numeric_limits<usize>::max()) ? e.c.nth_set(0) : n;

            const auto &p = this->wfc.patterns[n];
            if (!e.collapse(n, p.value)) {
                return result::Err(e.pos);
            }

            this->num_collapsed++;
            return result::Ok();
        }

        // observe specified wave element, collapsing it to one of its possible
        // values
        result::Result<void, V> observe(Element &e) {
            const auto n = this->wfc.pattern_fn(*this, e);
            return this->collapse(e, n);
        }

        // propagate the current value of the specified wave element
        // returns Ok on success, erroneous position V on failure
        // (contradiction)
        result::Result<void, V> propagate(Element &to_propagate) {
            // DFS elements to update
            std::stack<Element*> es;
            es.push(to_propagate);

            // if not std::nullopt, there is an unresolvable contradiction in
            // the wave
            std::optional<V> contradiction = std::nullopt;

            // propagate each stack entry to neighbors
            while (!es.empty()) {
                auto &e = *es.top();
                es.pop();

                // get only valid non-collapsed neighbors
                std::array<Element*, N * 2> neighbors;
                for (usize i = 0; i < N * 2; i++) {
                    neighbors[i] = nullptr;

                    const auto &n = Neighbors<N>::neighbors[i];
                    const auto pos_n = e.pos + n;

                    if (!ndarray::in_bounds(this->size_wave, pos_n)) {
                        continue;
                    }

                    auto &e_n =
                        ndarray::at(
                            this->size_wave,
                            &this->wave[0],
                            pos_n);

                    if (e_n.collapsed()) {
                        continue;
                    }

                    neighbors[i] = &e_n;
                }

                // compute superpatterns for valid neighbors
                std::array<Bitset<D>, N * 2> neighbor_patterns;
                for (
                    auto it = e.c.begin_on();
                    it != e.c.end_on();
                    it++) {
                    const auto &p = this->wfc.patterns[*it];
                    for (usize i = 0; i < N * 2; i++) {
                        if (neighbors[i]) {
                            neighbor_patterns[i] |= p.valid[i];
                        }
                    }
                }

                // apply superpatterns
                for (usize i = 0; i < N * 2; i++) {
                    if (!neighbors[i]) {
                        continue;
                    }

                    auto &e_n = *neighbors[i];
                    const auto popcnt_old = e_n.popcnt;
                    e_n.apply(*this, neighbor_patterns[i]);

                    if (e_n.popcnt != popcnt_old) {
                        if (e_n.popcnt == 0) {
                            // zero popcount = failure/contradiction
                            this->collapse(e_n, 0);
                            contradiction = e_n.pos;
                            break;
                        } else if (e_n.popcnt == 1) {
                            // one popcount = one remaining possibility/collapse
                            auto res = this->collapse(e_n);
                            if (res.isErr()) {
                                return res;
                            }
                        }

                        // propagate changed value
                        es.push(&e_n);
                    }
                }
            }

            if (contradiction) {
                return result::Err(*contradiction);
            }

            if (this->wfc.on_propagate) {
                (*this->wfc.on_propagate)(*this);
            }

            return result::Ok();
        }

        // collapse the wave
        result::Result<void, V> collapse() {
            // initialize the wave, all patterns are valid for each element
            this->wave = std::vector<Element>(math::prod(this->size_wave));

            // safe to parallelize, only one element modified at a time
            #pragma omp parallel for
            for (usize i = 0; i < this->wave.size(); i++) {
                auto &e = this->wave[i];
                e.pos = ndarray::unravel_index(this->size_wave, i);
                e.init(*this, this->wfc.mask_used);
            }

            // load preset values if present
            if (this->preset) {
                ndarray::each(
                    this->size_wave,
                    [&](const V &pos) {
                        const auto &p =
                            ndarray::at(this->size_wave, &this->preset[0], pos);
                        if (p) {
                            auto &e =
                                ndarray::at(
                                    this->size_wave,
                                    &this->wave[0],
                                    pos);
                            e.value = *p;
                        }
                    });
            }

            // collapse the wave
            while (this->num_collapsed != this->wave.size()) {
                auto &e = this->wfc.next_cell_fn(*this);

                const auto res_observe = this->observe(e);
                if (res_observe.isErr()) {
                    return res_observe;
                }

                const auto res_prop = this->propagate(e);
                if (res_prop.isErr()) {
                    return res_prop;
                }
            }

            // success!
            return result::Ok();
        }
    };

    // size of input data
    V size_in;

    // input data
    const T *in;

    // possible patterns
    std::vector<Pattern> patterns;

    // function to select pattern to collapse to
    PatternFn pattern_fn;

    // function to select next cell to collapse
    NextCellFn next_cell_fn;

    // behavior of patterns on borders
    BorderBehavior border_behavior;

    // mask of used bits of coefficient bitsets
    // bits are also zeroed for disallowed patterns (patterns without valid
    // neighbors)
    Bitset<D> mask_used;

    // random generator
    Rand rand;

    // options
    usize flags;

    // optional callback to be called after propagation of each collapsed wave
    // element
    std::optional<CallbackFn> on_propagate = std::nullopt;

    explicit WFC(
        const V &size_in,
        const T *in,
        PatternFunction pattern_function,
        NextCellFunction next_cell_function,
        BorderBehavior border_behavior,
        Rand &&rand,
        usize flags)
        : size_in(size_in),
          in(in),
          border_behavior(border_behavior),
          rand(std::move(rand)),
          flags(flags) {
        // get pattern data from input at specified location
        // returns std::nullopt if pattern is made illegal by border behavior
        const auto data_at =
            [&](const V &center) -> std::optional<std::array<T, VOL>> {
                const auto base = center - (V(S) / 2);

                if (border_behavior == BorderBehavior::EXCLUDE) {
                    if (!ndarray::in_bounds(size_in, base)
                            || !ndarray::in_bounds(size_in, base + V(S) - 1)) {
                        return std::nullopt;
                    }
                }

                std::array<T, VOL> dst;
                ndarray::each(
                    V(S),
                    [&](const V &offset) {
                        std::optional<T> override = std::nullopt;
                        V pos = base + offset;

                        switch (border_behavior) {
                            case BorderBehavior::EXCLUDE:
                                break;
                            case BorderBehavior::ZERO:
                                if (!ndarray::in_bounds(size_in, pos)) {
                                    override = T(0);
                                }
                                break;
                            case BorderBehavior::CLAMP:
                                pos = math::clamp(pos, V(0), size_in - V(1));
                                break;
                            case BorderBehavior::WRAP:
                                pos = ((pos % size_in) + size_in) % size_in;
                                break;
                        };

                        ndarray::at(V(S), &dst[0], offset) =
                            override ?
                                *override
                                : ndarray::at(size_in, in, pos);
                    });
                return std::make_optional(dst);
            };

        // slide along input data creating patterns
        ndarray::each(
            size_in,
            [&](const V &p) {
                if (const auto data = data_at(p)) {
                    this->patterns.emplace_back(Pattern(data));
                }
            });

        // calculate per-pattern frequencies (normalization occurs later once
        // patterns have been deduplicated)
        std::unordered_map<u64, usize> pattern_hash_to_freq;
        for (const auto &p : this->patterns) {
            pattern_hash_to_freq[p.hash()]++;
        }

        // assign base frequencies (again, not normalized)
        for (auto &p : this->patterns) {
            p.frequency = static_cast<f32>(pattern_hash_to_freq[p.hash()]);
        }

        // removes duplicate patterns from this->patterns
        const auto deduplicate =
            [&]() {
                std::sort(this->patterns.begin(), this->patterns.end());
                this->patterns.erase(
                    std::unique(this->patterns.begin(), this->patterns.end()),
                    this->patterns.end());
            };

        // remove duplicate patterns
        deduplicate();

        // add rotations, reflections
        const auto base_patterns = this->patterns;
        this->patterns.clear();
        for (const auto &p : base_patterns) {
            this->patterns.push_back(p);

            if (flags & FLAGS_ROTATE) {
                ASSERT(N == 2, "can only rotate patterns in 2 dimensions");
                if constexpr (N == 2) {
                    const auto permutations =
                        Rotator<N>::template permute<T, S>(p.data);
                    for (usize i = 1; i < permutations.size(); i++) {
                        auto &q =
                            this->patterns.emplace_back(
                                pattern(permutations[i]));
                        q.frequency = p.frequency;
                    }
                }
            }
            if (flags & FLAGS_REFLECT) {
                const auto permutations =
                    Reflector<N>::template permute<T, S>(p.data);
                for (usize i = 1; i < permutations.size(); i++) {
                    auto &q =
                        this->patterns.emplace_back(
                            pattern(permutations[i]));
                    q.frequency = p.frequency;
                }
            }
        }

        // deduplicate again to remove duplicates added by reflection/rotation
        deduplicate();

        // normalize pattern frequencies, pattern set will no longer change
        f32 frequency_total = 0.0f;
        for (const auto &p : this->patterns) {
            frequency_total += p.frequency;
        }

        for (auto &p : this->patterns) {
            p.frequency /= frequency_total;
        }

        // compute used mask, assign IDs
        this->mask_used.reset();
        for (usize i = 0; i < this->patterns.size(); i++) {
            this->patterns[i].id = i;
            this->mask_used.set(i);
        }

        // compuate valid patterns around each pattern (adjacency)
        // check against all overlapping slots for each pattern pair (p, q) for
        // every side
        //
        // safe to parallelize on p because only p is being written to and not
        // read from :)
        #pragma omp parallel for
        for (auto &p : this->patterns) {
            for (const auto &q : this->patterns) {
                for (usize i = 0; i < 2 * N; i++) {
                    const auto &n = Neighbors<N>::neighbors[i];
                    bool valid = true;
                    ndarray::each(
                        V(S),
                        [&](const auto &offset_q) {
                            if (!valid) {
                                return;
                            }

                            // compute offset into p's data
                            const auto offset_p = n + offset_q;
                            if (!ndarray::in_bounds(V(S), offset_p)) {
                                return;
                            }

                            const auto
                                v_p =
                                    ndarray::at(V(S), &p.data[0], offset_p),
                                v_q =
                                    ndarray::at(V(S), &q.data[0], offset_q);

                            // data must be equal at each offset for patterns to
                            // match (be valid neighbors)
                            if (v_p != v_q) {
                                valid = false;
                            }
                        });

                    if (valid) {
                        p.valid[i].set(q.id);
                    }
                }
            }
        }

        switch (pattern_function) {
            case PatternFunction::WEIGHTED:
                this->pattern_fn = this->pattern_weighted();
                break;
        }

        switch (next_cell_fn) {
            case NextCellFunction::MIN_ENTROPY:
                this->next_cell_fn = this->next_cell_min_entropy();
                break;
        }
    }

    // collapses a wave into the specified output
    // returns true on success
    bool collapse(
        const V &size_out,
        T *out,
        const std::optional<T> *preset = nullptr) const {
        auto w = Wave(*this, size_out, preset);
        const auto res = w.collapse();
        if (res.isErr()) {
            return false;
        }

        // map elements to output
        ndarray::each(
            w.size_wave,
            [&](const V &pos) {
                ndarray::at(size_out, out, pos) =
                    *ndarray::at(size_out, &w.wave, pos).value;
            });

        return true;
    }

    // returns a function which selects a pattern based on the distribution of
    // patterns in the input data
    PatternFn pattern_weighted() {
        return [&](Wave &w, Element &e) -> usize {
            f32 sum_cs = 0.0f;
            std::vector<f32> cs(this->patterns.size());
            for (
                auto it = e.c.begin_on();
                it != e.c.end_on();
                it++) {
                cs[*it] = this->patterns[*it].frequency;
                sum_cs += cs[*it];
            }

            const auto r = this->rand.template next<f32>(0.0f, sum_cs);
            f32 acc = 0.0f;
            for (usize i = 0; i < cs.size(); i++) {
                acc += cs[i];

                if (acc >= r) {
                    return i;
                }
            }

            ASSERT(false, "failed to select pattern for {}", e.pos);
            return e.c.nth_set(0);
        };
    }

    // returns a function which selects the next cell based on finding the cell
    // with the minimum entropy in those remaining
    NextCellFn next_cell_min_entropy() {
        return [&](Wave &w) -> Element& {
            f32 min = 1e4;
            Element *argmin = nullptr;

            // min entropy with a bit of noise
            for (auto &e : w.wave) {
                if (!e.collapsed()
                        && e.entropy < min
                        && e.entropy + rand.next<f32>(0.0f, 1e-6) < min) {
                    argmin = &e;
                }
            }

            ASSERT(argmin);
            return *argmin;
        };
    }
 };
 }
	#pragma once

	#include "util/types.hpp"
	#include "util/std.hpp"
	#include "util/ndarray.hpp"
	#include "util/collections.hpp"
	#include "util/rand.hpp"
	#include "util/hash.hpp"
	#include "util/assert.hpp"
	#include "util/bitset.hpp"
	#include "util/result.hpp"
	#include "util/omp.hpp"

	namespace wfc {
	enum Flags {
	// add rotations of default patterns to pattern set
	FLAGS_ROTATE = 1 << 0,

	// add reflections of default patterns to pattern set
	FLAGS_REFLECT = 1 << 1
	};

	// which pattern selection function to use
	enum class PatternFunction {
	WEIGHTED
	};

	// which cell selection function to use
	enum class NextCellFunction {
	MIN_ENTROPY
	};

	// behavior of borders for pattern calculation
	// EXCLUDE: any pattern which would include a border is not used
	// ZERO: borders have the value T(0)
	// CLAMP: borders are clamped to their nearest defined value
	// WRAP: borders are wrapped around the input data
	enum BorderBehavior {
	EXCLUDE,
	ZERO,
	CLAMP,
	WRAP
	};

	// permute n-dimensional rotations of an array
	template <usize N>
	struct Rotator {};

	template <>
	struct Rotator<2> {
	using V = ivec2;

	// permutes rotations of src
	template <typename T, usize S>
	static inline std::array<std::array<T, S * S>, 4> permute(
	const std::array<T, S * S> &src) {
	std::array<std::array<T, S * S>, 4> dst;
	dst[0] = src;
	dst[1] = rotate_ccw<T, S>(dst[0]);
	dst[2] = rotate_ccw<T, S>(dst[1]);
	dst[3] = rotate_ccw<T, S>(dst[2]);
	return dst;
	}

	private:
	template <typename T, usize S>
	static inline std::array<T, S * S> rotate(
	const std::array<T, S * S> &src) {
	std::array<T, S * S> dst;
	for (usize i = 0; i < S; i++) {
	for (usize j = 0; j < S; j++) {
	ndarray::at(V(S), &dst[0], { i, j }) =
	ndarray::at(V(S), &src[0], { S - j - 1, i });
	}
	}
	return dst;
	}
	};

	// permutes n-dimensional reflections of an array
	template <usize N>
	struct Reflector {
	using V = math::vec<N, int, math::defaultp>;

	template <typename T, usize S, usize VOL = math::cexp::pow(S, N)>
	static inline auto permute(
	const std::array<T, S * S> &src) {
	std::array<std::array<T, VOL>, math::cexp::pow(2, N)> dst;

	// iterate over possible permutations of axes
	usize i = 0;
	ndarray::each(
	V(2),
	[&](const V &v) {
	const auto which = math::vec<N, bool, math::defaultp>(v);
	dst[i] = src;
	for (usize j = 0; j < N; j++) {
	if (which[j]) {
	dst[i] = reflect_axis<T, S>(dst[i], j);
	}
	}
	i++;
	});

	return dst;
	}

	private:
	template <typename T, usize S, usize VOL = math::cexp::pow(S, N)>
	static inline auto reflect_axis(
	const std::array<T, VOL> &src, usize axis) {
	std::array<T, VOL> dst;
	ndarray::each(
	V(S),
	[&](const V &v) {
	V u = v;
	u[axis] = S - u[axis] - 1;
	ndarray::at(V(S), &dst[0], v) =
	ndarray::at(V(S), &src[0], u);
	});
	return dst;
	}
	};

	// neighbors in N dimensions
	template <usize N>
	struct Neighbors {};

	template <>
	struct Neighbors<2> {
	static constexpr std::array<ivec2, 4> neighbors =
	{
	ivec2(-1, 0),
	ivec2(1, 0),
	ivec2(0, -1),
	ivec2(0, 1),
	};
	};

	template <>
	struct Neighbors<3> {
	static constexpr std::array<ivec3, 9> neighbors =
	{
	ivec3(-1, 0, 0),
	ivec3(1, 0, 0),
	ivec3(0, -1, 0),
	ivec3(0, 1, 0),
	ivec3(0, 0, -1),
	ivec3(0, 0, 1),
	};
	};

	// implements n-dimensional wave function collapse with the "overlapping model"
	// T: pattern data type
	// N: number of dimension
	// S: size of patterns (MUST BE ODD!)
	// D: bitset size
	// V (defaulted): ivecN
	template <
	typename T,
	usize N,
	usize S,
	usize D,
	typename V = math::vec<N, int, math::defaultp>>
	requires ((S % 2) == 1)
	struct WFC {
	// forward declarations
	struct Pattern;
	struct Element;
	struct Wave;

	// function which takes current wave and returns next cell to collapse
	using NextCellFn = std::function<Element&(Wave&)>;

	// function which chooses the pattern to collapse a cell to
	using PatternFn = std::function<usize(Wave&, Element&)>;

	// function used for optional callbacks during collapse process
	using CallbackFn = std::function<void(const Wave&)>;

	// S ** N (volume of one pattern)
	static constexpr auto VOL = math::cexp::pow(S, N);

	// pattern derived from input data
	struct Pattern {
	// unique pattern ID
	usize id;

	// normalized frequency of this pattern in input data
	f32 frequency;

	// value at the center of this pattern
	T value;

	// pattern data
	std::array<T, VOL> data;

	// valid neighbors on each side of this pattern
	std::array<Bitset<D>, 2 * N> valid;

	explicit Pattern(const std::array<T, VOL> &data)
	: data(data),
	value(ndarray::at(V(S), &data[0], V(S) / 2)) {}

	// NOTE: comparison is on data only
	auto operator<=>(const Pattern &other) const {
	return (*this) == other ?
	std::strong_ordering::equal
	: this->hash() <=> other.hash();
	}

	// NOTE: equality comparison is on data only
	bool operator==(const Pattern &other) const {
	return this->hash() == other.hash()
	&& this->data == other.data;
	}

	inline u64 hash() const {
	if (this->_hash) {
	return this->_hash;
	}

	u64 v = 0x12345;
	for (const auto &x : this->data) {
	v ^= x + 0x9e3779b9 + (v << 6) + (v >> 2);
	}
	return this->_hash = v;
	}

	private:
	// stored hash, only calculated once
	mutable u64 _hash = 0;
	};

	// wave element
	struct Element {
	// position in output space
	V pos;

	// coefficient, marked bits are VALID choices
	Bitset<D> c;

	// number of valid bits remaining
	usize popcnt;

	// value post-collapse (std::nullopt if not collapsed)
	std::optional<T> value = std::nullopt;

	// memoized entropy values
	f32 sum_weights = 0.0f;
	f32 sum_weight_log_weights = 0.0f;
	f32 entropy = 0.0f;

	// intialize entropy values, coefficient
	void init(const Wave &w, const Bitset<D> &mask) {
	this->c = mask;
	this->popcnt = this->c.popcnt();

	for (
	auto it = this->c.begin_on();
	it != this->c.end_on();
	it++) {
	const auto weight = w.wfc.patterns[*it].frequency;
	this->sum_weights += weight;
	this->sum_weight_log_weights += weight * std::log(weight);
	}
	}

	// applies a mask to this element's coefficient, updating memoized
	// entropy values
	// returns false on failure (0 popcnt/contradiction)
	bool apply(const Wave &w, const Bitset<D> &mask) {
	// get bits which are on in the current coefficient but off in the
	// new mask. exit early if nothing changes.
	const auto diff = this->c & mask;
	if (diff.popcnt() == 0) {
	return true;
	}

	this->c &= mask;

	for (
	auto it = diff.begin_on();
	it != diff.end_on();
	it++) {
	const auto weight = w.wfc.patterns[*it].frequency;
	this->sum_weights -= weight;
	this->sum_weight_log_weights -= weight * std::log(weight);
	}

	this->entropy =
	std::log(this->sum_weights)
	- (this->sum_weight_log_weights / this->sum_weights);
	this->popcnt = this->c.popcnt();
	return this->popcnt != 0;
	}

	// collapse this element to pattern n
	bool collapse(usize n, const T &value) {
	if (!this->c[n]) {
	return false;
	}

	this->value = value;
	this->c.reset();
	this->c.set(n);
	this->popcnt = 1;
	this->entropy = 0.0f;
	this->sum_weights = 0.0f;
	this->sum_weight_log_weights = 0.0f;
	return true;
	}

	// returns true if this element is collapsed
	bool collapsed() const {
	return static_cast<bool>(this->value);
	}
	};

	// wave to collapse
	struct Wave {
	const WFC &wfc;

	// size of output wave
	V size_wave;

	// output wave elements
	std::vector<Element> wave;

	// optional preset values
	const std::optional<T> *preset;

	// total number of collapsed elements
	usize num_collapsed = 0;

	Wave(
	const WFC &wfc,
	const V &size_wave,
	const std::optional<T> *preset = nullptr)
	: wfc(wfc),
	size_wave(size_wave),
	preset(preset) {}

	// collapses the specified wave element to the only remaining
	// possibility, or n if specified
	result::Result<void, V> collapse(
	Element &e,
	usize n = std::numeric_limits<usize>::max()) {
	n = (n == std::numeric_limits<usize>::max()) ? e.c.nth_set(0) : n;

	const auto &p = this->wfc.patterns[n];
	if (!e.collapse(n, p.value)) {
	return result::Err(e.pos);
	}

	this->num_collapsed++;
	return result::Ok();
	}

	// observe specified wave element, collapsing it to one of its possible
	// values
	result::Result<void, V> observe(Element &e) {
	const auto n = this->wfc.pattern_fn(*this, e);
	return this->collapse(e, n);
	}

	// propagate the current value of the specified wave element
	// returns Ok on success, erroneous position V on failure
	// (contradiction)
	result::Result<void, V> propagate(Element &to_propagate) {
	// DFS elements to update
	std::stack<Element*> es;
	es.push(to_propagate);

	// if not std::nullopt, there is an unresolvable contradiction in
	// the wave
	std::optional<V> contradiction = std::nullopt;

	// propagate each stack entry to neighbors
	while (!es.empty()) {
	auto &e = *es.top();
	es.pop();

	// get only valid non-collapsed neighbors
	std::array<Element, N 2> neighbors;
	for (usize i = 0; i < N * 2; i++) {
	neighbors[i] = nullptr;

	const auto &n = Neighbors<N>::neighbors[i];
	const auto pos_n = e.pos + n;

	if (!ndarray::in_bounds(this->size_wave, pos_n)) {
	continue;
	}

	auto &e_n =
	ndarray::at(
	this->size_wave,
	&this->wave[0],
	pos_n);

	if (e_n.collapsed()) {
	continue;
	}

	neighbors[i] = &e_n;
	}

	// compute superpatterns for valid neighbors
	std::array<Bitset<D>, N * 2> neighbor_patterns;
	for (
	auto it = e.c.begin_on();
	it != e.c.end_on();
	it++) {
	const auto &p = this->wfc.patterns[*it];
	for (usize i = 0; i < N * 2; i++) {
	if (neighbors[i]) {
	neighbor_patterns[i] \|= p.valid[i];
	}
	}
	}

	// apply superpatterns
	for (usize i = 0; i < N * 2; i++) {
	if (!neighbors[i]) {
	continue;
	}

	auto &e_n = *neighbors[i];
	const auto popcnt_old = e_n.popcnt;
	e_n.apply(*this, neighbor_patterns[i]);

	if (e_n.popcnt != popcnt_old) {
	if (e_n.popcnt == 0) {
	// zero popcount = failure/contradiction
	this->collapse(e_n, 0);
	contradiction = e_n.pos;
	break;
	} else if (e_n.popcnt == 1) {
	// one popcount = one remaining possibility/collapse
	auto res = this->collapse(e_n);
	if (res.isErr()) {
	return res;
	}
	}

	// propagate changed value
	es.push(&e_n);
	}
	}
	}

	if (contradiction) {
	return result::Err(*contradiction);
	}

	if (this->wfc.on_propagate) {
	(this->wfc.on_propagate)(this);
	}

	return result::Ok();
	}

	// collapse the wave
	result::Result<void, V> collapse() {
	// initialize the wave, all patterns are valid for each element
	this->wave = std::vector<Element>(math::prod(this->size_wave));

	// safe to parallelize, only one element modified at a time
	#pragma omp parallel for
	for (usize i = 0; i < this->wave.size(); i++) {
	auto &e = this->wave[i];
	e.pos = ndarray::unravel_index(this->size_wave, i);
	e.init(*this, this->wfc.mask_used);
	}

	// load preset values if present
	if (this->preset) {
	ndarray::each(
	this->size_wave,
	[&](const V &pos) {
	const auto &p =
	ndarray::at(this->size_wave, &this->preset[0], pos);
	if (p) {
	auto &e =
	ndarray::at(
	this->size_wave,
	&this->wave[0],
	pos);
	e.value = *p;
	}
	});
	}

	// collapse the wave
	while (this->num_collapsed != this->wave.size()) {
	auto &e = this->wfc.next_cell_fn(*this);

	const auto res_observe = this->observe(e);
	if (res_observe.isErr()) {
	return res_observe;
	}

	const auto res_prop = this->propagate(e);
	if (res_prop.isErr()) {
	return res_prop;
	}
	}

	// success!
	return result::Ok();
	}
	};

	// size of input data
	V size_in;

	// input data
	const T *in;

	// possible patterns
	std::vector<Pattern> patterns;

	// function to select pattern to collapse to
	PatternFn pattern_fn;

	// function to select next cell to collapse
	NextCellFn next_cell_fn;

	// behavior of patterns on borders
	BorderBehavior border_behavior;

	// mask of used bits of coefficient bitsets
	// bits are also zeroed for disallowed patterns (patterns without valid
	// neighbors)
	Bitset<D> mask_used;

	// random generator
	Rand rand;

	// options
	usize flags;

	// optional callback to be called after propagation of each collapsed wave
	// element
	std::optional<CallbackFn> on_propagate = std::nullopt;

	explicit WFC(
	const V &size_in,
	const T *in,
	PatternFunction pattern_function,
	NextCellFunction next_cell_function,
	BorderBehavior border_behavior,
	Rand &&rand,
	usize flags)
	: size_in(size_in),
	in(in),
	border_behavior(border_behavior),
	rand(std::move(rand)),
	flags(flags) {
	// get pattern data from input at specified location
	// returns std::nullopt if pattern is made illegal by border behavior
	const auto data_at =
	[&](const V &center) -> std::optional<std::array<T, VOL>> {
	const auto base = center - (V(S) / 2);

	if (border_behavior == BorderBehavior::EXCLUDE) {
	if (!ndarray::in_bounds(size_in, base)
	\|\| !ndarray::in_bounds(size_in, base + V(S) - 1)) {
	return std::nullopt;
	}
	}

	std::array<T, VOL> dst;
	ndarray::each(
	V(S),
	[&](const V &offset) {
	std::optional<T> override = std::nullopt;
	V pos = base + offset;

	switch (border_behavior) {
	case BorderBehavior::EXCLUDE:
	break;
	case BorderBehavior::ZERO:
	if (!ndarray::in_bounds(size_in, pos)) {
	override = T(0);
	}
	break;
	case BorderBehavior::CLAMP:
	pos = math::clamp(pos, V(0), size_in - V(1));
	break;
	case BorderBehavior::WRAP:
	pos = ((pos % size_in) + size_in) % size_in;
	break;
	};

	ndarray::at(V(S), &dst[0], offset) =
	override ?
	*override
	: ndarray::at(size_in, in, pos);
	});
	return std::make_optional(dst);
	};

	// slide along input data creating patterns
	ndarray::each(
	size_in,
	[&](const V &p) {
	if (const auto data = data_at(p)) {
	this->patterns.emplace_back(Pattern(data));
	}
	});

	// calculate per-pattern frequencies (normalization occurs later once
	// patterns have been deduplicated)
	std::unordered_map<u64, usize> pattern_hash_to_freq;
	for (const auto &p : this->patterns) {
	pattern_hash_to_freq[p.hash()]++;
	}

	// assign base frequencies (again, not normalized)
	for (auto &p : this->patterns) {
	p.frequency = static_cast<f32>(pattern_hash_to_freq[p.hash()]);
	}

	// removes duplicate patterns from this->patterns
	const auto deduplicate =
	[&]() {
	std::sort(this->patterns.begin(), this->patterns.end());
	this->patterns.erase(
	std::unique(this->patterns.begin(), this->patterns.end()),
	this->patterns.end());
	};

	// remove duplicate patterns
	deduplicate();

	// add rotations, reflections
	const auto base_patterns = this->patterns;
	this->patterns.clear();
	for (const auto &p : base_patterns) {
	this->patterns.push_back(p);

	if (flags & FLAGS_ROTATE) {
	ASSERT(N == 2, "can only rotate patterns in 2 dimensions");
	if constexpr (N == 2) {
	const auto permutations =
	Rotator<N>::template permute<T, S>(p.data);
	for (usize i = 1; i < permutations.size(); i++) {
	auto &q =
	this->patterns.emplace_back(
	pattern(permutations[i]));
	q.frequency = p.frequency;
	}
	}
	}
	if (flags & FLAGS_REFLECT) {
	const auto permutations =
	Reflector<N>::template permute<T, S>(p.data);
	for (usize i = 1; i < permutations.size(); i++) {
	auto &q =
	this->patterns.emplace_back(
	pattern(permutations[i]));
	q.frequency = p.frequency;
	}
	}
	}

	// deduplicate again to remove duplicates added by reflection/rotation
	deduplicate();

	// normalize pattern frequencies, pattern set will no longer change
	f32 frequency_total = 0.0f;
	for (const auto &p : this->patterns) {
	frequency_total += p.frequency;
	}

	for (auto &p : this->patterns) {
	p.frequency /= frequency_total;
	}

	// compute used mask, assign IDs
	this->mask_used.reset();
	for (usize i = 0; i < this->patterns.size(); i++) {
	this->patterns[i].id = i;
	this->mask_used.set(i);
	}

	// compuate valid patterns around each pattern (adjacency)
	// check against all overlapping slots for each pattern pair (p, q) for
	// every side
	//
	// safe to parallelize on p because only p is being written to and not
	// read from :)
	#pragma omp parallel for
	for (auto &p : this->patterns) {
	for (const auto &q : this->patterns) {
	for (usize i = 0; i < 2 * N; i++) {
	const auto &n = Neighbors<N>::neighbors[i];
	bool valid = true;
	ndarray::each(
	V(S),
	[&](const auto &offset_q) {
	if (!valid) {
	return;
	}

	// compute offset into p's data
	const auto offset_p = n + offset_q;
	if (!ndarray::in_bounds(V(S), offset_p)) {
	return;
	}

	const auto
	v_p =
	ndarray::at(V(S), &p.data[0], offset_p),
	v_q =
	ndarray::at(V(S), &q.data[0], offset_q);

	// data must be equal at each offset for patterns to
	// match (be valid neighbors)
	if (v_p != v_q) {
	valid = false;
	}
	});

	if (valid) {
	p.valid[i].set(q.id);
	}
	}
	}
	}

	switch (pattern_function) {
	case PatternFunction::WEIGHTED:
	this->pattern_fn = this->pattern_weighted();
	break;
	}

	switch (next_cell_fn) {
	case NextCellFunction::MIN_ENTROPY:
	this->next_cell_fn = this->next_cell_min_entropy();
	break;
	}
	}

	// collapses a wave into the specified output
	// returns true on success
	bool collapse(
	const V &size_out,
	T *out,
	const std::optional<T> *preset = nullptr) const {
	auto w = Wave(*this, size_out, preset);
	const auto res = w.collapse();
	if (res.isErr()) {
	return false;
	}

	// map elements to output
	ndarray::each(
	w.size_wave,
	[&](const V &pos) {
	ndarray::at(size_out, out, pos) =
	*ndarray::at(size_out, &w.wave, pos).value;
	});

	return true;
	}

	// returns a function which selects a pattern based on the distribution of
	// patterns in the input data
	PatternFn pattern_weighted() {
	return [&](Wave &w, Element &e) -> usize {
	f32 sum_cs = 0.0f;
	std::vector<f32> cs(this->patterns.size());
	for (
	auto it = e.c.begin_on();
	it != e.c.end_on();
	it++) {
	cs[it] = this->patterns[it].frequency;
	sum_cs += cs[*it];
	}

	const auto r = this->rand.template next<f32>(0.0f, sum_cs);
	f32 acc = 0.0f;
	for (usize i = 0; i < cs.size(); i++) {
	acc += cs[i];

	if (acc >= r) {
	return i;
	}
	}

	ASSERT(false, "failed to select pattern for {}", e.pos);
	return e.c.nth_set(0);
	};
	}

	// returns a function which selects the next cell based on finding the cell
	// with the minimum entropy in those remaining
	NextCellFn next_cell_min_entropy() {
	return [&](Wave &w) -> Element& {
	f32 min = 1e4;
	Element *argmin = nullptr;

	// min entropy with a bit of noise
	for (auto &e : w.wave) {
	if (!e.collapsed()
	&& e.entropy < min
	&& e.entropy + rand.next<f32>(0.0f, 1e-6) < min) {
	argmin = &e;
	}
	}

	ASSERT(argmin);
	return *argmin;
	};
	}
	};
	}