#include <algorithm>
#include <array>
#include <atomic>
#include <cassert>
#include <cmath>
#include <complex>
#include <cstdint>
#include <cstring>
#include <memory>
#include <mutex>
#if __cpp_lib_math_constants
#include <numbers>
#endif // __cpp_lib_math_constants
#include <shared_mutex>
#include <thread>
#include <type_traits>
#include <unordered_map>
#include <vector>

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

#include "dfttest2_cpu.h"
#include "kernel.hpp"

#include <config.h> // generated by cmake, defines "VERSION"


template <typename T, typename T_in>
#if __cpp_concepts
    requires
        (std::is_same_v<T_in, T> || std::is_same_v<T_in, std::complex<T>>)
#endif // __cpp_concepts
static void dft(
    std::complex<T> * VS_RESTRICT dst,
    const T_in * VS_RESTRICT src,
    int n,
    int stride
) {
#if __cpp_lib_math_constants
    const auto pi = std::numbers::pi_v<T>;
#else // __cpp_lib_math_constants
    const auto pi = static_cast<T>(M_PI);
#endif // __cpp_lib_math_constants

    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 * pi / n;
            auto weight = std::complex(std::cos(imag), std::sin(imag));
            sum += src[j * stride] * weight;
        }
        dst[i * stride] = sum;
    }
}


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 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) {
        std::free(thread_data.padded2);
        std::free(thread_data.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);

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

    auto vi = vsapi->getVideoInfo(d->node);
    if (!isConstantFormat(vi)) {
        return set_error("only constant format input is supported");
    }
    if (vi->format->sampleType == stInteger && vi->format->bytesPerSample > 2) {
        return set_error("only 8-16 bit integer format input is supported");
    }
    if (vi->format->sampleType == stFloat && vi->format->bitsPerSample != 32) {
        return set_error("only 32-bit float format input is supported");
    }

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

    {
        auto ptr = vs_aligned_malloc<float>(
            (2 * d->radius + 1) * d->block_size * d->block_size * sizeof(float),
            64
        );
        if (ptr == nullptr) {
            return set_error("alloc error");
        }
        d->window.reset(ptr);
    }
    {
        auto window = vsapi->propGetFloatArray(in, "window", nullptr);
        for (int i = 0; i < (2 * d->radius + 1) * d->block_size * d->block_size / 16; i++) {
            for (int j = 0; j < 16; j++) {
                d->window[i * 16 + j] = static_cast<float>(window[i * 16 + j]);
            }
        }
    }

    {
        auto ptr = vs_aligned_malloc<float>(
            (2 * d->radius + 1) * d->block_size * (d->block_size / 2 + 1 + 15) * sizeof(float),
            64
        );
        if (ptr == nullptr) {
            return set_error("alloc error");
        }
        d->sigma.reset(ptr);
    }
    {
        auto sigma = vsapi->propGetFloatArray(in, "sigma", nullptr);
        for (int i = 0; i < (2 * d->radius + 1) * d->block_size; i++) {
            float sigma_padded[16] {};
            for (int j = 0; j < d->block_size / 2 + 1; j++) {
                sigma_padded[j] = static_cast<float>(sigma[i * (d->block_size / 2 + 1) + j]);
            }
            for (int j = 0; j < 16; j++) {
                d->sigma[i * 16 + j] = sigma_padded[j];
            }
        }
    }

    d->sigma2 = static_cast<float>(vsapi->propGetFloat(in, "sigma2", 0, nullptr));
    d->pmin = static_cast<float>(vsapi->propGetFloat(in, "pmin", 0, nullptr));
    d->pmax = static_cast<float>(vsapi->propGetFloat(in, "pmax", 0, nullptr));

    d->filter_type = static_cast<int>(vsapi->propGetInt(in, "filter_type", 0, nullptr));

    d->zero_mean = !!vsapi->propGetInt(in, "zero_mean", 0, &error);
    if (error) {
        d->zero_mean = true;
    }
    if (d->zero_mean) {
        {
            auto ptr = vs_aligned_malloc<float>(
                (2 * d->radius + 1) * d->block_size * (d->block_size / 2 + 1 + 15) * 2 * sizeof(float),
                64
            );
            if (ptr == nullptr) {
                return set_error("alloc error");
            }
            d->window_freq.reset(ptr);
        }
        auto window_freq = vsapi->propGetFloatArray(in, "window_freq", nullptr);
        for (int i = 0; i < (2 * d->radius + 1) * d->block_size; i++) {
            float sigma_padded[32] {};
            for (int j = 0; j < d->block_size / 2 + 1; j++) {
                sigma_padded[j] = static_cast<float>(window_freq[(i * (d->block_size / 2 + 1) + j) * 2]);
                sigma_padded[16 + j] = static_cast<float>(window_freq[(i * (d->block_size / 2 + 1) + j) * 2 + 1]);
            }
            for (int j = 0; j < 32; j++) {
                d->window_freq[i * 2 * 16 + j] = sigma_padded[j];
            }
        }
    }

    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<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 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_gcc",
        "dfttest2_gcc",
        "DFTTest2 (GCC vector extension)",
        VAPOURSYNTH_API_VERSION, 1, plugin
    );

    registerFunc(
        "DFTTest",
        "clip:clip;"
        "window:float[];"
        "sigma:float[];"
        "sigma2:float;"
        "pmin:float;"
        "pmax:float;"
        "filter_type:int;"
        "radius:int:opt;"
        "block_size:int:opt;"
        "block_step:int:opt;"
        "zero_mean:int:opt;"
        "window_freq:float[]:opt;"
        "planes:int[]:opt;",
        DFTTestCreate, nullptr, plugin
    );

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

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