#include <algorithm>
#include <array>
#include <atomic>
#include <complex>
#include <concepts>
#include <cstddef>
#include <cstdio>
#include <cstdlib>
#include <cstring>
#include <memory>
#include <numbers>
#include <shared_mutex>
#include <sstream>
#include <stdexcept>
#include <string>
#include <thread>
#include <type_traits>
#include <unordered_map>
#include <utility>
#include <variant>
#include <vector>

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

#include <cuda.h>
#include <cufft.h>
#include <nvrtc.h>

#include "kernel.hpp"

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

// real/complex-input DFT
template <typename T, typename T_in>
    requires
        (std::is_same_v<T_in, T> || std::is_same_v<T_in, std::complex<T>>)
static void dft(
    std::complex<T> * VS_RESTRICT dst,
    const T_in * VS_RESTRICT src,
    int n,
    int stride
) {

    int out_num = std::is_floating_point_v<T_in> ? (n / 2 + 1) : n;
    for (int i = 0; i < out_num; i++) {
        std::complex<T> sum {};
        for (int j = 0; j < n; j++) {
            auto imag = -2 * i * j * std::numbers::pi_v<T> / n;
            auto weight = std::complex(std::cos(imag), std::sin(imag));
            sum += src[j * stride] * weight;
        }
        dst[i * stride] = sum;
    }
}

static bool success(CUresult result) {
    return result == CUDA_SUCCESS;
}
static bool success(cufftResult_t result) {
    return result == CUFFT_SUCCESS;
}
static bool success(nvrtcResult result) {
    return result == NVRTC_SUCCESS;
}

static const char * get_error(CUresult error) {
    const char * error_message;
    if (cuGetErrorString(error, &error_message) == CUDA_SUCCESS) [[likely]] {
        return error_message;
    }
    return "unknown cuda error";
}

static const char * get_error(cufftResult_t error) {
    switch (error) {
        case CUFFT_SUCCESS:
            return "success";
        case CUFFT_INVALID_PLAN:
            return "invalid plan handle";
        case CUFFT_ALLOC_FAILED:
            return "failed to allocate memory";
        case CUFFT_INVALID_TYPE:
            return "invalid type";
        case CUFFT_INVALID_VALUE:
            return "invalid value";
        case CUFFT_INTERNAL_ERROR:
            return "internal error";
        case CUFFT_EXEC_FAILED:
            return "execution failed";
        case CUFFT_SETUP_FAILED:
            return "the cuFFT library failed to initialize";
        case CUFFT_INVALID_SIZE:
            return "invalid transform size";
        case CUFFT_UNALIGNED_DATA:
            return "unaligned data";
#if CUFFT_VERSION < 12000
        case CUFFT_INCOMPLETE_PARAMETER_LIST:
            return "missing parameters in call";
#endif
        case CUFFT_INVALID_DEVICE:
            return "invalid device: execution of a plan was on different GPU than plan creation";
#if CUFFT_VERSION < 12000
        case CUFFT_PARSE_ERROR:
            return "internal plan database error";
#endif
        case CUFFT_NO_WORKSPACE:
            return "no workspace has been provided prior to plan execution";
        case CUFFT_NOT_IMPLEMENTED:
            return "functionality not implemented";
        case CUFFT_NOT_SUPPORTED:
            return "operation not supported";
#if CUFFT_VERSION >= 12000
        case CUFFT_MISSING_DEPENDENCY:
            return "missing dependency";
        case CUFFT_NVRTC_FAILURE:
            return "nvrtc failure";
        case CUFFT_NVJITLINK_FAILURE:
            return "nvjitlink failure";
        case CUFFT_NVSHMEM_FAILURE:
            return "nvshmem failure";
#endif
        default:
            return "unknown";
    }
}

static const char * get_error(nvrtcResult error) {
    return nvrtcGetErrorString(error);
}

#define showError(expr) show_error_impl(expr, # expr, __LINE__)
template <typename T>
static void show_error_impl(T result, const char * source, int line_no) {
    if (!success(result)) [[unlikely]] {
        std::fprintf(stderr, "[%d] %s failed: %s\n", line_no, source, get_error(result));
    }
}

#define checkError(expr) do {                                                       \
    if (auto result = expr; !success(result)) [[unlikely]] {                        \
        std::ostringstream error;                                                   \
        error << '[' << __LINE__ << "] '" # expr "' failed: " << get_error(result); \
        return set_error(error.str().c_str());                                      \
    }                                                                               \
} while (0)

static void cuStreamDestroyCustom(CUstream stream) {
    showError(cuStreamDestroy(stream));
}

static void cuEventDestroyCustom(CUevent event) {
    showError(cuEventDestroy(event));
}

static void cuMemFreeCustom(CUdeviceptr p) {
    showError(cuMemFree(p));
}

static void cuMemFreeHostCustom(void * p) {
    showError(cuMemFreeHost(p));
}

static void cuModuleUnloadCustom(CUmodule module) {
    showError(cuModuleUnload(module));
}

static void cufftDestroyCustom(cufftHandle handle) {
    showError(cufftDestroy(handle));
}

static void nvrtcDestroyProgramCustom(nvrtcProgram * program) {
    showError(nvrtcDestroyProgram(program));
}

struct context_releaser {
    bool * context_retained {};
    CUdevice device {};
    void release() {
        context_retained = nullptr;
    }
    ~context_releaser() {
        if (context_retained && *context_retained) {
            showError(cuDevicePrimaryCtxRelease(device));
        }
    }
};

struct context_popper {
    bool * context_pushed {};
    ~context_popper() {
        if (!context_pushed || *context_pushed) {
            showError(cuCtxPopCurrent(nullptr));
        }
    }
};

struct node_freer {
    const VSAPI * & vsapi;
    VSNodeRef * node {};
    void release() {
        node = nullptr;
    }
    ~node_freer() {
        if (node) {
            vsapi->freeNode(node);
        }
    }
};

template <typename T, auto deleter, bool unsafe=false>
    requires
        std::default_initializable<T> &&
        std::is_trivially_copy_assignable_v<T> &&
        std::convertible_to<T, bool> &&
        std::invocable<decltype(deleter), T> &&
        (std::is_pointer_v<T> || unsafe) // e.g. CUdeviceptr is not a pointer
struct Resource {
    T data;

    [[nodiscard]] constexpr Resource() noexcept = default;

    [[nodiscard]] constexpr Resource(T x) noexcept : data(x) {}

    [[nodiscard]] constexpr Resource(Resource&& other) noexcept
            : data(std::exchange(other.data, T{}))
    { }

    Resource& operator=(Resource&& other) noexcept {
        if (this == &other) return *this;
        deleter_(data);
        data = std::exchange(other.data, T{});
        return *this;
    }

    Resource operator=(Resource other) = delete;

    Resource(const Resource& other) = delete;

    constexpr operator T() const noexcept {
        return data;
    }

    constexpr auto deleter_(T x) noexcept {
        if (x) {
            deleter(x);
            x = T{};
        }
    }

    Resource& operator=(T x) noexcept {
        deleter_(data);
        data = x;
        return *this;
    }

    constexpr ~Resource() noexcept {
        deleter_(data);
    }
};

template <typename T>
static T square(const T & x) {
    return x * x;
}

static int calc_pad_size(int size, int block_size, int block_step) {
    return (
        size
        + ((size % block_size) ? block_size - size % block_size : 0)
        + std::max(block_size - block_step, block_step) * 2
    );
}

static int calc_pad_num(int size, int block_size, int block_step) {
    return (calc_pad_size(size, block_size, block_step) - block_size) / block_step + 1;
}

template <typename T>
static void reflection_padding_impl(
    T * VS_RESTRICT dst, // shape: (pad_height, pad_width)
    const T * VS_RESTRICT src, // shape: (height, stride)
    int width, int height, int stride,
    int block_size, int block_step
) {

    int pad_width = calc_pad_size(width, block_size, block_step);
    int pad_height = calc_pad_size(height, block_size, block_step);

    int offset_y = (pad_height - height) / 2;
    int offset_x = (pad_width - width) / 2;

    vs_bitblt(
        &dst[offset_y * pad_width + offset_x], pad_width * sizeof(T),
        src, stride * sizeof(T),
        width * sizeof(T), height
    );

    // copy left and right regions
    for (int y = offset_y; y < offset_y + height; y++) {
        auto dst_line = &dst[y * pad_width];

        for (int x = 0; x < offset_x; x++) {
            dst_line[x] = dst_line[offset_x * 2 - x];
        }

        for (int x = offset_x + width; x < pad_width; x++) {
            dst_line[x] = dst_line[2 * (offset_x + width) - 2 - x];
        }
    }

    // copy top region
    for (int y = 0; y < offset_y; y++) {
        std::memcpy(
            &dst[y * pad_width],
            &dst[(offset_y * 2 - y) * pad_width],
            pad_width * sizeof(T)
        );
    }

    // copy bottom region
    for (int y = offset_y + height; y < pad_height; y++) {
        std::memcpy(
            &dst[y * pad_width],
            &dst[(2 * (offset_y + height) - 2 - y) * pad_width],
            pad_width * sizeof(T)
        );
    }
}

static void reflection_padding(
    uint8_t * VS_RESTRICT dst, // shape: (pad_height, pad_width)
    const uint8_t * VS_RESTRICT src, // shape: (height, stride)
    int width, int height, int stride,
    int block_size, int block_step,
    int bytes_per_sample
) {

    if (bytes_per_sample == 1) {
        reflection_padding_impl(
            static_cast<uint8_t *>(dst),
            static_cast<const uint8_t *>(src),
            width, height, stride,
            block_size, block_step
        );
    } else if (bytes_per_sample == 2) {
        reflection_padding_impl(
            reinterpret_cast<uint16_t *>(dst),
            reinterpret_cast<const uint16_t *>(src),
            width, height, stride,
            block_size, block_step
        );
    } else if (bytes_per_sample == 4) {
        reflection_padding_impl(
            reinterpret_cast<uint32_t *>(dst),
            reinterpret_cast<const uint32_t *>(src),
            width, height, stride,
            block_size, block_step
        );
    }
}

static std::variant<CUmodule, std::string> compile(
    const char * user_kernel,
    CUdevice device,
    int radius,
    int block_size,
    int block_step,
    bool in_place,
    int warp_size,
    int warps_per_block,
    int sample_type,
    int bits_per_sample
) {

    auto set_error = [](const char * error_message) -> std::string {
        return std::string{ error_message };
    };

    int major;
    checkError(cuDeviceGetAttribute(&major, CU_DEVICE_ATTRIBUTE_COMPUTE_CAPABILITY_MAJOR, device));
    int minor;
    checkError(cuDeviceGetAttribute(&minor, CU_DEVICE_ATTRIBUTE_COMPUTE_CAPABILITY_MINOR, device));
    int compute_capability = major * 10 + minor;

    // find maximum supported architecture
    int num_archs;
    checkError(nvrtcGetNumSupportedArchs(&num_archs));
    const auto supported_archs = std::make_unique<int []>(num_archs);
    checkError(nvrtcGetSupportedArchs(supported_archs.get()));

    bool generate_cubin = compute_capability <= supported_archs[num_archs - 1];

    std::ostringstream kernel_source;
    kernel_source << "#define RADIUS " << radius << '\n';
    kernel_source << "#define BLOCK_SIZE " << block_size << '\n';
    kernel_source << "#define BLOCK_STEP " << block_step << '\n';
    kernel_source << "#define IN_PLACE " << (int) in_place << '\n';
    kernel_source << "#define WARP_SIZE " << warp_size << '\n';
    kernel_source << "#define WARPS_PER_BLOCK " << warps_per_block << '\n';
    if (sample_type == stInteger) {
        int bytes_per_sample = bits_per_sample / 8;
        const char * type {};
        if (bytes_per_sample == 1) {
            type = "unsigned char";
        } else if (bytes_per_sample == 2) {
            type = "unsigned short";
        } else if (bytes_per_sample == 4) {
            type = "unsigned int";
        }
        kernel_source << "#define TYPE " << type << '\n';
        kernel_source << "#define SCALE " << 1.0 / (1 << (bits_per_sample - 8)) << '\n';
        kernel_source << "#define PEAK " << ((1 << bits_per_sample) - 1) << '\n';
    } else if (sample_type == stFloat) {
        if (bits_per_sample == 32) {
            kernel_source << "#define TYPE float\n";
        }
        kernel_source << "#define SCALE 255.0\n";
    }
    kernel_source << user_kernel << '\n';
    kernel_source << kernel_implementation;

    nvrtcProgram program;
    checkError(nvrtcCreateProgram(&program, kernel_source.str().c_str(), nullptr, 0, nullptr, nullptr));
    Resource<nvrtcProgram *, nvrtcDestroyProgramCustom> destroyer { &program };

    const std::string arch_str = {
        generate_cubin ?
        "-arch=sm_" + std::to_string(compute_capability) :
        "-arch=compute_" + std::to_string(supported_archs[num_archs - 1])
    };

    const char * opts[] = {
        arch_str.c_str(),
        "-use_fast_math",
        "-std=c++17",
        "-modify-stack-limit=false"
    };

    auto compilation = nvrtcCompileProgram(program, (int) std::extent_v<decltype(opts)>, opts);

    size_t log_size;
    showError(nvrtcGetProgramLogSize(program, &log_size));

    std::string error_message;
    if (log_size > 1) {
        error_message.resize(log_size);
        showError(nvrtcGetProgramLog(program, error_message.data()));
    }

    if (success(compilation)) {
        if (log_size > 1) {
            std::fprintf(stderr, "nvrtc: %s\n", error_message.c_str());
        }
    } else {
        return error_message;
    }

    std::unique_ptr<char []> image;
    if (generate_cubin) {
        size_t cubin_size;
        checkError(nvrtcGetCUBINSize(program, &cubin_size));
        image = std::make_unique<char[]>(cubin_size);
        checkError(nvrtcGetCUBIN(program, image.get()));
    } else {
        size_t ptx_size;
        checkError(nvrtcGetPTXSize(program, &ptx_size));
        image = std::make_unique<char[]>(ptx_size);
        checkError(nvrtcGetPTX(program, image.get()));
    }

    CUmodule module;
    checkError(cuModuleLoadData(&module, image.get()));

    return module;
}


struct DFTTestThreadData {
    uint8_t * h_padded; // shape: (pad_height, pad_width)
};


struct DFTTestData {
    VSNodeRef * node;
    int radius;
    int block_size;
    int block_step;
    std::array<bool, 3> process;
    CUdevice device; // device_id
    bool in_place;

    int warp_size;

    // most existing devices contain four schedulers per sm
    int warps_per_block = 4;

    CUcontext context; // use primary stream for interoperability
    Resource<CUstream, cuStreamDestroyCustom> stream;

    Resource<CUevent, cuEventDestroyCustom> event;

    // shape: (vertical_num, horizontal_num, 2*radius+1, block_size, block_size)
    Resource<CUdeviceptr, cuMemFreeCustom, true> d_spatial;

    // shape: (vertical_num, horizontal_num, 2*radius+1, block_size, block_size/2+1)
    Resource<CUdeviceptr, cuMemFreeCustom, true> d_frequency;

    std::mutex lock;

    // padded shape: (pad_height, pad_width)
    // NOTE: must be declared before cuFFT handles so it outlives them
    // (cuFFT handles reference this memory via cufftSetWorkArea)
    Resource<CUdeviceptr, cuMemFreeCustom, true> d_work_area_or_padded;

    // 2-D or 3-D, depends on radius
    Resource<cufftHandle, cufftDestroyCustom, true> rfft_handle;
    Resource<cufftHandle, cufftDestroyCustom, true> irfft_handle;
    Resource<cufftHandle, cufftDestroyCustom, true> subsampled_rfft_handle;
    Resource<cufftHandle, cufftDestroyCustom, true> subsampled_irfft_handle;

    Resource<CUmodule, cuModuleUnloadCustom> module;
    CUfunction filter_kernel;
    int filter_num_blocks;
    CUfunction im2col_kernel;
    int im2col_num_blocks;
    CUfunction col2im_kernel;

    std::atomic<int> num_uninitialized_threads;
    std::unordered_map<std::thread::id, DFTTestThreadData> thread_data;
    std::shared_mutex thread_data_lock;
};

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

    auto d = static_cast<const DFTTestData *>(*instanceData);

    auto vi = vsapi->getVideoInfo(d->node);
    vsapi->setVideoInfo(vi, 1, node);
}

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

    auto d = static_cast<DFTTestData *>(*instanceData);

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

    auto set_error = [vsapi, frameCtx](const char * error_message) -> std::nullptr_t {
        vsapi->setFilterError(error_message, frameCtx);
        return nullptr;
    };

    checkError(cuCtxPushCurrent(d->context));
    context_popper context_popper;

    auto vi = vsapi->getVideoInfo(d->node);

    DFTTestThreadData thread_data;

    auto thread_id = std::this_thread::get_id();
    if (d->num_uninitialized_threads.load(std::memory_order::acquire) == 0) {
        const auto & const_data = d->thread_data;
        thread_data = const_data.at(thread_id);
    } else {
        bool initialized = true;

        d->thread_data_lock.lock_shared();
        try {
            const auto & const_data = d->thread_data;
            thread_data = const_data.at(thread_id);
        } catch (const std::out_of_range &) {
            initialized = false;
        }
        d->thread_data_lock.unlock_shared();

        if (!initialized) {
            auto padded_size = (
                (2 * d->radius + 1) *
                calc_pad_size(vi->height, d->block_size, d->block_step) *
                calc_pad_size(vi->width, d->block_size, d->block_step) *
                vi->format->bytesPerSample
            );

            checkError(cuMemHostAlloc((void **) &thread_data.h_padded, padded_size, 0));

            {
                std::lock_guard _ { d->thread_data_lock };
                d->thread_data.emplace(thread_id, thread_data);
            }

            d->num_uninitialized_threads.fetch_sub(1, std::memory_order::release);
        }
    }

    std::vector<std::unique_ptr<const VSFrameRef, decltype(vsapi->freeFrame)>> src_frames;
    src_frames.reserve(2 * d->radius + 1);
    for (int i = n - d->radius; i <= n + d->radius; i++) {
        src_frames.emplace_back(
            vsapi->getFrameFilter(std::clamp(i, 0, vi->numFrames - 1), d->node, frameCtx),
            vsapi->freeFrame
        );
    }

    auto & src_center_frame = src_frames[d->radius];
    auto format = vsapi->getFrameFormat(src_center_frame.get());

    const VSFrameRef * fr[] {
        d->process[0] ? nullptr : src_center_frame.get(),
        d->process[1] ? nullptr : src_center_frame.get(),
        d->process[2] ? nullptr : src_center_frame.get()
    };
    const int pl[] { 0, 1, 2 };
    std::unique_ptr<VSFrameRef, decltype(vsapi->freeFrame)> dst_frame {
        vsapi->newVideoFrame2(format, vi->width, vi->height, fr, pl, src_center_frame.get(), core),
        vsapi->freeFrame
    };

    for (int plane = 0; plane < format->numPlanes; plane++) {
        if (!d->process[plane]) {
            continue;
        }

        int width = vsapi->getFrameWidth(src_center_frame.get(), plane);
        int height = vsapi->getFrameHeight(src_center_frame.get(), plane);
        int stride = vsapi->getStride(src_center_frame.get(), plane) / vi->format->bytesPerSample;

        bool subsampled = vi->format->subSamplingW != 0 || vi->format->subSamplingW != 0;
        auto & rfft_handle = (plane == 0 || !subsampled) ? d->rfft_handle : d->subsampled_rfft_handle;
        auto & irfft_handle = (plane == 0 || !subsampled) ? d->irfft_handle : d->subsampled_irfft_handle;

        int padded_size_spatial = (
            calc_pad_size(height, d->block_size, d->block_step) *
            calc_pad_size(width, d->block_size, d->block_step)
        );

        for (int i = 0; i < 2 * d->radius + 1; i++) {
            auto srcp = vsapi->getReadPtr(src_frames[i].get(), plane);
            reflection_padding(
                &thread_data.h_padded[(i * padded_size_spatial) * vi->format->bytesPerSample],
                srcp,
                width, height, stride,
                d->block_size, d->block_step,
                vi->format->bytesPerSample
            );
        }

        {
            std::lock_guard lock { d->lock };

            CUdeviceptr d_buffer = d->in_place ? d->d_frequency.data : d->d_spatial.data;

            int padded_bytes = (2 * d->radius + 1) * padded_size_spatial * vi->format->bytesPerSample;
            checkError(cuMemcpyHtoDAsync(d->d_work_area_or_padded.data, thread_data.h_padded, padded_bytes, d->stream));
            {
                void * params[] { &d_buffer, &d->d_work_area_or_padded.data, &width, &height };
                checkError(cuLaunchKernel(
                    d->im2col_kernel,
                    d->im2col_num_blocks, 1, 1,
                    d->warps_per_block * d->warp_size, 1, 1,
                    0,
                    d->stream,
                    params, nullptr
                ));
            }
            checkError(cufftExecR2C(
                rfft_handle,
                (cufftReal *) d_buffer,
                (cufftComplex *) d->d_frequency.data
            ));
            {
                int num_blocks = (
                    calc_pad_num(height, d->block_size, d->block_step) *
                    calc_pad_num(width, d->block_size, d->block_step)
                );
                void * params[] { &d->d_frequency.data, &num_blocks };
                checkError(cuLaunchKernel(
                    d->filter_kernel,
                    d->filter_num_blocks, 1, 1,
                    d->warps_per_block * d->warp_size, 1, 1,
                    0,
                    d->stream,
                    params, nullptr
                ));
            }
            checkError(cufftExecC2R(
                irfft_handle,
                (cufftComplex *) d->d_frequency.data,
                (cufftReal *) d_buffer
            ));
            {
                void * params[] { &d->d_work_area_or_padded.data, &d_buffer, &width, &height };
                unsigned int vertical_size = calc_pad_size(height, d->block_size, d->block_step);
                unsigned int horizontal_size = calc_pad_size(width, d->block_size, d->block_step);
                unsigned int grid_x = (horizontal_size + d->warp_size - 1) / d->warp_size;
                unsigned int grid_y = (vertical_size + d->warps_per_block - 1) / d->warps_per_block;
                checkError(cuLaunchKernel(
                    d->col2im_kernel,
                    grid_x, grid_y, 1,
                    d->warp_size, d->warps_per_block, 1,
                    0,
                    d->stream,
                    params, nullptr
                ));
            }
            {
                size_t pad_width = calc_pad_size(width, d->block_size, d->block_step);
                size_t pad_height = calc_pad_size(height, d->block_size, d->block_step);
                const CUDA_MEMCPY3D config {
                    .srcXInBytes = (pad_width - width) / 2 * vi->format->bytesPerSample,
                    .srcY = (pad_height - height) / 2,
                    .srcZ = (size_t) d->radius,
                    .srcMemoryType = CU_MEMORYTYPE_DEVICE,
                    .srcDevice = d->d_work_area_or_padded.data,
                    .srcPitch = pad_width * vi->format->bytesPerSample,
                    .srcHeight = pad_height,
                    .dstXInBytes = (pad_width - width) / 2 * vi->format->bytesPerSample,
                    .dstY = (pad_height - height) / 2,
                    .dstZ = 0, // vs_bitblt(dstp) copies from the 0-th slice
                    .dstMemoryType = CU_MEMORYTYPE_HOST,
                    .dstHost = thread_data.h_padded,
                    .dstPitch = pad_width * vi->format->bytesPerSample,
                    .dstHeight = pad_height,
                    .WidthInBytes = (size_t) width * vi->format->bytesPerSample,
                    .Height = (size_t) height,
                    .Depth = 1
                };
                checkError(cuMemcpy3DAsync(&config, d->stream));
            }

            checkError(cuEventRecord(d->event, d->stream));
            checkError(cuEventSynchronize(d->event));
        }

        int pad_width = calc_pad_size(width, d->block_size, d->block_step);
        int pad_height = calc_pad_size(height, d->block_size, d->block_step);
        int offset_y = (pad_height - height) / 2;
        int offset_x = (pad_width - width) / 2;

        auto dstp = vsapi->getWritePtr(dst_frame.get(), plane);
        auto input = &thread_data.h_padded[(offset_y * pad_width + offset_x) * vi->format->bytesPerSample];
        vs_bitblt(
            dstp, stride * vi->format->bytesPerSample,
            input, pad_width * vi->format->bytesPerSample,
            width * vi->format->bytesPerSample, height
        );
    }

    return dst_frame.release();
}

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

    auto d = static_cast<const DFTTestData *>(instanceData);

    vsapi->freeNode(d->node);

    showError(cuCtxPushCurrent(d->context));

    for (const auto & [_, thread_data] : d->thread_data) {
        showError(cuMemFreeHost(thread_data.h_padded));
    }

    auto device = d->device;

    delete d;

    showError(cuCtxPopCurrent(nullptr));

    showError(cuDevicePrimaryCtxRelease(device));
}

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

    bool context_retained = false;
    bool context_pushed = false;

    context_releaser context_releaser { &context_retained };

    // pop context before release
    context_popper context_popper { &context_pushed };

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

    d->node = vsapi->propGetNode(in, "clip", 0, nullptr);
    node_freer node_freer { vsapi, d->node };

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

    auto vi = vsapi->getVideoInfo(d->node);

    auto user_kernel = vsapi->propGetData(in, "kernel", 0, nullptr);

    int error;

    d->radius = int64ToIntS(vsapi->propGetInt(in, "radius", 0, &error));
    if (error) {
        d->radius = 0;
    }

    d->block_size = int64ToIntS(vsapi->propGetInt(in, "block_size", 0, &error));
    if (error) {
        d->block_size = 8;
    }

    d->block_step = int64ToIntS(vsapi->propGetInt(in, "block_step", 0, &error));
    if (error) {
        d->block_step = d->block_size;
    }

    int num_planes_args = vsapi->propNumElements(in, "planes");
    d->process.fill(num_planes_args <= 0);
    for (int i = 0; i < num_planes_args; ++i) {
        int plane = static_cast<int>(vsapi->propGetInt(in, "planes", i, nullptr));

        if (plane < 0 || plane >= vi->format->numPlanes) {
            return set_error("plane index out of range");
        }

        if (d->process[plane]) {
            return set_error("plane specified twice");
        }

        d->process[plane] = true;
    }

    d->in_place = !!(vsapi->propGetInt(in, "in_place", 0, &error));
    if (error) {
        d->in_place = false;
    }

    int device_id = int64ToIntS(vsapi->propGetInt(in, "device_id", 0, &error));
    if (error) {
        device_id = 0;
    }

    checkError(cuInit(0));
    checkError(cuDeviceGet(&d->device, device_id));

    checkError(cuDevicePrimaryCtxRetain(&d->context, d->device));
    context_retained = true;
    context_releaser.device = d->device;

    checkError(cuCtxPushCurrent(d->context));
    context_pushed = true;

    checkError(cuDeviceGetAttribute(&d->warp_size, CU_DEVICE_ATTRIBUTE_WARP_SIZE, d->device));

    auto compilation = compile(
        user_kernel,
        d->device,
        d->radius, d->block_size, d->block_step, d->in_place,
        d->warp_size, d->warps_per_block,
        vi->format->sampleType, vi->format->bitsPerSample
    );
    if (std::holds_alternative<std::string>(compilation)) {
        std::ostringstream message;
        message << '[' << __LINE__ << "] compile(): " << std::get<std::string>(compilation);
        vsapi->setError(out, message.str().c_str());
        return ;
    }
    d->module = std::get<CUmodule>(compilation);

    int num_sms;
    checkError(cuDeviceGetAttribute(&num_sms, CU_DEVICE_ATTRIBUTE_MULTIPROCESSOR_COUNT, d->device));

    checkError(cuModuleGetFunction(&d->filter_kernel, d->module, "frequency_filtering"));
    {
        int max_blocks_per_sm;
        checkError(cuOccupancyMaxActiveBlocksPerMultiprocessor(
            &max_blocks_per_sm,
            d->filter_kernel,
            d->warps_per_block * d->warp_size,
            0
        ));
        d->filter_num_blocks = num_sms * max_blocks_per_sm;
    }

    checkError(cuModuleGetFunction(&d->im2col_kernel, d->module, "im2col"));
    {
        int max_blocks_per_sm;
        checkError(cuOccupancyMaxActiveBlocksPerMultiprocessor(
            &max_blocks_per_sm,
            d->im2col_kernel,
            d->warps_per_block * d->warp_size,
            0
        ));
        d->im2col_num_blocks = num_sms * max_blocks_per_sm;
    }

    checkError(cuModuleGetFunction(&d->col2im_kernel, d->module, "col2im"));

    checkError(cuStreamCreate(&d->stream.data, CU_STREAM_NON_BLOCKING));

    checkError(cuEventCreate(
        &d->event.data,
        CU_EVENT_BLOCKING_SYNC | CU_EVENT_DISABLE_TIMING
    ));

    size_t padded_bytes = (
        (2 * d->radius + 1) *
        calc_pad_size(vi->height, d->block_size, d->block_step) *
        calc_pad_size(vi->width, d->block_size, d->block_step) *
        vi->format->bytesPerSample
    );
    // merge allocation to fft's work_area
    // checkError(cuMemAlloc(&d->d_padded.data, padded_bytes));

    if (!d->in_place) {
        size_t spatial_bytes = (
            calc_pad_num(vi->height, d->block_size, d->block_step) *
            calc_pad_num(vi->width, d->block_size, d->block_step) *
            (2 * d->radius + 1) *
            square(d->block_size) *
            sizeof(cufftReal)
        );
        checkError(cuMemAlloc(&d->d_spatial.data, spatial_bytes));
    }

    size_t frequency_bytes = (
        (2 * d->radius + 1) *
        calc_pad_num(vi->height, d->block_size, d->block_step) *
        calc_pad_num(vi->width, d->block_size, d->block_step) *
        d->block_size *
        (d->block_size / 2 + 1) *
        sizeof(cufftComplex)
    );
    checkError(cuMemAlloc(&d->d_frequency.data, frequency_bytes));

    // init cufft
    {
        size_t max_work_size { padded_bytes };

        int batch = (
            calc_pad_num(vi->height, d->block_size, d->block_step) *
            calc_pad_num(vi->width, d->block_size, d->block_step)
        );
        std::array<int, 3> fft_size { 2 * d->radius + 1, d->block_size, d->block_size };
        int is_spatial = d->radius == 0;
        int fft_rank = 3 - is_spatial;
        auto fft_n = &fft_size[is_spatial];

        checkError(cufftPlanMany(
            &d->rfft_handle.data,
            fft_rank, fft_n,
            nullptr, 1, 0,
            nullptr, 1, 0,
            CUFFT_R2C, batch
        ));
        {
            size_t work_size;
            checkError(cufftGetSize(d->rfft_handle, &work_size));
            max_work_size = std::max(max_work_size, work_size);
        }
        checkError(cufftSetWorkArea(d->rfft_handle, nullptr)); // free work area
        checkError(cufftSetStream(d->rfft_handle, d->stream));

        checkError(cufftPlanMany(
            &d->irfft_handle.data,
            fft_rank, fft_n,
            nullptr, 1, 0,
            nullptr, 1, 0,
            CUFFT_C2R, batch
        ));
        {
            size_t work_size;
            checkError(cufftGetSize(d->irfft_handle, &work_size));
            max_work_size = std::max(max_work_size, work_size);
        }
        checkError(cufftSetWorkArea(d->irfft_handle, nullptr)); // free work area
        checkError(cufftSetStream(d->irfft_handle, d->stream));

        if (vi->format->subSamplingW != 0 || vi->format->subSamplingH != 0) {
            int subsampled_batch = (
                calc_pad_num(vi->height >> vi->format->subSamplingH, d->block_size, d->block_step) *
                calc_pad_num(vi->width >> vi->format->subSamplingW, d->block_size, d->block_step)
            );

            checkError(cufftPlanMany(
                &d->subsampled_rfft_handle.data,
                fft_rank, fft_n,
                nullptr, 1, 0,
                nullptr, 1, 0,
                CUFFT_R2C, subsampled_batch
            ));
            {
                size_t work_size;
                checkError(cufftGetSize(d->subsampled_rfft_handle, &work_size));
                max_work_size = std::max(max_work_size, work_size);
            }
            checkError(cufftSetWorkArea(d->subsampled_rfft_handle, nullptr)); // free work area
            checkError(cufftSetStream(d->subsampled_rfft_handle, d->stream));

            checkError(cufftPlanMany(
                &d->subsampled_irfft_handle.data,
                fft_rank, fft_n,
                nullptr, 1, 0,
                nullptr, 1, 0,
                CUFFT_C2R, subsampled_batch
            ));
            {
                size_t work_size;
                checkError(cufftGetSize(d->subsampled_irfft_handle, &work_size));
                max_work_size = std::max(max_work_size, work_size);
            }
            checkError(cufftSetWorkArea(d->subsampled_irfft_handle, nullptr)); // free work area
            checkError(cufftSetStream(d->subsampled_irfft_handle, d->stream));
        }

        checkError(cuMemAlloc(&d->d_work_area_or_padded.data, max_work_size));
        checkError(cufftSetWorkArea(d->rfft_handle, (void *) d->d_work_area_or_padded.data));
        checkError(cufftSetWorkArea(d->irfft_handle, (void *) d->d_work_area_or_padded.data));
        if (vi->format->subSamplingW != 0 || vi->format->subSamplingH != 0) {
            checkError(cufftSetWorkArea(d->subsampled_rfft_handle, (void *) d->d_work_area_or_padded.data));
            checkError(cufftSetWorkArea(d->subsampled_irfft_handle, (void *) d->d_work_area_or_padded.data));
        }
    }

    VSCoreInfo info;
    vsapi->getCoreInfo2(core, &info);
    d->num_uninitialized_threads.store(info.numThreads, std::memory_order_relaxed);
    d->thread_data.reserve(info.numThreads);

    vsapi->createFilter(
        in, out, "DFTTest",
        DFTTestInit, DFTTestGetFrame, DFTTestFree,
        fmParallel, 0, d.release(), core
    );

    node_freer.release();
    context_releaser.release();
}

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

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

    int ndim = vsapi->propNumElements(in, "shape");
    if (ndim != 1 && ndim != 2 && ndim != 3) {
        return set_error("\"shape\" must be an array of ints with 1, 2 or 3 values");
    }

    std::array<int, 3> shape {};
    {
        auto shape_array = vsapi->propGetIntArray(in, "shape", nullptr);
        for (int i = 0; i < ndim; i++) {
            shape[i] = int64ToIntS(shape_array[i]);
        }
    }

    int size = 1;
    for (int i = 0; i < ndim; i++) {
        size *= shape[i];
    }
    if (vsapi->propNumElements(in, "data") != size) {
        return set_error("cannot reshape array");
    }

    int complex_size = shape[ndim - 1] / 2 + 1;
    for (int i = 0; i < ndim - 1; i++) {
        complex_size *= shape[i];
    }

    auto input = vsapi->propGetFloatArray(in, "data", nullptr);

    auto output = std::make_unique<std::complex<double> []>(complex_size);

    if (ndim == 1) {
        dft(output.get(), input, size, 1);
        vsapi->propSetFloatArray(out, "ret", (const double *) output.get(), complex_size * 2);
    } else if (ndim == 2) {
        for (int i = 0; i < shape[0]; i++) {
            dft(&output[i * (shape[1] / 2 + 1)], &input[i * shape[1]], shape[1], 1);
        }

        auto output2 = std::make_unique<std::complex<double> []>(complex_size);

        for (int i = 0; i < shape[1] / 2 + 1; i++) {
            dft(&output2[i], &output[i], shape[0], shape[1] / 2 + 1);
        }

        vsapi->propSetFloatArray(out, "ret", (const double *) output2.get(), complex_size * 2);
    } else {
        for (int i = 0; i < shape[0] * shape[1]; i++) {
            dft(&output[i * (shape[2] / 2 + 1)], &input[i * shape[2]], shape[2], 1);
        }

        auto output2 = std::make_unique<std::complex<double> []>(complex_size);

        for (int i = 0; i < shape[0]; i++) {
            for (int j = 0; j < shape[2] / 2 + 1; j++) {
                dft(
                    &output2[i * shape[1] * (shape[2] / 2 + 1) + j],
                    &output[i * shape[1] * (shape[2] / 2 + 1) + j],
                    shape[1],
                    (shape[2] / 2 + 1)
                );
            }
        }

        for (int i = 0; i < shape[1] * (shape[2] / 2 + 1); i++) {
            dft(&output[i], &output2[i], shape[0], shape[1] * (shape[2] / 2 + 1));
        }

        vsapi->propSetFloatArray(out, "ret", (const double *) output.get(), complex_size * 2);
    }
}

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

    auto data = vsapi->propGetFloatArray(in, "data", nullptr);
    int num = vsapi->propNumElements(in, "data");

    auto converted_data = std::make_unique<double []>(num);
    for (int i = 0; i < num; i++) {
        converted_data[i] = static_cast<float>(data[i]);
    }

    if (num == 1) {
        vsapi->propSetFloat(out, "ret", converted_data[0], paReplace);
    } else {
        vsapi->propSetFloatArray(out, "ret", converted_data.get(), num);
    }
}

static void Version(const VSMap *, VSMap * out, void *, VSCore *, const VSAPI *vsapi) {
    vsapi->propSetData(out, "version", VERSION, -1, paReplace);

    vsapi->propSetInt(out, "cufft_version_build", CUFFT_VERSION, paReplace);

    int cufft_version;
    if (cufftGetVersion(&cufft_version) == CUFFT_SUCCESS) {
        vsapi->propSetInt(out, "cufft_version", cufft_version, paReplace);
    }
};

VS_EXTERNAL_API(void)
VapourSynthPluginInit(VSConfigPlugin configFunc, VSRegisterFunction registerFunc, VSPlugin *plugin) {
    configFunc(
        "io.github.amusementclub.dfttest2_cuda",
        "dfttest2_cuda",
        "DFTTest2 (CUDA)",
        VAPOURSYNTH_API_VERSION, 1, plugin
    );

    registerFunc(
        "DFTTest",
        "clip:clip;"
        "kernel:data[];"
        "radius:int:opt;"
        "block_size:int:opt;"
        "block_step:int:opt;"
        "planes:int[]:opt;"
        "in_place:int:opt;"
        "device_id:int:opt;",
        DFTTestCreate, nullptr, plugin
    );

    registerFunc(
        "RDFT",
        "data:float[];"
        "shape:int[];",
        RDFT, nullptr, plugin
    );

    registerFunc(
        "ToSingle",
        "data:float[];",
        ToSingle, nullptr, plugin
    );

    registerFunc(
        "Version",
        "",
        Version, nullptr, plugin
    );
}
