#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 <variant>
#include <utility>
#include <vector>

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

#include <hip/hip_runtime.h>
#include <hip/hiprtc.h>

#include "dft_kernels.hpp"
#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(hipError_t result) {
    return result == hipSuccess;
}
static bool success(hiprtcResult result) {
    return result == HIPRTC_SUCCESS;
}

static const char * get_error(hipError_t error) {
    return hipGetErrorString(error);
}

static const char * get_error(hiprtcResult error) {
    return hiprtcGetErrorString(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 hipStreamDestroyCustom(hipStream_t stream) {
    showError(hipStreamDestroy(stream));
}

static void hipEventDestroyCustom(hipEvent_t event) {
    showError(hipEventDestroy(event));
}

static void hipFreeCustom(hipDeviceptr_t p) {
    showError(hipFree(p));
}

static void hipModuleUnloadCustom(hipModule_t module) {
    showError(hipModuleUnload(module));
}

static void hiprtcDestroyProgramCustom(hiprtcProgram * program) {
    showError(hiprtcDestroyProgram(program));
}

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

template <typename T, auto deleter>
    requires
        std::default_initializable<T> &&
        std::is_trivially_copy_assignable_v<T> &&
        std::convertible_to<T, bool> &&
        std::invocable<decltype(deleter), T>
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<hipModule_t, std::string> compile(
    const char * user_kernel,
    hipDevice_t 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 };
    };

    hipDeviceProp_t prop;
    checkError(hipGetDeviceProperties(&prop, device));

    constexpr bool generate_bitcode = false;

    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 << fft_header;
    for (const auto & impl : rdft_implementations) {
        kernel_source << impl;
    }
    for (const auto & impl : dft_implementations) {
        kernel_source << impl;
    }
    for (const auto & impl : idft_implementations) {
        kernel_source << impl;
    }
    for (const auto & impl : irdft_implementations) {
        kernel_source << impl;
    }
    kernel_source << kernel_implementation;

    hiprtcProgram program;
    checkError(hiprtcCreateProgram(&program, kernel_source.str().c_str(), nullptr, 0, nullptr, nullptr));
    Resource<hiprtcProgram *, hiprtcDestroyProgramCustom> destroyer { &program };

    const std::string arch_str = std::string("--offload-arch=") + prop.gcnArchName;

    const char * opts[] = {
        arch_str.c_str(),
        "-std=c++17",
        "-ffast-math",
        "-mno-wavefrontsize64", // rdna only
    };

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

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

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

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

    std::unique_ptr<char []> image;
    if (generate_bitcode) {
        size_t bitcode_size;
        checkError(hiprtcGetBitcodeSize(program, &bitcode_size));
        image = std::make_unique_for_overwrite<char[]>(bitcode_size);
        checkError(hiprtcGetBitcode(program, image.get()));
    } else {
        size_t code_size;
        checkError(hiprtcGetCodeSize(program, &code_size));
        image = std::make_unique_for_overwrite<char[]>(code_size);
        checkError(hiprtcGetCode(program, image.get()));
    }

    hipModule_t module;
    checkError(hipModuleLoadData(&module, image.get()));

    return module;
}


struct ticket_semaphore {
    std::atomic<intptr_t> ticket {};
    std::atomic<intptr_t> current {};

    void acquire() noexcept {
        intptr_t tk { ticket.fetch_add(1, std::memory_order::acquire) };
        while (true) {
            intptr_t curr { current.load(std::memory_order::acquire) };
            if (tk <= curr) {
                return;
            }
            current.wait(curr, std::memory_order::relaxed);
        }
    }

    void release() noexcept {
        current.fetch_add(1, std::memory_order::release);
        current.notify_all();
    }
};


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


struct DFTTestStreamData {
    Resource<hipStream_t, hipStreamDestroyCustom> stream;

    Resource<hipEvent_t, hipEventDestroyCustom> event;

    // shape: (vertical_num, horizontal_num, 2*radius+1, block_size, block_size)
    Resource<hipDeviceptr_t, hipFreeCustom> d_spatial;

    Resource<hipDeviceptr_t, hipFreeCustom> d_padded; // shape: (pad_height, pad_width)
};


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

    int warp_size;

    int warps_per_block = 1;

    ticket_semaphore semaphore;
    std::vector<DFTTestStreamData> stream_data;
    std::vector<int> ticket;
    std::mutex ticket_lock;

    Resource<hipModule_t, hipModuleUnloadCustom> module;
    hipFunction_t fused_kernel;
    int fused_num_blocks;
    hipFunction_t 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(hipSetDevice(d->device));

    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(hipHostMalloc((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;

        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
            );
        }

        {
            d->semaphore.acquire();

            int ticket;
            {
                std::lock_guard lock { d->ticket_lock };
                ticket = d->ticket.back();
                d->ticket.pop_back();
            }

            auto & stream_data = d->stream_data[ticket];

            int padded_bytes = (2 * d->radius + 1) * padded_size_spatial * vi->format->bytesPerSample;
            checkError(hipMemcpyHtoDAsync(stream_data.d_padded.data, thread_data.h_padded, padded_bytes, stream_data.stream));
            {
                void * params[] { &stream_data.d_spatial.data, &stream_data.d_padded.data, &width, &height };
                checkError(hipModuleLaunchKernel(
                    d->fused_kernel,
                    d->fused_num_blocks, 1, 1,
                    d->warps_per_block * d->warp_size, 1, 1,
                    0,
                    stream_data.stream,
                    params, nullptr
                ));
            }
            {
                void * params[] { &stream_data.d_padded.data, &stream_data.d_spatial.data, &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(hipModuleLaunchKernel(
                    d->col2im_kernel,
                    grid_x, grid_y, 1,
                    d->warp_size, d->warps_per_block, 1,
                    0,
                    stream_data.stream,
                    params, nullptr
                ));
            }
            {
                unsigned int pad_width = calc_pad_size(width, d->block_size, d->block_step);
                unsigned int pad_height = calc_pad_size(height, d->block_size, d->block_step);
                const HIP_MEMCPY3D config {
                    .srcXInBytes = (pad_width - width) / 2 * vi->format->bytesPerSample,
                    .srcY = (pad_height - height) / 2,
                    .srcZ = (unsigned int) d->radius,
                    .srcMemoryType = hipMemoryTypeDevice,
                    .srcDevice = stream_data.d_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 = hipMemoryTypeHost,
                    .dstHost = thread_data.h_padded,
                    .dstPitch = pad_width * vi->format->bytesPerSample,
                    .dstHeight = pad_height,
                    .WidthInBytes = (unsigned int) width * vi->format->bytesPerSample,
                    .Height = (unsigned int) height,
                    .Depth = 1
                };
                checkError(hipDrvMemcpy3DAsync(&config, stream_data.stream));
            }

            checkError(hipEventRecord(stream_data.event, stream_data.stream));
            checkError(hipEventSynchronize(stream_data.event));

            {
                std::lock_guard lock { d->ticket_lock };
                d->ticket.emplace_back(ticket);
            }
            d->semaphore.release();
        }

        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);

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

    delete d;
}

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

    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;
    }

    if (d->radius < 0 || d->radius > 3) {
        return set_error("\"radius\" must be in [0, 1, 2, 3]");
    }

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

    if (d->block_size != 16) {
        return set_error("\"block_size\" must be 16");
    }

    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 = true;
    }
    if (d->in_place) {
        return set_error("\"in_place\" not supported yet");
    }

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

    int num_streams = int64ToIntS(vsapi->propGetInt(in, "num_streams", 0, &error));
    if (error) {
        num_streams = 1;
    }
    d->semaphore.current.store(num_streams - 1, std::memory_order::relaxed);
    d->ticket.reserve(num_streams);
    for (int i = 0; i < num_streams; i++) {
        d->ticket.emplace_back(i);
    }

    checkError(hipSetDevice(d->device));

    checkError(hipDeviceGetAttribute(&d->warp_size, hipDeviceAttributeWarpSize, 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<hipModule_t>(compilation);

    int num_sms;
    checkError(hipDeviceGetAttribute(&num_sms, hipDeviceAttributeMultiprocessorCount, d->device));

    checkError(hipModuleGetFunction(&d->fused_kernel, d->module, "fused"));
    {
        int max_blocks_per_sm;
        checkError(hipModuleOccupancyMaxActiveBlocksPerMultiprocessor(
            &max_blocks_per_sm,
            d->fused_kernel,
            d->warps_per_block * d->warp_size,
            0
        ));
        d->fused_num_blocks = num_sms * max_blocks_per_sm;
    }

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

    d->stream_data.resize(num_streams);
    for (int i = 0; i < num_streams; i++) {
        auto & stream_data = d->stream_data[i];

        checkError(hipStreamCreateWithFlags(&stream_data.stream.data, hipStreamNonBlocking));

        checkError(hipEventCreateWithFlags(
            &stream_data.event.data,
            hipEventBlockingSync | hipEventDisableTiming
        ));

        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
        );
        checkError(hipMalloc(&stream_data.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(float)
            );
            checkError(hipMalloc(&stream_data.d_spatial.data, spatial_bytes));
        }
    }

    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
    );
}

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_for_overwrite<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_for_overwrite<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_for_overwrite<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_for_overwrite<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);
};

VS_EXTERNAL_API(void)
VapourSynthPluginInit(VSConfigPlugin configFunc, VSRegisterFunction registerFunc, VSPlugin *plugin) {
    configFunc(
        "io.github.amusementclub.dfttest2_hiprtc",
        "dfttest2_hiprtc",
        "DFTTest2 (HIPRTC)",
        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;"
        "num_streams: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
    );
}
