// based on KNLMeansCL by Khanattila

#include <algorithm>
#include <array>
#include <cstdint>
#include <cstdlib>
#include <cstring>
#include <limits>
#include <memory>
#include <mutex>
#include <shared_mutex>
#include <stdexcept>
#include <thread>
#include <type_traits>
#include <unordered_map>

#include <VapourSynth.h>
#include <VSHelper.h>

#include <nlm_ispc.h> // generated by the ispc compiler

#include <config.h> // generated by cmake and git

namespace {
enum struct ChannelMode { Y, UV, YUV, RGB };

struct NLMData {
    VSNodeRef * node; // clip
    const VSVideoInfo *vi;
    int d;
    int a;
    int s;
    float h;
    ChannelMode channels;
    decltype(&ispc::nlmVerticalWelsch) nlm_vertical; // wmode
    float wref;
    VSNodeRef * ref_node; // rclip

    // run-time resources
    std::shared_mutex workspaces_lock;
    std::unordered_map<std::thread::id, float *> workspaces;
};
}

template <typename T>
static inline auto castVoidPtr(T * p) noexcept {
    if constexpr (std::is_const_v<T>) {
        return reinterpret_cast<const void *>(p);
    } else {
        return reinterpret_cast<void *>(p);
    }
}

template <typename T1, typename T2>
static inline std::array<T1 *, 3> castPtrs(std::array<T2 *, 3> ptrs) {
    return {
        (T1 *) ptrs[0],
        (T1 *) ptrs[1],
        (T1 *) ptrs[2],
    };
}

template <typename T>
static inline constexpr T square(T x) noexcept {
    return x * x;
}

// T: (const) VSFrameRef
template <typename T>
static inline auto getPtrs(
    T * frame,
    ChannelMode channels,
    const VSAPI * vsapi
) noexcept {

    using value_type = std::conditional_t<std::is_const_v<T>, const void, void>;

    std::array<value_type *, 3> ptrs {};

    auto get_ptr = [frame, vsapi](int plane) {
        if constexpr (std::is_const_v<T>) {
            return castVoidPtr(vsapi->getReadPtr(frame, plane));
        } else {
            return castVoidPtr(vsapi->getWritePtr(frame, plane));
        }
    };

    switch (channels) {
        case ChannelMode::Y:
            ptrs[0] = get_ptr(0);
            break;
        case ChannelMode::UV:
            ptrs[1] = get_ptr(1);
            ptrs[2] = get_ptr(2);
            break;
        case ChannelMode::YUV:
        case ChannelMode::RGB:
            ptrs[0] = get_ptr(0);
            ptrs[1] = get_ptr(1);
            ptrs[2] = get_ptr(2);
            break;
    }

    return ptrs;
}

static void VS_CC nlmInit(
    VSMap * in,
    VSMap * out,
    void ** instanceData,
    VSNode * node,
    VSCore * core,
    const VSAPI * vsapi
) noexcept {

    const auto * d = reinterpret_cast<const NLMData *>(*instanceData);
    vsapi->setVideoInfo(vsapi->getVideoInfo(d->node), 1, node);
}

static inline void nlmDistanceDispatch_f32(
    float * temp0,
    std::array<const float *, 3> centerp,
    std::array<const float *, 3> neighborp,
    int offset_x,
    int offset_y,
    int width,
    int height,
    int stride,
    ChannelMode channels
) noexcept {

    switch (channels) {
        case ChannelMode::Y:
            ispc::nlmDistanceLuma_f32(
                temp0,
                centerp[0],
                neighborp[0],
                offset_x, offset_y,
                width, height, stride
            );
            break;
        case ChannelMode::UV:
            ispc::nlmDistanceChroma_f32(
                temp0,
                centerp[1], centerp[2],
                neighborp[1], neighborp[2],
                offset_x, offset_y,
                width, height, stride
            );
            break;
        case ChannelMode::YUV:
            ispc::nlmDistanceYUV_f32(
                temp0,
                centerp[0], centerp[1], centerp[2],
                neighborp[0], neighborp[1], neighborp[2],
                offset_x, offset_y,
                width, height, stride
            );
            break;
        case ChannelMode::RGB:
            ispc::nlmDistanceRGB_f32(
                temp0,
                centerp[0], centerp[1], centerp[2],
                neighborp[0], neighborp[1], neighborp[2],
                offset_x, offset_y,
                width, height, stride
            );
            break;
    }
}

static inline void nlmDistanceDispatch_u8(
    float * temp0,
    std::array<const uint8_t *, 3> centerp,
    std::array<const uint8_t *, 3> neighborp,
    int offset_x,
    int offset_y,
    int width,
    int height,
    int stride,
    ChannelMode channels,
    float inv_divisor
) noexcept {

    switch (channels) {
        case ChannelMode::Y:
            ispc::nlmDistanceLuma_u8(
                temp0,
                centerp[0],
                neighborp[0],
                offset_x, offset_y,
                width, height, stride,
                inv_divisor
            );
            break;
        case ChannelMode::UV:
            ispc::nlmDistanceChroma_u8(
                temp0,
                centerp[1], centerp[2],
                neighborp[1], neighborp[2],
                offset_x, offset_y,
                width, height, stride,
                inv_divisor
            );
            break;
        case ChannelMode::YUV:
            ispc::nlmDistanceYUV_u8(
                temp0,
                centerp[0], centerp[1], centerp[2],
                neighborp[0], neighborp[1], neighborp[2],
                offset_x, offset_y,
                width, height, stride,
                inv_divisor
            );
            break;
        case ChannelMode::RGB:
            ispc::nlmDistanceRGB_u8(
                temp0,
                centerp[0], centerp[1], centerp[2],
                neighborp[0], neighborp[1], neighborp[2],
                offset_x, offset_y,
                width, height, stride,
                inv_divisor
            );
            break;
    }
}

static inline void nlmDistanceDispatch_u16(
    float * temp0,
    std::array<const uint16_t *, 3> centerp,
    std::array<const uint16_t *, 3> neighborp,
    int offset_x,
    int offset_y,
    int width,
    int height,
    int stride,
    ChannelMode channels,
    float inv_divisor
) noexcept {

    switch (channels) {
        case ChannelMode::Y:
            ispc::nlmDistanceLuma_u16(
                temp0,
                centerp[0],
                neighborp[0],
                offset_x, offset_y,
                width, height, stride,
                inv_divisor
            );
            break;
        case ChannelMode::UV:
            ispc::nlmDistanceChroma_u16(
                temp0,
                centerp[1], centerp[2],
                neighborp[1], neighborp[2],
                offset_x, offset_y,
                width, height, stride,
                inv_divisor
            );
            break;
        case ChannelMode::YUV:
            ispc::nlmDistanceYUV_u16(
                temp0,
                centerp[0], centerp[1], centerp[2],
                neighborp[0], neighborp[1], neighborp[2],
                offset_x, offset_y,
                width, height, stride,
                inv_divisor
            );
            break;
        case ChannelMode::RGB:
            ispc::nlmDistanceRGB_u16(
                temp0,
                centerp[0], centerp[1], centerp[2],
                neighborp[0], neighborp[1], neighborp[2],
                offset_x, offset_y,
                width, height, stride,
                inv_divisor
            );
            break;
    }
}

static inline void nlmDistance(
    float * temp0,
    std::array<const void *, 3> centerp,
    std::array<const void *, 3> neighborp,
    int offset_x,
    int offset_y,
    int width,
    int height,
    int stride,
    ChannelMode channels,
    int bits
) noexcept {

    if (bits == 32) {
        nlmDistanceDispatch_f32(
            temp0,
            castPtrs<const float>(centerp), castPtrs<const float>(neighborp),
            offset_x, offset_y,
            width, height, stride, channels
        );
    } else if (bits <= 8) {
        float inv_divisor = 1.0f / ((1 << bits) - 1);
        nlmDistanceDispatch_u8(
            temp0,
            castPtrs<const uint8_t>(centerp), castPtrs<const uint8_t>(neighborp),
            offset_x, offset_y,
            width, height, stride, channels, inv_divisor
        );
    } else if (bits <= 16) {
        float inv_divisor = 1.0f / ((1 << bits) - 1);
        nlmDistanceDispatch_u16(
            temp0,
            castPtrs<const uint16_t>(centerp), castPtrs<const uint16_t>(neighborp),
            offset_x, offset_y,
            width, height, stride, channels, inv_divisor
        );
    } else {
        assert(false);
    }
}

static inline void nlmAccumulationDispatch_f32(
    float * weightp,
    std::array<float *, 3> wdstp,
    float * max_weightp,
    std::array<const float *, 3> srcp_bwd,
    std::array<const float *, 3> srcp_fwd,
    const float * temp_bwd,
    const float * temp_fwd,
    int offset_x,
    int offset_y,
    int width,
    int height,
    int stride,
    ChannelMode channels
) noexcept {

    switch (channels) {
        case ChannelMode::Y:
            ispc::nlmAccumulationCh1_f32(
                weightp, wdstp[0], max_weightp,
                srcp_bwd[0],
                srcp_fwd[0],
                temp_bwd, temp_fwd,
                offset_x, offset_y,
                width, height, stride
            );
            break;
        case ChannelMode::UV:
            ispc::nlmAccumulationCh2_f32(
                weightp, wdstp[0], wdstp[1], max_weightp,
                srcp_bwd[1], srcp_bwd[2],
                srcp_fwd[1], srcp_fwd[2],
                temp_bwd, temp_fwd,
                offset_x, offset_y,
                width, height, stride
            );
            break;
        case ChannelMode::YUV:
        case ChannelMode::RGB:
            ispc::nlmAccumulationCh3_f32(
                weightp, wdstp[0], wdstp[1], wdstp[2], max_weightp,
                srcp_bwd[0], srcp_bwd[1], srcp_bwd[2],
                srcp_fwd[0], srcp_fwd[1], srcp_fwd[2],
                temp_bwd, temp_fwd,
                offset_x, offset_y,
                width, height, stride
            );
            break;
    }
}

static inline void nlmAccumulationDispatch_u8(
    float * weightp,
    std::array<float *, 3> wdstp,
    float * max_weightp,
    std::array<const uint8_t *, 3> srcp_bwd,
    std::array<const uint8_t *, 3> srcp_fwd,
    const float * temp_bwd,
    const float * temp_fwd,
    int offset_x,
    int offset_y,
    int width,
    int height,
    int stride,
    ChannelMode channels
) noexcept {

    switch (channels) {
        case ChannelMode::Y:
            ispc::nlmAccumulationCh1_u8(
                weightp, wdstp[0], max_weightp,
                srcp_bwd[0],
                srcp_fwd[0],
                temp_bwd, temp_fwd,
                offset_x, offset_y,
                width, height, stride
            );
            break;
        case ChannelMode::UV:
            ispc::nlmAccumulationCh2_u8(
                weightp, wdstp[0], wdstp[1], max_weightp,
                srcp_bwd[1], srcp_bwd[2],
                srcp_fwd[1], srcp_fwd[2],
                temp_bwd, temp_fwd,
                offset_x, offset_y,
                width, height, stride
            );
            break;
        case ChannelMode::YUV:
        case ChannelMode::RGB:
            ispc::nlmAccumulationCh3_u8(
                weightp, wdstp[0], wdstp[1], wdstp[2], max_weightp,
                srcp_bwd[0], srcp_bwd[1], srcp_bwd[2],
                srcp_fwd[0], srcp_fwd[1], srcp_fwd[2],
                temp_bwd, temp_fwd,
                offset_x, offset_y,
                width, height, stride
            );
            break;
    }
}

static inline void nlmAccumulationDispatch_u16(
    float * weightp,
    std::array<float *, 3> wdstp,
    float * max_weightp,
    std::array<const uint16_t *, 3> srcp_bwd,
    std::array<const uint16_t *, 3> srcp_fwd,
    const float * temp_bwd,
    const float * temp_fwd,
    int offset_x,
    int offset_y,
    int width,
    int height,
    int stride,
    ChannelMode channels
) noexcept {

    switch (channels) {
        case ChannelMode::Y:
            ispc::nlmAccumulationCh1_u16(
                weightp, wdstp[0], max_weightp,
                srcp_bwd[0],
                srcp_fwd[0],
                temp_bwd, temp_fwd,
                offset_x, offset_y,
                width, height, stride
            );
            break;
        case ChannelMode::UV:
            ispc::nlmAccumulationCh2_u16(
                weightp, wdstp[0], wdstp[1], max_weightp,
                srcp_bwd[1], srcp_bwd[2],
                srcp_fwd[1], srcp_fwd[2],
                temp_bwd, temp_fwd,
                offset_x, offset_y,
                width, height, stride
            );
            break;
        case ChannelMode::YUV:
        case ChannelMode::RGB:
            ispc::nlmAccumulationCh3_u16(
                weightp, wdstp[0], wdstp[1], wdstp[2], max_weightp,
                srcp_bwd[0], srcp_bwd[1], srcp_bwd[2],
                srcp_fwd[0], srcp_fwd[1], srcp_fwd[2],
                temp_bwd, temp_fwd,
                offset_x, offset_y,
                width, height, stride
            );
            break;
    }
}

static inline void nlmAccumulation(
    float * weightp,
    std::array<float *, 3> wdstp,
    float * max_weightp,
    std::array<const void *, 3> srcp_bwd,
    std::array<const void *, 3> srcp_fwd,
    const float * temp_bwd,
    const float * temp_fwd,
    int offset_x,
    int offset_y,
    int width,
    int height,
    int stride,
    ChannelMode channels,
    int bits
) noexcept {

    if (bits == 32) {
        nlmAccumulationDispatch_f32(
            weightp, wdstp, max_weightp,
            castPtrs<const float>(srcp_bwd), castPtrs<const float>(srcp_fwd), temp_bwd, temp_fwd,
            offset_x, offset_y, width, height, stride, channels
        );
    } else if (bits <= 8) {
        nlmAccumulationDispatch_u8(
            weightp, wdstp, max_weightp,
            castPtrs<const uint8_t>(srcp_bwd), castPtrs<const uint8_t>(srcp_fwd), temp_bwd, temp_fwd,
            offset_x, offset_y, width, height, stride, channels
        );
    } else if (bits <= 16) {
        nlmAccumulationDispatch_u16(
            weightp, wdstp, max_weightp,
            castPtrs<const uint16_t>(srcp_bwd), castPtrs<const uint16_t>(srcp_fwd), temp_bwd, temp_fwd,
            offset_x, offset_y, width, height, stride, channels
        );
    } else {
        assert(false);
    }
}

static inline void nlmFinishDispatch_f32(
    std::array<float *, 3> dstp,
    std::array<const float *, 3> srcp,
    const float * weightp,
    std::array<float *, 3> wdstp,
    const float * max_weightp,
    float wref,
    int width,
    int height,
    int stride,
    ChannelMode channels
) noexcept {

    switch (channels) {
        case ChannelMode::Y:
            ispc::nlmFinishCh1_f32(
                dstp[0],
                srcp[0],
                weightp, wdstp[0],
                max_weightp, wref,
                width, height, stride
            );
            break;
        case ChannelMode::UV:
            ispc::nlmFinishCh2_f32(
                dstp[1], dstp[2],
                srcp[1], srcp[2],
                weightp, wdstp[0], wdstp[1],
                max_weightp, wref,
                width, height, stride
            );
            break;
        case ChannelMode::YUV:
        case ChannelMode::RGB:
            ispc::nlmFinishCh3_f32(
                dstp[0], dstp[1], dstp[2],
                srcp[0], srcp[1], srcp[2],
                weightp, wdstp[0], wdstp[1], wdstp[2],
                max_weightp, wref,
                width, height, stride
            );
            break;
    }
}

static inline void nlmFinishDispatch_u8(
    std::array<uint8_t *, 3> dstp,
    std::array<const uint8_t *, 3> srcp,
    const float * weightp,
    std::array<float *, 3> wdstp,
    const float * max_weightp,
    float wref,
    int width,
    int height,
    int stride,
    ChannelMode channels,
    int peak
) noexcept {

    switch (channels) {
        case ChannelMode::Y:
            ispc::nlmFinishCh1_u8(
                dstp[0],
                srcp[0],
                weightp, wdstp[0],
                max_weightp, wref,
                width, height, stride,
                peak
            );
            break;
        case ChannelMode::UV:
            ispc::nlmFinishCh2_u8(
                dstp[1], dstp[2],
                srcp[1], srcp[2],
                weightp, wdstp[0], wdstp[1],
                max_weightp, wref,
                width, height, stride,
                peak
            );
            break;
        case ChannelMode::YUV:
        case ChannelMode::RGB:
            ispc::nlmFinishCh3_u8(
                dstp[0], dstp[1], dstp[2],
                srcp[0], srcp[1], srcp[2],
                weightp, wdstp[0], wdstp[1], wdstp[2],
                max_weightp, wref,
                width, height, stride,
                peak
            );
            break;
    }
}

static inline void nlmFinishDispatch_u16(
    std::array<uint16_t *, 3> dstp,
    std::array<const uint16_t *, 3> srcp,
    const float * weightp,
    std::array<float *, 3> wdstp,
    const float * max_weightp,
    float wref,
    int width,
    int height,
    int stride,
    ChannelMode channels,
    int peak
) noexcept {

    switch (channels) {
        case ChannelMode::Y:
            ispc::nlmFinishCh1_u16(
                dstp[0],
                srcp[0],
                weightp, wdstp[0],
                max_weightp, wref,
                width, height, stride,
                peak
            );
            break;
        case ChannelMode::UV:
            ispc::nlmFinishCh2_u16(
                dstp[1], dstp[2],
                srcp[1], srcp[2],
                weightp, wdstp[0], wdstp[1],
                max_weightp, wref,
                width, height, stride,
                peak
            );
            break;
        case ChannelMode::YUV:
        case ChannelMode::RGB:
            ispc::nlmFinishCh3_u16(
                dstp[0], dstp[1], dstp[2],
                srcp[0], srcp[1], srcp[2],
                weightp, wdstp[0], wdstp[1], wdstp[2],
                max_weightp, wref,
                width, height, stride,
                peak
            );
            break;
    }
}

static inline void nlmFinish(
    std::array<void *, 3> dstp,
    std::array<const void *, 3> srcp,
    const float * weightp,
    std::array<float *, 3> wdstp,
    const float * max_weightp,
    float wref,
    int width,
    int height,
    int stride,
    ChannelMode channels,
    int bits
) noexcept {

    if (bits == 32) {
        nlmFinishDispatch_f32(
            castPtrs<float>(dstp), castPtrs<const float>(srcp),
            weightp, wdstp, max_weightp, wref, width, height, stride, channels
        );
    } else if (bits <= 8) {
        int peak = (1 << bits) - 1;
        nlmFinishDispatch_u8(
            castPtrs<uint8_t>(dstp), castPtrs<const uint8_t>(srcp),
            weightp, wdstp, max_weightp, wref, width, height, stride, channels, peak
        );
    } else if (bits <= 16) {
        int peak = (1 << bits) - 1;
        nlmFinishDispatch_u16(
            castPtrs<uint16_t>(dstp), castPtrs<const uint16_t>(srcp),
            weightp, wdstp, max_weightp, wref, width, height, stride, channels, peak
        );
    } else {
        assert(false);
    }
}

static const VSFrameRef *VS_CC nlmGetFrame(
    int n,
    int activationReason,
    void ** instanceData,
    void ** frameData,
    VSFrameContext * frameCtx,
    VSCore * core,
    const VSAPI * vsapi
) noexcept {

    auto * d = reinterpret_cast<NLMData *>(*instanceData);

    if (activationReason == arInitial) {
        int start = std::max(0, n - d->d);
        int end = std::min(n + d->d, d->vi->numFrames - 1);
        for (int i = start; i <= end; i++) {
            vsapi->requestFrameFilter(i, d->node, frameCtx);
            if (d->ref_node) {
                vsapi->requestFrameFilter(i, d->ref_node, frameCtx);
            }
        }
        return nullptr;
    } else if (activationReason != arAllFramesReady) {
        return nullptr;
    }

    // activationReason == arAllFramesReady

    int nlm_d = d->d;
    int nlm_a = d->a;
    int nlm_s = d->s;
    float nlm_h2_inv_norm = square(255.0f) / (3.0f * square(d->h) * square(2 * nlm_s + 1));
    float nlm_wref = d->wref;
    ChannelMode channels = d->channels;

    const auto & ref_node = d->ref_node ? d->ref_node : d->node;
    auto ref_frame = vsapi->getFrameFilter(n, ref_node, frameCtx);

    int bits = d->vi->format->bitsPerSample;
    int width, height, stride; // dimensions of the plane to be processed, not the video dimension
    if (channels == ChannelMode::UV) {
        width = d->vi->width >> d->vi->format->subSamplingW;
        height = d->vi->height >> d->vi->format->subSamplingH;
        stride = vsapi->getStride(ref_frame, 1) / d->vi->format->bytesPerSample;
    } else {
        width = d->vi->width;
        height = d->vi->height;
        stride = vsapi->getStride(ref_frame, 0) / d->vi->format->bytesPerSample;
    }

    int size = height * stride; // size of each plane in quad-bytes
    // number of input channels
    int num_input_channels = [channels]() {
        if (channels == ChannelMode::Y) {
            return 1;
        } else if (channels == ChannelMode::UV) {
            return 2;
        } else {
            // channels == ChannelMode::YUV || channels == ChannelMode::RGB
            return 3;
        }
    }();
    // size in quad-bytes: size * (4 + num_input_channels + (nlm_d != 0)) + width
    float * workspace;
    {
        auto thread_id = std::this_thread::get_id();
        d->workspaces_lock.lock_shared();
        bool init = true;
        try {
            const auto & const_workspaces = d->workspaces;
            workspace = const_workspaces.at(thread_id);
        } catch (const std::out_of_range &) {
            init = false;
        }
        d->workspaces_lock.unlock_shared();

        if (!init) {
            auto workspace_size = size * (4 + num_input_channels + (nlm_d != 0)) + width;
            auto workspace_bytes = workspace_size * sizeof(float);
            workspace = vs_aligned_malloc<float>(workspace_bytes, 256);

            if (!workspace) {
                vsapi->freeFrame(ref_frame);
                vsapi->setFilterError("nlm_ispc: malloc() failed", frameCtx);
                return nullptr;
            }

            std::lock_guard _ { d->workspaces_lock };
            d->workspaces.emplace(thread_id, workspace);
        }
    }

    // zero-initialize aggregation buffers
    std::memset(workspace, 0, (1 + num_input_channels) * size * sizeof(float));
    // stores the sum of weights of each pixel
    float * weightp = workspace;
    std::array<float *, 3> wdstp {
        // stores the weighted sum of pixel values of the first processed plane
        workspace + size,
        // stores the weighted sum of pixel values of the second processed plane
        num_input_channels <= 1 ? nullptr : workspace + 2 * size,
        // stores the weighted sum of pixel values of the third processed plane
        num_input_channels <= 2 ? nullptr : workspace + 3 * size
    };

    // stores the maximum weight encountered of each pixel
    float * max_weightp = workspace + (1 + num_input_channels) * size;
    for (int i = 0; i < size; i++) {
        max_weightp[i] = std::numeric_limits<float>::epsilon();
    }

    // temporary storage for the calculation of patch distances
    float * temp = workspace + (2 + num_input_channels) * size;
    float * temp_bwd = workspace + (3 + num_input_channels) * size;
    float * temp_fwd = nlm_d == 0 ? nullptr : workspace + (4 + num_input_channels) * size;

    // buffer for the vertical box filter during patch distance calculation
    // size in quad-bytes: width
    float * buffer = workspace + (4 + num_input_channels + (nlm_d != 0)) * size;

    std::array refp { getPtrs(ref_frame, channels, vsapi) };

    for (int i = -nlm_d; i <= 0; i++) {
        auto bwd_n = std::max(n + i, 0);
        auto fwd_n = std::min(n - i, d->vi->numFrames - 1);
        auto src_frame_bwd = vsapi->getFrameFilter(bwd_n, d->node, frameCtx);
        auto src_frame_fwd = vsapi->getFrameFilter(fwd_n, d->node, frameCtx);
        auto ref_frame_bwd = vsapi->getFrameFilter(bwd_n, ref_node, frameCtx);
        auto ref_frame_fwd = vsapi->getFrameFilter(fwd_n, ref_node, frameCtx);

        std::array srcp_bwd { getPtrs(src_frame_bwd, channels, vsapi) };
        std::array srcp_fwd { getPtrs(src_frame_fwd, channels, vsapi) };
        std::array refp_bwd { getPtrs(ref_frame_bwd, channels, vsapi) };
        std::array refp_fwd { getPtrs(ref_frame_fwd, channels, vsapi) };

        for (int offset_y = -nlm_a; offset_y <= nlm_a; offset_y++) {
            for (int offset_x = -nlm_a; offset_x <= nlm_a; offset_x++) {
                if (i * square(2 * nlm_a + 1) + offset_y * (2 * nlm_a + 1) + offset_x >= 0) {
                    continue;
                }

                nlmDistance(
                    temp_bwd,
                    refp, refp_bwd,
                    offset_x, offset_y, width, height, stride, channels, bits
                );

                ispc::nlmHorizontal(
                    temp,
                    temp_bwd,
                    nlm_s, width, height, stride
                );

                d->nlm_vertical(
                    temp_bwd,
                    temp,
                    nlm_s, nlm_h2_inv_norm, width, height, stride, buffer
                );

                // jump at the end of this basic block
                if (i == 0) {
                    // bwd == fwd
                    nlmAccumulation(
                        weightp, wdstp, max_weightp,
                        srcp_bwd, srcp_bwd, temp_bwd, temp_bwd,
                        offset_x, offset_y, width, height, stride, channels, bits
                    );
                    continue;
                }

                // i != 0
                nlmDistance(
                    temp_fwd,
                    refp_fwd, refp,
                    offset_x, offset_y, width, height, stride, channels, bits
                );

                ispc::nlmHorizontal(
                    temp,
                    temp_fwd,
                    nlm_s, width, height, stride
                );

                d->nlm_vertical(
                    temp_fwd,
                    temp,
                    nlm_s, nlm_h2_inv_norm, width, height, stride, buffer
                );

                nlmAccumulation(
                    weightp, wdstp, max_weightp,
                    srcp_bwd, srcp_fwd, temp_bwd, temp_fwd,
                    offset_x, offset_y, width, height, stride, channels, bits
                );
            }
        }

        vsapi->freeFrame(src_frame_fwd);
        vsapi->freeFrame(src_frame_bwd);
        vsapi->freeFrame(ref_frame_fwd);
        vsapi->freeFrame(ref_frame_bwd);
    }

    vsapi->freeFrame(ref_frame);

    auto src_frame = vsapi->getFrameFilter(n, d->node, frameCtx);
    std::array srcp { getPtrs(src_frame, channels, vsapi) };

    VSFrameRef * dst_frame;
    if (channels == ChannelMode::Y && d->vi->format->numPlanes > 1) {
        const VSFrameRef * fr[3] { nullptr, src_frame, src_frame };
        constexpr int pl[3] { 0, 1, 2 };
        dst_frame = vsapi->newVideoFrame2(d->vi->format, d->vi->width, d->vi->height, fr, pl, src_frame, core);
    } else if (channels == ChannelMode::UV && d->vi->format->numPlanes > 1) {
        const VSFrameRef * fr[3] { src_frame, nullptr, nullptr };
        constexpr int pl[3] { 0, 1, 2 };
        dst_frame = vsapi->newVideoFrame2(d->vi->format, d->vi->width, d->vi->height, fr, pl, src_frame, core);
    } else {
        dst_frame = vsapi->newVideoFrame(d->vi->format, d->vi->width, d->vi->height, src_frame, core);
    }
    std::array dstp { getPtrs(dst_frame, channels, vsapi) };

    nlmFinish(dstp, srcp, weightp, wdstp, max_weightp, nlm_wref, width, height, stride, channels, bits);

    vsapi->freeFrame(src_frame);

    return dst_frame;
}

static void VS_CC nlmFree(
    void * instanceData,
    VSCore * core,
    const VSAPI * vsapi
) noexcept {

    auto * d = reinterpret_cast<NLMData *>(instanceData);

    vsapi->freeNode(d->node);
    if (d->ref_node) {
        vsapi->freeNode(d->ref_node);
    }

    for (const auto & [_, ptr] : d->workspaces) {
        vs_aligned_free(ptr);
    }

    delete d;
}

static void VS_CC nlmCreate(
    const VSMap * in,
    VSMap * out,
    void * userData,
    VSCore * core,
    const VSAPI * vsapi
) noexcept {

    auto d = std::make_unique<NLMData>();

    d->node = vsapi->propGetNode(in, "clip", 0, nullptr);
    d->vi = vsapi->getVideoInfo(d->node);

    auto set_error = [vsapi, out, &d](const char * error_message) -> void {
        vsapi->setError(out, error_message);
        vsapi->freeNode(d->node);
    };

    if ((d->vi->format->sampleType == stInteger && d->vi->format->bitsPerSample > 16) ||
        (d->vi->format->sampleType == stFloat && d->vi->format->bitsPerSample != 32)
     ) {
        return set_error("only 1-16 bit integer or 32-bit float supported");
    }

    int err;

    d->d = int64ToIntS(vsapi->propGetInt(in, "d", 0, &err));
    if (err) {
        d->d = 1;
    }
    if (d->d < 0) {
        return set_error("\"d\" must be non-negative");
    }

    d->a = int64ToIntS(vsapi->propGetInt(in, "a", 0, &err));
    if (err) {
        d->a = 2;
    }
    if (d->a <= 0) {
        return set_error("\"a\" must be positive");
    }

    d->s = int64ToIntS(vsapi->propGetInt(in, "s", 0, &err));
    if (err) {
        d->s = 4;
    }
    if (d->s < 0) {
        return set_error("\"s\" must be non-negative");
    }

    d->h = static_cast<float>(vsapi->propGetFloat(in, "h", 0, &err));
    if (err) {
        d->h = 1.2f;
    }
    if (d->h <= 0.0f) {
        return set_error("\"h\" must be positive");
    }

    auto wmode = vsapi->propGetInt(in, "wmode", 0, &err);
    if (err) {
        wmode = 0;
    }
    if (wmode < 0 || wmode > 3) {
        return set_error("\"wmode\" must be 0, 1, 2 or 3");
    }
    decltype(d->nlm_vertical) nlmVerticalKernels[] {
        &ispc::nlmVerticalWelsch,
        &ispc::nlmVerticalBisquareA,
        &ispc::nlmVerticalBisquareB,
        &ispc::nlmVerticalBisquareC
    };
    d->nlm_vertical = nlmVerticalKernels[wmode];

    auto channels = vsapi->propGetData(in, "channels", 0, &err);
    if (err) {
        channels = "AUTO";
    }
    auto channels_len = std::strlen(channels);
    if (channels_len == 1 && *channels == 'Y') {
        d->channels = ChannelMode::Y;
    } else if (channels_len == 2 && std::strncmp(channels, "UV", 2) == 0) {
        d->channels = ChannelMode::UV;
    } else if (channels_len == 3 && std::strncmp(channels, "YUV", 3) == 0) {
        d->channels = ChannelMode::YUV;
    } else if (channels_len == 3 && std::strncmp(channels, "RGB", 3) == 0) {
        d->channels = ChannelMode::RGB;
    } else if (channels_len == 4 && std::strncmp(channels, "AUTO", 4) == 0) {
        if (d->vi->format->colorFamily == cmRGB) {
            d->channels = ChannelMode::RGB;
        } else {
            d->channels = ChannelMode::Y;
        }
    } else {
        return set_error("\"channels\" must be \"Y\", \"UV\', \"YUV\", \"RGB\" or \"AUTO\"");
    }

    if (d->channels == ChannelMode::Y) {
        if (d->vi->format->colorFamily != cmGray && d->vi->format->colorFamily != cmYUV) {
            return set_error("color family must be Gray or YUV for \"channels\" == \"Y\"");
        }
    } else if (d->channels == ChannelMode::UV) {
        if (d->vi->format->colorFamily != cmYUV) {
            return set_error("color family must be YUV for \"channels\" == \"UV\"");
        }
    } else if (d->channels == ChannelMode::YUV) {
        if (d->vi->format->colorFamily != cmYUV || d->vi->format->subSamplingW || d->vi->format->subSamplingH) {
            return set_error("color family must be YUV444 for \"channels\" == \"YUV\"");
        }
    } else if (d->channels == ChannelMode::RGB) {
        if (d->vi->format->colorFamily != cmRGB) {
            return set_error("color family must be RGB for \"channels\" == \"RGB\"");
        }
    }

    d->wref = static_cast<float>(vsapi->propGetFloat(in, "wref", 0, &err));
    if (err) {
        d->wref = 1.0f;
    }

    d->ref_node = vsapi->propGetNode(in, "rclip", 0, &err);
    if (err) {
        d->ref_node = nullptr;
    }
    if (d->ref_node) {
        const auto ref_vi = vsapi->getVideoInfo(d->ref_node);
        if (!isSameFormat(d->vi, ref_vi) || d->vi->numFrames != ref_vi->numFrames) {
            vsapi->freeNode(d->ref_node);
            return set_error("\"rclip\" must be of the same format as \"clip\"");
        }
    }

    VSCoreInfo core_info;
    vsapi->getCoreInfo2(core, &core_info);
    d->workspaces.reserve(core_info.numThreads);

    vsapi->createFilter(
        in, out,
        "NLMeans", nlmInit, nlmGetFrame, nlmFree,
        fmParallel, 0, d.release(), core
    );
}

VS_EXTERNAL_API(void) VapourSynthPluginInit(
    VSConfigPlugin configFunc,
    VSRegisterFunction registerFunc,
    VSPlugin * plugin
) noexcept {

    configFunc(
        "io.github.amusementclub.vs-nlm-ispc",
        "nlm_ispc",
        "Non-local means denoise filter implemented in ISPC",
        VAPOURSYNTH_API_VERSION, 1, plugin
    );

    registerFunc(
        "NLMeans",
        "clip:clip;"
        "d:int:opt;"
        "a:int:opt;"
        "s:int:opt;"
        "h:float:opt;"
        "channels:data:opt;"
        "wmode:int:opt;"
        "wref:float:opt;"
        "rclip:clip:opt;",
        nlmCreate,
        nullptr, plugin
    );

    auto getVersion = [](const VSMap *, VSMap * out, void *, VSCore *, const VSAPI *vsapi) {
        vsapi->propSetData(out, "version", VERSION, -1, paReplace);
    };
    registerFunc("Version", "", getVersion, nullptr, plugin);
}
