#pragma once
#include <cstring>
#include <cuda_helper.hpp>


template<typename T>
__global__ void generateCacheImportanceSDF(T* dataDevice, int3 size, T* result) {
    CUDA::ThreadInfo ti;
    if (ti.threadId >= size.x * size.y * size.z)
        return;

    unsigned int i = ti.threadId % size.x;
    unsigned int index = ti.threadId / size.x;
    unsigned int j = index % size.y;
    unsigned int k = index / size.y;
    float temp = 0;
    float v = 0;
    int radius = 2;
    for (int a = -radius; a <= radius; ++a) {
        for (int b = -radius; b <= radius; ++b) {
            for (int c = -radius; c <= radius; ++c) {
                int pos_x = i + a;
                int pos_y = j + b;
                int pos_z = k + c;
                if (pos_x > 0 && pos_x < size.x && pos_y > 0 && pos_y < size.y && pos_z > 0 && pos_z < size.z) {
                    float ex = dataDevice[pos_z * size.x * size.y + pos_y * size.x + pos_x];
                    temp += ex;
                    v++;
                }
            }
        }
    }
    result[k * size.x * size.y + j * size.x + i] = temp / v;
}

class Volume {
public:
    template<typename T>
    Volume(int xBound, int yBound, int zBound, float3 spacing, T* data) {
        this->xBound = xBound;
        this->yBound = yBound;
        this->zBound = zBound;
        this->spacing = spacing;

        float x = xBound * spacing.x;
        float y = yBound * spacing.y;
        float z = zBound * spacing.z;

        float max = std::max(x, std::max(y, z));

        aspectX = x / max;
        aspectY = y / max;
        aspectZ = z / max;

        auto extend = make_cudaExtent(xBound, yBound, zBound);
        cudaChannelFormatDesc channelDesc = cudaCreateChannelDesc<T>();
        CUDA::checkCudaError(cudaMalloc3DArray(&d_volumeArray, &channelDesc, extend));

        cudaMemcpy3DParms copyParams = {nullptr};
        copyParams.srcPtr = make_cudaPitchedPtr(data, xBound * sizeof(T), xBound, yBound);
        copyParams.dstArray = d_volumeArray;
        copyParams.extent = extend;
        copyParams.kind = cudaMemcpyHostToDevice;
        CUDA::checkCudaError(cudaMemcpy3D(&copyParams));

        cudaResourceDesc resDescArray{};
        memset(&resDescArray, 0, sizeof(cudaResourceDesc));
        resDescArray.resType = cudaResourceTypeArray;
        resDescArray.res.array.array = d_volumeArray;

        cudaTextureDesc texDescArray{};
        memset(&texDescArray, 0, sizeof(cudaTextureDesc));
        texDescArray.addressMode[0] = cudaAddressModeClamp;
        texDescArray.addressMode[1] = cudaAddressModeClamp;
        texDescArray.addressMode[2] = cudaAddressModeClamp;
        texDescArray.normalizedCoords = true;
        texDescArray.readMode = cudaReadModeNormalizedFloat;
        texDescArray.filterMode = cudaFilterModeLinear;

        CUDA::checkCudaError(cudaCreateTextureObject(&texture, &resDescArray, &texDescArray, nullptr));
    }

    Volume(int xBound, int yBound, int zBound, float3 spacing, float* data) {
        readFloat = true;

        this->xBound = xBound;
        this->yBound = yBound;
        this->zBound = zBound;
        this->spacing = spacing;

        float x = xBound * spacing.x;
        float y = yBound * spacing.y;
        float z = zBound * spacing.z;

        float max = std::max(x, std::max(y, z));

        aspectX = x / max;
        aspectY = y / max;
        aspectZ = z / max;

        auto extend = make_cudaExtent(xBound, yBound, zBound);
        cudaChannelFormatDesc channelDesc = cudaCreateChannelDesc<float>();
        CUDA::checkCudaError(cudaMalloc3DArray(&d_volumeArray, &channelDesc, extend));

        cudaMemcpy3DParms copyParams = {nullptr};
        copyParams.srcPtr = make_cudaPitchedPtr(data, xBound * sizeof(float), xBound, yBound);
        copyParams.dstArray = d_volumeArray;
        copyParams.extent = extend;
        copyParams.kind = cudaMemcpyHostToDevice;
        CUDA::checkCudaError(cudaMemcpy3D(&copyParams));

        cudaResourceDesc resDescArray{};
        memset(&resDescArray, 0, sizeof(cudaResourceDesc));
        resDescArray.resType = cudaResourceTypeArray;
        resDescArray.res.array.array = d_volumeArray;

        cudaTextureDesc texDescArray{};
        memset(&texDescArray, 0, sizeof(cudaTextureDesc));
        texDescArray.addressMode[0] = cudaAddressModeClamp;
        texDescArray.addressMode[1] = cudaAddressModeClamp;
        texDescArray.addressMode[2] = cudaAddressModeClamp;
        texDescArray.normalizedCoords = true;
        texDescArray.readMode = cudaReadModeElementType;
        texDescArray.filterMode = cudaFilterModeLinear;

        CUDA::checkCudaError(cudaCreateTextureObject(&texture, &resDescArray, &texDescArray, nullptr));
    }

    template<typename T>
    Volume(int xBound, int yBound, int zBound, float3 spacing, T* data, bool filter) {
        this->xBound = xBound;
        this->yBound = yBound;
        this->zBound = zBound;
        this->spacing = spacing;

        float x = xBound * spacing.x;
        float y = yBound * spacing.y;
        float z = zBound * spacing.z;

        float max = std::max(x, std::max(y, z));

        aspectX = x / max;
        aspectY = y / max;
        aspectZ = z / max;

        createTexture(data);
    }

    ~Volume() {
        CUDA::checkCudaError(cudaDestroyTextureObject(texture));
        CUDA::freeArray(d_volumeArray);
    }

    [[nodiscard]] const cudaTextureObject_t& getTexture() const {
        return texture;
    }

    [[nodiscard]] float3 getMinBB() const {
        return make_float3(-aspectX, -aspectY, -aspectZ);
    }

    [[nodiscard]] float3 getMaxBB() const {
        return make_float3(aspectX, aspectY, aspectZ);
    }

    [[nodiscard]] int3 getSize() const {
        return make_int3(xBound, yBound, zBound);
    }

    [[nodiscard]] float3 getGridSpacing() const {
        return make_float3(2.f * aspectX / (float) xBound, 2.f * aspectY / (float) yBound,
                           2.f * aspectZ / (float) zBound);
    }

    void updateFiltering(cudaTextureFilterMode mode) {
        CUDA::checkCudaError(cudaDestroyTextureObject(texture));
        cudaTextureReadMode readMode;
        if (readFloat) {
            readMode = cudaReadModeElementType;
        } else {
            readMode = cudaReadModeNormalizedFloat;
        }

        cudaResourceDesc resDescArray{};
        memset(&resDescArray, 0, sizeof(cudaResourceDesc));
        resDescArray.resType = cudaResourceTypeArray;
        resDescArray.res.array.array = d_volumeArray;

        cudaTextureDesc texDescArray{};
        memset(&texDescArray, 0, sizeof(cudaTextureDesc));
        texDescArray.addressMode[0] = cudaAddressModeClamp;
        texDescArray.addressMode[1] = cudaAddressModeClamp;
        texDescArray.addressMode[2] = cudaAddressModeClamp;
        texDescArray.normalizedCoords = true;
        texDescArray.readMode = readMode;
        texDescArray.filterMode = mode;

        CUDA::checkCudaError(cudaCreateTextureObject(&texture, &resDescArray, &texDescArray, nullptr));
    }

private:
    void createTexture(unsigned char* data);
    void createTexture(unsigned short* data);
    void createTexture(short* data);
    int xBound;
    int yBound;
    int zBound;
    float3 spacing;

    float aspectX;
    float aspectY;
    float aspectZ;

    cudaArray* d_volumeArray = nullptr;
    cudaTextureObject_t texture = 0;

    bool readFloat = false;
};
