#pragma once

#include <device_launch_parameters.h>
#include "helper_math.hpp"
#include <iostream>

// pi
#define Pi (3.14159265358979323846)
// 1/pi
#define InvPi (0.31830988618379067154)
// 1/2pi
#define Inv2Pi (0.15915494309189533577)
// 1/4pi
#define Inv4Pi (0.07957747154594766788)
// pi/2
#define PiOver2 (1.57079632679489661923)
// pi/4
#define PiOver4 (0.78539816339744830961)
// sqrt(2)
#define Sqrt2 (1.41421356237309504880)

#define EPS (0.00001f)

// shift from interval [a,b] to [c,d]
template<typename T>
__host__ __device__ T shiftInterval(T t, T a, T b, T c, T d) {
    return c + (d - c) / (b - a) * (t - a);
}

// shift from interval [a,b] to [0,1]
template<typename T>
__host__ __device__ T shiftInterval01(T t, T a, T b) {
    return (t - a) / (b - a);
}

inline __host__ __device__ float avg(float3 vec) {
    return (vec.x + vec.y + vec.z) / 3.f;
}

inline __host__ __device__ float3 abs(float3 vec) {
    return make_float3(fabsf(vec.x), fabsf(vec.y), fabsf(vec.z));
}

inline __host__ __device__ float sum(float3 vec) {
    return vec.x + vec.y + vec.z;
}

inline __host__ __device__ float4 sqrt(float4 vec) {
    return make_float4(sqrtf(vec.x), sqrtf(vec.y), sqrtf(vec.z), sqrtf(vec.w));
}

inline __host__ __device__ bool operator==(float3 a, float3 b) {
    return a.x == b.x && a.y == b.y && a.z == b.z;
}

inline __host__ __device__ void operator^=(uint2& a, uint2 b) {
    a = make_uint2(a.x ^ b.x, a.y ^ b.y);
}

inline __host__ __device__ uint2 operator^(uint2 a, uint2 b) {
    return make_uint2(a.x ^ b.x, a.y ^ b.y);
}

inline __host__ __device__ uint2 operator>>(uint2 a, uint b) {
    return make_uint2(a.x >> b, a.y >> b);
}

// 3x3 Matrix
struct mat3 {
    float3 x, y, z;
};

inline __host__ __device__ mat3 make_mat3(const float xx, const float xy, const float xz, const float yx,
                                          const float yy, const float yz, const float zx, const float zy,
                                          const float zz) {
    mat3 matrix{};
    matrix.x = make_float3(xx, xy, xz);
    matrix.y = make_float3(yx, yy, yz);
    matrix.z = make_float3(zx, zy, zz);
    return matrix;
}

inline __host__ __device__ mat3 make_mat3_identity() {
    mat3 matrix{};
    matrix.x = make_float3(1, 0, 0);
    matrix.y = make_float3(0, 1, 0);
    matrix.z = make_float3(0, 0, 1);
    return matrix;
}

inline __host__ __device__ mat3 make_mat3(const float3 x, const float3 y, const float3 z) {
    mat3 matrix{};
    matrix.x = x;
    matrix.y = y;
    matrix.z = z;
    return matrix;
}

inline __host__ __device__ mat3 make_mat3(const float value) {
    return make_mat3(make_float3(value), make_float3(value), make_float3(value));
}

inline __host__ __device__ mat3 operator*(const mat3& first, const mat3& second) {
    return make_mat3(first.x * second.x, first.y * second.y, first.z * second.z);
}

inline __host__ __device__ float3 operator*(const mat3& first, const float3& second) {
    return make_float3(first.x.x * second.x + first.x.y * second.y + first.x.z * second.z,
                       first.y.x * second.x + first.y.y * second.y + first.y.z * second.z,
                       first.z.x * second.x + first.z.y * second.y + first.z.z * second.z);
}

namespace CUDA {
    constexpr int THREADS_PER_BLOCK = 64;

    struct ThreadInfo {
        unsigned int localThreadId; // local thread id in the block
        unsigned int threadId; // global thread index
        unsigned int blockId; // id of this block
        unsigned int threadsPerBlock; // id of this block

        __device__ ThreadInfo() {
            localThreadId = threadIdx.x;
            blockId = blockIdx.x;
            threadsPerBlock = blockDim.x;
            threadId = threadsPerBlock * blockId + localThreadId;
        }
    };

    struct Ray {
        float3 eye;
        float3 dir;
    };

    struct Intersection {
        bool hit;
        float tMin;
        float tMax;
    };

    struct LightSample {
        float3 Li;
        float3 wi;
        float pdf;
        float3 pos;
    };

    struct PhaseSample {
        float3 dir;
        float val;
        float pdf;
    };

    struct PosSample {
        float3 color;
        float3 pos;
        bool absorb;
    };

    struct TF {
        float minValue = 0.f;
        float maxValue = 1.f;
        int size = 0;
        float4* data = nullptr;
    };

    struct Volume {
        cudaTextureObject_t volumeData;
        cudaTextureObject_t importanceSDF;
        TF transferFunction;
        float3 bbMin;
        float3 bbMax;
        float density;
        float absorption;
        float scattering;
        float g;
        float3 gridSize;
    };

    struct Camera {
        float3 eye;
        float3 U;
        float3 V;
        float3 W;
    };

    struct PointLight {
        float intensity;
        float3 pos;
    };

    struct OptimizationParameters {
        float p;
        float q;
        float r;
        float lambda;
        float thresholdMin;
        float thresholdMax;
        bool view;
        bool light;
        bool phong;
        float shiny;
        float scale;
        bool sdf;
        bool position;
    };

    inline __host__ __device__ int getBlockCount(int problemSize) {
        return (problemSize + (THREADS_PER_BLOCK - 1)) / (THREADS_PER_BLOCK);
    }

    inline __host__ void checkCudaError(cudaError_t error) {
        if (error != cudaSuccess) {
            std::cerr << "CUDA error: " << cudaGetErrorString(error) << std::endl;
        }
    }

    inline __host__ void checkCudaKernelError(cudaError_t error) {
        if (error != cudaSuccess) {
            std::cerr << "CUDA kernel failed: " << cudaGetErrorString(error) << std::endl;
        }
    }


    template<typename T>
    void free(T*& devPtr) {
        if (devPtr != nullptr) {
            checkCudaError(cudaFree(devPtr));
            devPtr = nullptr;
        }
    }

    template<typename T>
    void freeAsync(T*& devPtr, cudaStream_t stream) {
        if (devPtr != nullptr) {
            checkCudaError(cudaFreeAsync(devPtr, stream));
            devPtr = nullptr;
        }
    }

    template<typename T>
    void malloc(T*& devPtr, size_t size) {
        checkCudaError(cudaMalloc(&devPtr, size * sizeof(T)));
    }

    template<typename T>
    void mallocAsync(T*& devPtr, size_t size, cudaStream_t stream) {
        checkCudaError(cudaMallocAsync(&devPtr, size * sizeof(T), stream));
    }

    template<typename T>
    void freeArray(T*& devPtr) {
        if (devPtr != nullptr) {
            checkCudaError(cudaFreeArray(devPtr));
            devPtr = nullptr;
        }
    }

    template<typename T>
    void memcopy(T*& destPtr, T* srcPtr, size_t size, cudaMemcpyKind kind) {
        checkCudaError(cudaMemcpy(destPtr, srcPtr, size * sizeof(T), kind));
    }

    template<typename T>
    void memcopyToHost(T*& destPtr, T* srcPtr, size_t size) {
        memcopy(destPtr, srcPtr, size, cudaMemcpyDeviceToHost);
    }

    template<typename T>
    void memcopyToDevice(T*& destPtr, T* srcPtr, size_t size) {
        memcopy(destPtr, srcPtr, size, cudaMemcpyHostToDevice);
    }

    template<typename T>
    void memcopyAsync(T*& destPtr, T* srcPtr, size_t size, cudaMemcpyKind kind, cudaStream_t stream) {
        checkCudaError(cudaMemcpyAsync(destPtr, srcPtr, size * sizeof(T), kind, stream));
    }

    template<typename T>
    void memcopyToHostAsync(T*& destPtr, T* srcPtr, size_t size, cudaStream_t stream) {
        memcopyAsync(destPtr, srcPtr, size, cudaMemcpyDeviceToHost, stream);
    }

    template<typename T>
    void memcopyToDeviceAsync(T*& destPtr, T* srcPtr, size_t size, cudaStream_t stream) {
        memcopyAsync(destPtr, srcPtr, size, cudaMemcpyHostToDevice, stream);
    }

    template<typename T>
    void memset(T*& devPtr, size_t size) {
        checkCudaError(cudaMemset(devPtr, 0, size * sizeof(T)));
    }
} // namespace CUDA
