#include<stdio.h>
#include<cuda_runtime.h>
#include"cudaMath.h"
#include "../base/utils.h"
#include "../base/datatypes.h"

__global__ void amxpy_out_gpu(
    const cuDoubleComplex a,
    const cuDoubleComplex* x, const int x_size,
    const cuDoubleComplex* y, const int y_size,
    const uint64_t* indptr, const uint64_t* indices, const cuDoubleComplex* values,
    cuDoubleComplex* out, const int out_size, const int dim)
{
    int kernelIndx = blockIdx.x * blockDim.x + threadIdx.x;
    if (kernelIndx < x_size) {
        uint64_t j = 0; // offset of indices and values
        uint64_t k = 0; // column index
        cuDoubleComplex sum = y[kernelIndx];
        uint64_t start = indptr[kernelIndx];
        uint64_t end = indptr[kernelIndx + 1];
        for (j = start; j < end; ++j) {
            k = indices[j];
            cuDoubleComplex prod = cuCmul(values[j], x[k]);
            sum = cuCadd(sum, cuCmul(a, prod));
        }
        out[kernelIndx] = sum;
    }
}

__global__ void amv_kernel(
    const cuDoubleComplex a,
    const cuDoubleComplex* x, const int x_size,
    const uint64_t* indptr_, const uint64_t* indices_, const cuDoubleComplex* values_,
    cuDoubleComplex* y, const int y_size)
{
    int kernelIndx = blockIdx.x * blockDim.x + threadIdx.x;
    if (kernelIndx < x_size) {
        cuDoubleComplex accum = make_cuDoubleComplex(0, 0);
        for (uint64_t j = indptr_[kernelIndx]; j < indptr_[kernelIndx + 1]; ++j) {
            uint64_t k = indices_[j];
            cuDoubleComplex product = cuCmul(cuCmul(a, values_[j]), x[k]);
            accum = cuCadd(accum, product);
        }
        y[kernelIndx] = accum;
    }
}

__global__ void mv_kernel(
    const cuDoubleComplex* x, const int x_size,
    const uint64_t* indptr_, const uint64_t* indices_, const cuDoubleComplex* values_,
    cuDoubleComplex* y, const int y_size)
{
    int kernelIndx = blockIdx.x * blockDim.x + threadIdx.x;
    if (kernelIndx < x_size) {
        cuDoubleComplex accum = make_cuDoubleComplex(0, 0);
        for (uint64_t j = indptr_[kernelIndx]; j < indptr_[kernelIndx + 1]; ++j) {
            uint64_t k = indices_[j];
            cuDoubleComplex product = cuCmul(values_[j], x[k]);
            accum = cuCadd(accum, product);
        }
        y[kernelIndx] = accum;
    }
}

__global__ void amxsy_out_kernel(
    const cuDoubleComplex a,
    const cuDoubleComplex* x, 
    const cuDoubleComplex* y,
    const uint64_t* indptr, const uint64_t* indices, const cuDoubleComplex* values,
    cuDoubleComplex* out, const int dim)
{
    int kernelIndx = blockIdx.x * blockDim.x + threadIdx.x;
    if (kernelIndx < dim) {
        uint64_t j = 0; // offset of indices and values
        uint64_t k = 0; // column index
        cuDoubleComplex sum = cuCmul(y[kernelIndx], make_cuDoubleComplex(-1, 0));
        uint64_t start = indptr[kernelIndx];
        uint64_t end = indptr[kernelIndx + 1];
        for (j = start; j < end; ++j) {
            k = indices[j];
            sum = cuCadd(sum, cuCmul(cuCmul(a, values[j]), x[k]));
        }
        out[kernelIndx] = sum;
    }
}

__global__ void axpbypcz_out_kernel(
    const cuDoubleComplex a, const cuDoubleComplex* x, const int x_size,
    const cuDoubleComplex b, const cuDoubleComplex* y, const int y_size,
    const cuDoubleComplex c, const cuDoubleComplex* z, const int z_size,
    cuDoubleComplex* out, int dim)
{
    int kernelIndx = blockIdx.x * blockDim.x + threadIdx.x;
    if (kernelIndx < dim) {
        cuDoubleComplex ax = cuCmul(a, x[kernelIndx]);
        cuDoubleComplex by = cuCmul(b, y[kernelIndx]);
        cuDoubleComplex cz = cuCmul(c, z[kernelIndx]);
        out[kernelIndx] = cuCadd(cuCadd(ax, by), cz);
    }
}

__global__ void amxpy_gpu(
    const cuDoubleComplex a,
    const cuDoubleComplex* x, const int x_size,
    cuDoubleComplex* y, const int y_size,
    const int dim,
    const uint64_t* indptr, const uint64_t* indices, const cuDoubleComplex* values)
{
    int kernelIndx = blockIdx.x * blockDim.x + threadIdx.x;
    if (kernelIndx < dim) {
#ifdef _MSC_VER
        uint64_t j = 0; // offset of indices and values
        uint64_t k = 0; // column index
        for (j = indptr[kernelIndx]; j < indptr[kernelIndx + 1]; ++j) {
            k = indices[j];
            y[kernelIndx] = cuCadd(y[kernelIndx], cuCmul(cuCmul(a, values[j]), x[k]));
        }
#else
        cuDoubleComplex temp_sum = make_cuDoubleComplex(0, 0); // Local accumulator
        uint64_t start = indptr[kernelIndx];
        uint64_t end = indptr[kernelIndx + 1];
        for (uint64_t j = start; j < end; ++j) {
            uint64_t k = indices[j];
            cuDoubleComplex prod = cuCmul(cuCmul(a, values[j]), x[k]);
            temp_sum = cuCadd(temp_sum, prod);
        }
        atomicAdd(&y[kernelIndx].x, temp_sum.x); // Using atomic operations if necessary
        atomicAdd(&y[kernelIndx].y, temp_sum.y);
#endif
    }
}

__global__ void amxsy_kernel(
    const double a,
    const cuDoubleComplex* x, const int x_size,
    cuDoubleComplex* y, const int y_size,
    const int dim,
    const uint64_t* indptr, const uint64_t* indices, const cuDoubleComplex* values)
{
    int kernelIndx = blockIdx.x * blockDim.x + threadIdx.x;
    if (kernelIndx < dim) {
        uint64_t j = 0; // offset of indices and values
        uint64_t k = 0; // column index
        y[kernelIndx] = make_cuDoubleComplex(-1 * y[kernelIndx].x, -1 * y[kernelIndx].y);
        for (j = indptr[kernelIndx]; j < indptr[kernelIndx + 1]; ++j) {
            k = indices[j];
            y[kernelIndx] = cuCadd(y[kernelIndx], cuCmul(make_cuDoubleComplex(values[j].x * a, values[j].y * a), x[k]));
        }
    }
}

__global__ void axpy_kernel(
    const cuDoubleComplex a, const cuDoubleComplex* x,
    const int dim,
    cuDoubleComplex* y
)
{
    int kernelIndx = blockIdx.x * blockDim.x + threadIdx.x;
    if (kernelIndx < dim) {
        y[kernelIndx] = cuCadd(y[kernelIndx], cuCmul(a, x[kernelIndx]));
    }
}

__global__ void axpy_kernel(
    const double a, const cuDoubleComplex* x,
    const int dim,
    cuDoubleComplex* y
)
{
    int kernelIndx = blockIdx.x * blockDim.x + threadIdx.x;
    if (kernelIndx < dim) {
        y[kernelIndx] = cuCadd(y[kernelIndx], make_cuDoubleComplex(a*x[kernelIndx].x, a * x[kernelIndx].y));
    }
}

__global__ void xmy_kernel(const cuDoubleComplex* x, const int x_size,
    const cuDoubleComplex* y, const int y_size, const int dim,
    const uint64_t* indptr, const uint64_t* indices, const cuDoubleComplex* values,
    cuDoubleComplex* product)
{
    __shared__  cuDoubleComplex result[1024];
    cuDoubleComplex tmp = make_cuDoubleComplex(0, 0);
    int kernelIndx = blockIdx.x * blockDim.x + threadIdx.x;
    if (kernelIndx < dim) {
        uint64_t j = 0; // offset of indices and values
        uint64_t k = 0; // column index
        for (j = indptr[kernelIndx]; j < indptr[kernelIndx + 1]; ++j) {
            k = indices[j];
            //cuDoubleComplex tempConj = x[kernelIndx];
            cuDoubleComplex tempConj=cuConj(x[kernelIndx]);
            cuDoubleComplex prod = cuCmul(tempConj, values[j]);
            cuDoubleComplex prod2 = cuCmul(prod, y[k]);
            tmp = cuCadd(tmp, prod2);
        }
    }
    result[threadIdx.x] = make_cuDoubleComplex(tmp.x, tmp.y);

    if (threadIdx.x == 0)
    {
        for (size_t i = 0; i < 1024; i++) {
            product[blockIdx.x] = cuCadd(result[i], product[blockIdx.x]);
        }
    }

    return;
}

__global__ void vec_dot_kernel(const cuDoubleComplex* x, 
    const cuDoubleComplex* y, const int size, cuDoubleComplex* out)
{
    __shared__  cuDoubleComplex result[1024];
    int kernelIndx = blockIdx.x * blockDim.x + threadIdx.x;
    cuDoubleComplex tmp = { 0,0 };
    if (kernelIndx < size) {
        tmp = cuCmul(cuConj(x[kernelIndx]), y[kernelIndx]);
    }
    result[threadIdx.x] = make_cuDoubleComplex(tmp.x, tmp.y);
    __syncthreads();
    if (threadIdx.x == 0)
    {
        for (size_t i = 0; i < 1024; i++) {
            out[blockIdx.x] = cuCadd(result[i], out[blockIdx.x]);
        }
    }
    return;
}

__global__ void vec_dot_kernel2(
    const cuDoubleComplex* buffer_D,
    cuDoubleComplex** wf_vb_tn_D,
    const size_t  row,
    const size_t  buffer_size,
    const size_t  num_kernel,
    const bool isRowMajor,cuDoubleComplex* mu_sample_D)
{
	int kernelIndx = blockIdx.x * blockDim.x + threadIdx.x;
	cuDoubleComplex tmp = { 0,0 };        
    cuDoubleComplex* wf_vb_tn=wf_vb_tn_D[kernelIndx];
	if (kernelIndx < num_kernel) {
        int jIndex = kernelIndx;
		for (int idx = 0; idx < buffer_size; idx++)
		{
			tmp = cuCadd(tmp, cuCmul(cuConj(buffer_D[idx]), wf_vb_tn[idx]));
		}
		if (isRowMajor)
			mu_sample_D[num_kernel * row + jIndex] = tmp;
		else
			mu_sample_D[row + num_kernel * jIndex] = tmp;
	}
	return;
}

__global__ void xmy_diag_kernel(
    const cuDoubleComplex* x,const cuDoubleComplex* y, 
    const cuDoubleComplex* m,  const int dim,
    cuDoubleComplex* product)
{
    __shared__  cuDoubleComplex result[1024];
    cuDoubleComplex tmp = make_cuDoubleComplex(0, 0);
    int kernelIndx = blockIdx.x * blockDim.x + threadIdx.x;
    if (kernelIndx < dim) {
		cuDoubleComplex tempConj = cuConj(x[kernelIndx]);
		cuDoubleComplex prod = cuCmul(tempConj, m[kernelIndx]);
        prod = cuCmul(prod, y[kernelIndx]);
		tmp = cuCadd(tmp, prod);
    }
    result[threadIdx.x] = make_cuDoubleComplex(tmp.x, tmp.y);
    if (threadIdx.x == 0)
    {
        for (size_t i = 0; i < 1024; i++) {
            product[blockIdx.x] = cuCadd(result[i], product[blockIdx.x]);
        }
    }
    return;
}

__global__ void axpbypz_kernel(
    const cuDoubleComplex a, const cuDoubleComplex* x,
    const cuDoubleComplex b, const cuDoubleComplex* y,
    const int dim,
    cuDoubleComplex* z)
{
    int kernelIndx = blockIdx.x * blockDim.x + threadIdx.x;
    if (kernelIndx < dim) {
        z[kernelIndx] = cuCadd(z[kernelIndx], cuCadd(cuCmul(a, x[kernelIndx]), cuCmul(b, y[kernelIndx])));
    }
}

__global__ void vec_norm2_kernel(const cuDoubleComplex* x, const int size, cuDoubleComplex* out){
    __shared__  cuDoubleComplex result[1024];
    int kernelIndx = blockIdx.x * blockDim.x + threadIdx.x;
    cuDoubleComplex tmp = { 0,0 };
    if (kernelIndx < size) {
        tmp =cuCmul(cuConj(x[kernelIndx]), x[kernelIndx]);
    }
    result[threadIdx.x] = make_cuDoubleComplex(tmp.x, tmp.y);
    __syncthreads();
    if (threadIdx.x == 0)
    {
        for (size_t i = 0; i < 1024; i++) {
            out[blockIdx.x] = cuCadd(result[i], out[blockIdx.x]);
        }
    }
    return;
}

__global__ void vec_div_sca_kernel(
    cuDoubleComplex* x,const int size,const double a){
    int kernelIndx = blockIdx.x * blockDim.x + threadIdx.x;
    if (kernelIndx < size) {
        x[kernelIndx].x = x[kernelIndx].x / a;
        x[kernelIndx].y = x[kernelIndx].y / a;
    }
}
