#include "cuda_runtime.h"
#include "device_launch_parameters.h"
#include <stdio.h>
#include <opencv2\opencv.hpp>
#include "refocus.h"

using namespace std;
using namespace cv;

__device__ int index(int n, int m)
{
	if (0 <= n && n < m)
		return n;
	else if (n < 0)
		return 0;
	else
		return m - 1;
}

__global__ void refocusKernel(uchar* pLFData, uchar* pRefocusLFData, uchar* pRefocusData,float alpha,
                              int imgH, int imgW, int radius)
{
    const int tidx=blockDim.x*blockIdx.x+threadIdx.x;
    const int tidy=blockDim.y*blockIdx.y+threadIdx.y;
    int diameter=2*radius+1;
    int heightLF=imgH*diameter;
    int widthLF=imgW*diameter;
    int windowSize=diameter*diameter;
    if (tidx<imgW && tidy<imgH)
    {
        float output_color_B = 0;
		float output_color_G = 0;
		float output_color_R = 0;

		for (int j = -radius; j < radius + 1; j++)
		{
			for (int i = -radius; i < radius + 1; i++)
			{
				float x_ind = i*(1 - 1 / alpha) + tidx;   //参考原文献公式,平移
				float y_ind = j*(1 - 1 / alpha) + tidy;

				int x_floor = floor(x_ind); //向下取整得到整数坐标
				int y_floor = floor(y_ind);

				int x_1 = index(x_floor, imgW); //边界判断，左上角a坐标
				int y_1 = index(y_floor, imgH);
				int x_2 = index(x_floor + 1, imgW); //边界判断，右下角d坐标
				int y_2 = index(y_floor + 1, imgH);

				float x_1_w = 1 - (x_ind - x_floor); //插值权重
				float x_2_w = 1 - x_1_w;
				float y_1_w = 1 - (y_ind - y_floor);
				float y_2_w = 1 - y_1_w;

				/*根据a和d像素点在当前单个微透镜下坐标(i,j)得到其在光场图像下坐标*/
				int x_1_index = i + radius + (x_1) * diameter;
				int y_1_index = j + radius + (y_1) * diameter;
				int x_2_index = i + radius + (x_2) * diameter;
				int y_2_index = j + radius + (y_2) * diameter;

				/*a, b, c, d 四点插值*/
				float interp_color_B = y_1_w * x_1_w * pLFData[y_1_index * widthLF * 3 + x_1_index * 3 + 0 ] +
									y_2_w * x_1_w * pLFData[y_1_index * widthLF * 3 + x_2_index * 3 + 0 ] +
									y_1_w * x_2_w * pLFData[y_2_index * widthLF * 3 + x_1_index * 3 + 0 ] +
									y_2_w * x_2_w * pLFData[y_2_index * widthLF * 3 + x_2_index * 3 + 0 ];

				float interp_color_G = y_1_w * x_1_w * pLFData[y_1_index * widthLF * 3 + x_1_index * 3 + 1] +
									y_2_w * x_1_w * pLFData[y_1_index * widthLF * 3 + x_2_index * 3 + 1] +
									y_1_w * x_2_w * pLFData[y_2_index * widthLF * 3 + x_1_index * 3 + 1] +
									y_2_w * x_2_w * pLFData[y_2_index * widthLF * 3 + x_2_index * 3 + 1];

				float interp_color_R = y_1_w * x_1_w * pLFData[y_1_index * widthLF * 3 + x_1_index * 3 + 2] +
									y_2_w * x_1_w * pLFData[y_1_index * widthLF * 3 + x_2_index * 3 + 2] +
									y_1_w * x_2_w * pLFData[y_2_index * widthLF * 3 + x_1_index * 3 + 2] +
									y_2_w * x_2_w * pLFData[y_2_index * widthLF * 3 + x_2_index * 3 + 2];

				// CORRESPONDENCE ANALYSIS
				/*得到当前像素点在光场图像下的坐标*/
				int x_index_remap = i + radius + (tidx) * diameter;
				int y_index_remap = j + radius + (tidy) * diameter;
				/*赋值插值结果，得到重聚焦后的光场图像*/
				pRefocusLFData[x_index_remap * 3 + y_index_remap * widthLF * 3 + 0] = interp_color_B;
				pRefocusLFData[x_index_remap * 3 + y_index_remap * widthLF * 3 + 1] = interp_color_G;
				pRefocusLFData[x_index_remap * 3 + y_index_remap * widthLF * 3 + 2] = interp_color_R;

				// DEFOCUS ANALYSIS
				output_color_R = interp_color_R + output_color_R;
				output_color_G = interp_color_G + output_color_G;
				output_color_B = interp_color_B + output_color_B;
			}
		}
		/*插值结果求平均，得到宏像素值*/
		pRefocusData[tidx * 3 + tidy * imgW * 3 + 0] = output_color_B / windowSize;
		pRefocusData[tidx * 3 + tidy * imgW * 3 + 1] = output_color_G / windowSize;
		pRefocusData[tidx * 3 + tidy * imgW * 3 + 2] = output_color_R / windowSize;
    }
}

void CRefocus::cudaRefocus(Mat &LFRemapImg, Mat &refocusImgLF, Mat &refocusImg, float alpha)
{
    if (true==isFirst)
    {
        int sizeLF=imgHeightLF*imgWidthLF*sizeof(uchar3);
        int size=imgHeight*imgWidth*sizeof(uchar3);
        cudaError err;
        err=cudaMalloc(&pLFImg,sizeLF);
        err=cudaMalloc(&pRefocusLFImg, sizeLF);
        err=cudaMalloc(&pRefocusImg, size);
        err=cudaMemcpy(pLFImg, LFRemapImg.data, sizeLF,cudaMemcpyHostToDevice);
        isFirst=false;
    }

    dim3 block(8,8);
    dim3 grid( (imgWidth+block.x-1)/block.x, (imgHeight+block.y-1)/block.y);
    refocusKernel<<<grid, block>>>(pLFImg, pRefocusLFImg, pRefocusImg,alpha, imgHeight, imgWidth, radius);
    cudaThreadSynchronize();

    cudaError err;
    err=cudaGetLastError();
    if (err!=cudaSuccess)
    {
        cout<<"err="<<err<<endl;
        getchar();
    }

    err=cudaMemcpy(refocusImg.data, pRefocusImg, imgWidth*imgHeight*sizeof(uchar3), cudaMemcpyDeviceToHost);
}