#include <stdlib.h>
#include <string.h>
#include <signal.h>
#include <cassert>

#include "src/acc/acc_backprojector.h"
#include "src/acc/cpu/cpu_kernels/helper.h"

namespace CpuKernels
{

// 2019Mar28: Takanori defined CPU_BP_INITIALIZE here to avoid template instantiation
// problems with CTF_PREMULTIPLIED.
// The result looks OK; I don't know why this is here in the first place.
#define CPU_BP_INITIALIZE

#ifndef CPU_BP_INITIALIZE //TODO Clean this up
template < bool CTF_PREMULTIPLIED >
	void backproject2D(
		unsigned long imageCount,
		int     block_size,
		XFLOAT *g_img_real,
		XFLOAT *g_img_imag,
		XFLOAT *g_trans_x,
		XFLOAT *g_trans_y,
		XFLOAT* g_weights,
		XFLOAT* g_Minvsigma2s,
		XFLOAT* g_ctfs,
		unsigned long translation_num,
		XFLOAT significant_weight,
		XFLOAT weight_norm,
		XFLOAT *g_eulers,
		XFLOAT *g_model_real,
		XFLOAT *g_model_imag,
		XFLOAT *g_model_weight,
		int max_r,
		int max_r2,
		XFLOAT padding_factor,
		unsigned img_x,
		unsigned img_y,
		unsigned img_xy,
		unsigned mdl_x,
		int mdl_inity,
		tbb::spin_mutex *mutexes);
#else
template < bool CTF_PREMULTIPLIED >
void backproject2D(
		unsigned long imageCount,
		int     block_size,
		XFLOAT *g_img_real,
		XFLOAT *g_img_imag,
		XFLOAT *g_trans_x,
		XFLOAT *g_trans_y,
		XFLOAT* g_weights,
		XFLOAT* g_Minvsigma2s,
		XFLOAT* g_ctfs,
		unsigned long translation_num,
		XFLOAT significant_weight,
		XFLOAT weight_norm,
		XFLOAT *g_eulers,
		XFLOAT *g_model_real,
		XFLOAT *g_model_imag,
		XFLOAT *g_model_weight,
		int max_r,
		int max_r2,
		XFLOAT padding_factor,
		unsigned img_x,
		unsigned img_y,
		unsigned img_xy,
		unsigned mdl_x,
		int mdl_inity,
		tbb::spin_mutex *mutexes)
{
	int img_y_half = img_y / 2;

	int max_r2_out = max_r2 * padding_factor * padding_factor;

	for (unsigned long img=0; img<imageCount; img++) {
		XFLOAT s_eulers[4];

		s_eulers[0] = g_eulers[img*9+0] * padding_factor;
		s_eulers[1] = g_eulers[img*9+1] * padding_factor;
		s_eulers[2] = g_eulers[img*9+3] * padding_factor;
		s_eulers[3] = g_eulers[img*9+4] * padding_factor;

		XFLOAT weight_norm_inverse = (XFLOAT) 1.0 / weight_norm;
		XFLOAT inv_minsigma_ctf;

		XFLOAT minvsigma2, ctf, img_real, img_imag, Fweight, real, imag, weight;

		int pixel_pass_num(ceilf((float)img_xy/(float)block_size));

		for(int tid=0; tid<block_size; tid++)
		{
			for (unsigned pass = 0; pass < pixel_pass_num; pass++)
			{
				unsigned long pixel = (pass * block_size) + tid;

				if (pixel >= img_xy)
					continue;

				int x = pixel % img_x;
				int y = (int)floorf( (float)pixel / (float)img_x);

				if (y > img_y_half)
				{
					y -= img_y;
				}

				//WAVG
				minvsigma2 = g_Minvsigma2s[pixel];
				ctf = g_ctfs[pixel];
				img_real = g_img_real[pixel];
				img_imag = g_img_imag[pixel];
				Fweight = (XFLOAT) 0.0;
				real = (XFLOAT) 0.0;
				imag = (XFLOAT) 0.0;
				if(CTF_PREMULTIPLIED)
					inv_minsigma_ctf = weight_norm_inverse * minvsigma2;
				else
					inv_minsigma_ctf = weight_norm_inverse * ctf * minvsigma2;

				XFLOAT temp_real, temp_imag;
				for (unsigned long itrans = 0; itrans < translation_num; itrans++)
				{
					weight = g_weights[img * translation_num + itrans];

					if (weight >= significant_weight)
					{
						weight = weight * inv_minsigma_ctf;
						if(CTF_PREMULTIPLIED)
							Fweight += weight * ctf * ctf;
						else
							Fweight += weight * ctf;


					CpuKernels::translatePixel(x, y, g_trans_x[itrans], g_trans_y[itrans], img_real, img_imag, temp_real, temp_imag);

					real += temp_real * weight;
					imag += temp_imag * weight;
					}
				}

				if (Fweight > (XFLOAT) 0.0)
				{

					// Get logical coordinates in the 3D map
					XFLOAT xp = (s_eulers[0] * x + s_eulers[1] * y );
					XFLOAT yp = (s_eulers[2] * x + s_eulers[3] * y );

					// Only consider pixels that are projected inside the allowed circle in output coordinates.
					//     --JZ, Nov. 26th 2018
					if ( ( xp * xp + yp * yp ) > max_r2_out)
						continue;

					// Only asymmetric half is stored
					if (xp < 0)
					{
						// Get complex conjugated hermitian symmetry pair
						xp = -xp;
						yp = -yp;
						imag = -imag;
					}

					int x0 = floorf(xp);
					XFLOAT fx = xp - x0;
					int x1 = x0 + 1;

					int y0 = floorf(yp);
					XFLOAT fy = yp - y0;
					y0 -= mdl_inity;
					int y1 = y0 + 1;

					XFLOAT mfx = (XFLOAT) 1.0 - fx;
					XFLOAT mfy = (XFLOAT) 1.0 - fy;

					XFLOAT dd00 = mfy * mfx;
					XFLOAT dd01 = mfy *  fx;
					XFLOAT dd10 =  fy * mfx;
					XFLOAT dd11 =  fy *  fx;

					// Locking necessary since all threads share the same back projector
					{
						tbb::spin_mutex::scoped_lock lock(mutexes[y0]);
						g_model_real  [y0 * mdl_x + x0]+=dd00 * real;
						g_model_imag  [y0 * mdl_x + x0]+=dd00 * imag;
						g_model_weight[y0 * mdl_x + x0]+=dd00 * Fweight;

						g_model_real  [y0 * mdl_x + x1]+=dd01 * real;
						g_model_imag  [y0 * mdl_x + x1]+=dd01 * imag;
						g_model_weight[y0 * mdl_x + x1]+=dd01 * Fweight;
					}

					{
						tbb::spin_mutex::scoped_lock lock(mutexes[y1]);
						g_model_real  [y1 * mdl_x + x0]+=dd10 * real;
						g_model_imag  [y1 * mdl_x + x0]+=dd10 * imag;
						g_model_weight[y1 * mdl_x + x0]+=dd10 * Fweight;

						g_model_real  [y1 * mdl_x + x1]+=dd11 * real;
						g_model_imag  [y1 * mdl_x + x1]+=dd11 * imag;
						g_model_weight[y1 * mdl_x + x1]+=dd11 * Fweight;
					}
				}  // Fweight > (RFLOAT) 0.0
			} // for pass
		}  // for tid
	} // img
}
#endif // CPU_BP_INITIALIZE

template < bool DATA3D, bool CTF_PREMULTIPLIED >
void backproject3D(
		unsigned long imageCount,
		int     block_size,
		XFLOAT *g_img_real,
		XFLOAT *g_img_imag,
		XFLOAT *g_trans_x,
		XFLOAT *g_trans_y,
		XFLOAT *g_trans_z,
		XFLOAT* g_weights,
		XFLOAT* g_Minvsigma2s,
		XFLOAT* g_ctfs,
		unsigned long translation_num,
		XFLOAT significant_weight,
		XFLOAT weight_norm,
		XFLOAT *g_eulers,
		XFLOAT *g_model_real,
		XFLOAT *g_model_imag,
		XFLOAT *g_model_weight,
		int max_r,
		int max_r2,
		XFLOAT padding_factor,
		unsigned img_x,
		unsigned img_y,
		unsigned img_z,
		size_t   img_xyz,
		unsigned mdl_x,
		unsigned mdl_y,
		int mdl_inity,
		int mdl_initz,
		tbb::spin_mutex *mutexes)
{
	int img_y_half = img_y / 2;
	int img_z_half = img_z / 2;

	int max_r2_vol = max_r2 * padding_factor * padding_factor;

	for (unsigned long img=0; img<imageCount; img++) {
		XFLOAT s_eulers[9];

		 for (int i = 0; i < 9; i++)
			 s_eulers[i] = g_eulers[img*9+i];

		XFLOAT weight_norm_inverse = (XFLOAT) 1.0 / weight_norm;

		int pixel_pass_num(0);
		if(DATA3D)
			pixel_pass_num = (ceilf((float)img_xyz/(float)block_size));
		else
			pixel_pass_num = (ceilf((float)img_xyz/(float)block_size));

		// We collect block_size number of values before storing the results to
		// help vectorization and control memory accesses
		XFLOAT real[block_size], imag[block_size], Fweight[block_size];
		XFLOAT xp[block_size], yp[block_size], zp[block_size];

		for (unsigned pass = 0; pass < pixel_pass_num; pass++)
		{
			memset(Fweight,0,sizeof(XFLOAT)*block_size);
			#pragma omp simd
			for(int tid=0; tid<block_size; tid++)
			{
				int ok_for_next(1);  // This flag avoids continues, helping the vectorizer

				size_t pixel(0);
				if(DATA3D)
					pixel = ((size_t)pass * (size_t)block_size) + (size_t)tid;
				else
					pixel = ((size_t)pass * (size_t)block_size) + (size_t)tid;

				if (pixel >= img_xyz)
					continue; // just doesn't make sense to proceed in this case

				int x,y,z,xy;
				XFLOAT minvsigma2, ctf, img_real, img_imag, weight;

				if(DATA3D)
				{
					z =  CpuKernels::floorfracf(pixel, (size_t)img_x*(size_t)img_y);
					xy = pixel % (img_x*img_y);
					x =             xy  % img_x;
					y = CpuKernels::floorfracf( xy,   (size_t)img_x);

					if (z > img_z_half)
					{
						z = z - img_z;

						if(x==0)
							ok_for_next=0;
					}
				}
				else
				{
					x =             pixel % img_x;
					y = CpuKernels::floorfracf( pixel , (size_t)img_x);
				}
				if (y > img_y_half)
				{
					y = y - img_y;
				}
				// Get logical coordinates in the 3D map

				if(DATA3D)
				{
					xp[tid] = (s_eulers[0] * x + s_eulers[1] * y + s_eulers[2] * z) * padding_factor;
					yp[tid] = (s_eulers[3] * x + s_eulers[4] * y + s_eulers[5] * z) * padding_factor;
					zp[tid] = (s_eulers[6] * x + s_eulers[7] * y + s_eulers[8] * z) * padding_factor;
				}
				else
				{
					xp[tid] = (s_eulers[0] * x + s_eulers[1] * y ) * padding_factor;
					yp[tid] = (s_eulers[3] * x + s_eulers[4] * y ) * padding_factor;
					zp[tid] = (s_eulers[6] * x + s_eulers[7] * y ) * padding_factor;
				}

				// Only consider pixels that are projected inside the sphere in output coordinates.
				//     --JZ, Nov. 26th 2018
				if ( ( xp[tid] * xp[tid] + yp[tid] * yp[tid] + zp[tid] * zp[tid] ) > max_r2_vol)
				{
					ok_for_next = 0;
				}

				if(ok_for_next)
				{
					//WAVG
					minvsigma2 = g_Minvsigma2s[pixel];
					ctf = g_ctfs[pixel];
					img_real = g_img_real[pixel];
					img_imag = g_img_imag[pixel];
					Fweight[tid] = (XFLOAT) 0.0;
					real[tid] = (XFLOAT) 0.0;
					imag[tid] = (XFLOAT) 0.0;
					XFLOAT inv_minsigma_ctf;
					if(CTF_PREMULTIPLIED)
						inv_minsigma_ctf = weight_norm_inverse * minvsigma2;
					else
						inv_minsigma_ctf = weight_norm_inverse * ctf * minvsigma2;

					XFLOAT temp_real, temp_imag;

					for (unsigned long itrans = 0; itrans < translation_num; itrans++)
					{
						weight = g_weights[img * translation_num + itrans];

						if (weight >= significant_weight)
						{
							weight = weight * inv_minsigma_ctf;
							if(CTF_PREMULTIPLIED)
								Fweight[tid] += weight * ctf * ctf;
							else
								Fweight[tid] += weight * ctf;

							if(DATA3D)
								CpuKernels::translatePixel(x, y, z, g_trans_x[itrans], g_trans_y[itrans], g_trans_z[itrans], img_real, img_imag, temp_real, temp_imag);
							else
								CpuKernels::translatePixel(x, y,    g_trans_x[itrans], g_trans_y[itrans],                    img_real, img_imag, temp_real, temp_imag);

							real[tid] += temp_real * weight;
							imag[tid] += temp_imag * weight;
						}
					}

					//BP
					if (Fweight[tid] > (XFLOAT) 0.0)
					{

						// Only asymmetric half is stored

						if (xp[tid] < (XFLOAT) 0.0)
						{
							// Get complex conjugated hermitian symmetry pair
							xp[tid] = -xp[tid];
							yp[tid] = -yp[tid];
							zp[tid] = -zp[tid];
							imag[tid] = -imag[tid];
						}
					} // Fweight[tid] > (RFLOAT) 0.0
				}// end for if_ok_for_next
			}  // for tid

			for(int tid=0; tid<block_size; tid++)
			{
				if (Fweight[tid] > (XFLOAT) 0.0)
				{
					int x0 = floorf(xp[tid]);
					XFLOAT fx = xp[tid] - x0;
					int x1 = x0 + 1;

					int y0 = floorf(yp[tid]);
					XFLOAT fy = yp[tid] - y0;
					y0 -= mdl_inity;
					int y1 = y0 + 1;

					int z0 = floorf(zp[tid]);
					XFLOAT fz = zp[tid] - z0;
					z0 -= mdl_initz;
					int z1 = z0 + 1;

					XFLOAT mfx = (XFLOAT)1.0 - fx;
					XFLOAT mfy = (XFLOAT)1.0 - fy;
					XFLOAT mfz = (XFLOAT)1.0 - fz;

					// Locking necessary since all threads share the same back projector
					XFLOAT dd000 = mfz * mfy * mfx;
					XFLOAT dd001 = mfz * mfy *  fx;

					size_t z0MdlxMdly = (size_t)z0 * (size_t)mdl_x * (size_t)mdl_y;

					{
						tbb::spin_mutex::scoped_lock lock(mutexes[z0 * mdl_y + y0]);

						g_model_real  [z0MdlxMdly + y0 * mdl_x + x0]+=dd000 * real[tid];
						g_model_imag  [z0MdlxMdly + y0 * mdl_x + x0]+=dd000 * imag[tid];
						g_model_weight[z0MdlxMdly + y0 * mdl_x + x0]+=dd000 * Fweight[tid];

						g_model_real  [z0MdlxMdly + y0 * mdl_x + x1]+=dd001 * real[tid];
						g_model_imag  [z0MdlxMdly + y0 * mdl_x + x1]+=dd001 * imag[tid];
						g_model_weight[z0MdlxMdly + y0 * mdl_x + x1]+=dd001 * Fweight[tid];
					}

					XFLOAT dd010 = mfz *  fy * mfx;
					XFLOAT dd011 = mfz *  fy *  fx;

					{
						tbb::spin_mutex::scoped_lock lock(mutexes[z0 * mdl_y + y1]);

						g_model_real  [z0MdlxMdly + y1 * mdl_x + x0]+=dd010 * real[tid];
						g_model_imag  [z0MdlxMdly + y1 * mdl_x + x0]+=dd010 * imag[tid];
						g_model_weight[z0MdlxMdly + y1 * mdl_x + x0]+=dd010 * Fweight[tid];

						g_model_real  [z0MdlxMdly + y1 * mdl_x + x1]+=dd011 * real[tid];
						g_model_imag  [z0MdlxMdly + y1 * mdl_x + x1]+=dd011 * imag[tid];
						g_model_weight[z0MdlxMdly + y1 * mdl_x + x1]+=dd011 * Fweight[tid];
					}

					XFLOAT dd100 =  fz * mfy * mfx;
					XFLOAT dd101 =  fz * mfy *  fx;

					size_t z1MdlxMdly = (size_t)z1 * (size_t)mdl_x * (size_t)mdl_y;

					{
						tbb::spin_mutex::scoped_lock lock(mutexes[z1 * mdl_y + y0]);

						g_model_real  [z1MdlxMdly + y0 * mdl_x + x0]+=dd100 * real[tid];
						g_model_imag  [z1MdlxMdly + y0 * mdl_x + x0]+=dd100 * imag[tid];
						g_model_weight[z1MdlxMdly + y0 * mdl_x + x0]+=dd100 * Fweight[tid];

						g_model_real  [z1MdlxMdly + y0 * mdl_x + x1]+=dd101 * real[tid];
						g_model_imag  [z1MdlxMdly + y0 * mdl_x + x1]+=dd101 * imag[tid];
						g_model_weight[z1MdlxMdly + y0 * mdl_x + x1]+=dd101 * Fweight[tid];
					}

					XFLOAT dd110 =  fz *  fy * mfx;
					XFLOAT dd111 =  fz *  fy *  fx;

					{
						tbb::spin_mutex::scoped_lock lock(mutexes[z1 * mdl_y + y1]);

						g_model_real  [z1MdlxMdly + y1 * mdl_x + x0]+=dd110 * real[tid];
						g_model_imag  [z1MdlxMdly + y1 * mdl_x + x0]+=dd110 * imag[tid];
						g_model_weight[z1MdlxMdly + y1 * mdl_x + x0]+=dd110 * Fweight[tid];

						g_model_real  [z1MdlxMdly + y1 * mdl_x + x1]+=dd111 * real[tid];
						g_model_imag  [z1MdlxMdly + y1 * mdl_x + x1]+=dd111 * imag[tid];
						g_model_weight[z1MdlxMdly + y1 * mdl_x + x1]+=dd111 * Fweight[tid];
					}
				}  // Fweight[tid] > (RFLOAT) 0.0
			}  // for tid
		} // for pass
	} // for img
}

// sincos lookup table optimization. Function translatePixel calls
// sincos(x*tx + y*ty). We precompute 2D lookup tables for x and y directions.
// The first dimension is x or y pixel index, and the second dimension is x or y
// translation index. Since sin(a+B) = sin(A) * cos(B) + cos(A) * sin(B), and
// cos(A+B) = cos(A) * cos(B) - sin(A) * sin(B), we can use lookup table to
// compute sin(x*tx + y*ty) and cos(x*tx + y*ty).
#ifndef CPU_BP_INITIALIZE //TODO Clean this up
template < bool CTF_PREMULTIPLIED >
void backprojectRef3D(
		unsigned long imageCount,
		XFLOAT *g_img_real,
		XFLOAT *g_img_imag,
		XFLOAT *g_trans_x,
		XFLOAT *g_trans_y,
		XFLOAT* g_weights,
		XFLOAT* g_Minvsigma2s,
		XFLOAT* g_ctfs,
		unsigned long trans_num,
		XFLOAT significant_weight,
		XFLOAT weight_norm,
		XFLOAT *g_eulers,
		XFLOAT *g_model_real,
		XFLOAT *g_model_imag,
		XFLOAT *g_model_weight,
		int     max_r,
		int     max_r2,
		XFLOAT   padding_factor,
		unsigned img_x,
		unsigned img_y,
		unsigned img_z,
		size_t   img_xyz,
		unsigned mdl_x,
		unsigned mdl_y,
		int      mdl_inity,
		int      mdl_initz,
		tbb::spin_mutex *mutexes);
#else
template < bool CTF_PREMULTIPLIED >
void backprojectRef3D(
		unsigned long imageCount,
		XFLOAT *g_img_real,
		XFLOAT *g_img_imag,
		XFLOAT *g_trans_x,
		XFLOAT *g_trans_y,
		XFLOAT* g_weights,
		XFLOAT* g_Minvsigma2s,
		XFLOAT* g_ctfs,
		unsigned long trans_num,
		XFLOAT significant_weight,
		XFLOAT weight_norm,
		XFLOAT *g_eulers,
		XFLOAT *g_model_real,
		XFLOAT *g_model_imag,
		XFLOAT *g_model_weight,
		int     max_r,
		int     max_r2,
		XFLOAT   padding_factor,
		unsigned img_x,
		unsigned img_y,
		unsigned img_z,
		size_t   img_xyz,
		unsigned mdl_x,
		unsigned mdl_y,
		int      mdl_inity,
		int      mdl_initz,
		tbb::spin_mutex *mutexes)
{
	int img_y_half = img_y / 2;
	int img_y_half_2 = img_y_half * img_y_half;
	int img_z_half = img_z / 2;

	int max_r2_vol = max_r2 * padding_factor * padding_factor;

	for (unsigned long img=0; img<imageCount; img++) {
		XFLOAT s_eulers[9];

		for(int i = 0; i < 9; i++)
			s_eulers[i] = g_eulers[img*9+i];

		XFLOAT sin_x[trans_num][img_x], cos_x[trans_num][img_x];
		XFLOAT sin_y[trans_num][img_y], cos_y[trans_num][img_y];

		CpuKernels::computeSincosLookupTable2D(trans_num, g_trans_x, g_trans_y, img_x, img_y,
								   &sin_x[0][0], &cos_x[0][0],
								   &sin_y[0][0], &cos_y[0][0]);

		XFLOAT weight_norm_inverse = (XFLOAT) 1.0 / weight_norm;

		XFLOAT xp[img_x], yp[img_x], zp[img_x];
		XFLOAT real[img_x], imag[img_x], Fweight[img_x];

		size_t mdl_x_mdl_y = (size_t)mdl_x * (size_t)mdl_y;
		size_t pixel = 0;
		for(int iy = 0; iy < img_y; iy++) {
			int y = iy;
			if (iy > img_y_half) {
				y = iy - img_y;
			}

			int y2 = y * y;
			int xmax = sqrt((XFLOAT)(img_y_half_2 - y2));

			memset(Fweight,0,sizeof(XFLOAT)*img_x);
			memset(real,   0,sizeof(XFLOAT)*img_x);
			memset(imag,   0,sizeof(XFLOAT)*img_x);

			for (unsigned long itrans = 0; itrans < trans_num; itrans++)
			{
				XFLOAT weight = g_weights[img * trans_num + itrans];
				if (weight < significant_weight)
					continue;

				XFLOAT trans_cos_y, trans_sin_y;
				if ( y < 0) {
					trans_cos_y =  cos_y[itrans][-y];
					trans_sin_y = -sin_y[itrans][-y];
				}
				else {
					trans_cos_y = cos_y[itrans][y];
					trans_sin_y = sin_y[itrans][y];
				}

				XFLOAT *trans_cos_x = &cos_x[itrans][0];
				XFLOAT *trans_sin_x = &sin_x[itrans][0];

				#pragma omp simd
				for(int x=0; x<xmax; x++) {
					XFLOAT minvsigma2 = g_Minvsigma2s[pixel + x];
					XFLOAT ctf        = g_ctfs       [pixel + x];
					XFLOAT inv_minsigma_ctf;
					XFLOAT my_weight;
					if(CTF_PREMULTIPLIED)
					{
						inv_minsigma_ctf = weight_norm_inverse * minvsigma2;
						my_weight = weight * inv_minsigma_ctf;
						Fweight[x] += my_weight  * ctf * ctf;
					}
					else
					{
						inv_minsigma_ctf = weight_norm_inverse * ctf * minvsigma2;
						my_weight = weight * inv_minsigma_ctf;
						Fweight[x] += my_weight  * ctf;
					}

					XFLOAT img_real   = g_img_real   [pixel + x];
					XFLOAT img_imag   = g_img_imag   [pixel + x];

					XFLOAT ss = trans_sin_x[x] * trans_cos_y + trans_cos_x[x] * trans_sin_y;
					XFLOAT cc = trans_cos_x[x] * trans_cos_y - trans_sin_x[x] * trans_sin_y;

					XFLOAT shifted_real = cc * img_real - ss * img_imag;
					XFLOAT shifted_imag = cc * img_imag + ss * img_real;
					/*
					XFLOAT shifted_real, shifted_imag;
					CpuKernels::translatePixel(x, y,  g_trans_x[itrans], g_trans_y[itrans],
								   img_real, img_imag, shifted_real, shifted_imag);
					*/
					real[x] += shifted_real * my_weight;
					imag[x] += shifted_imag * my_weight;
				}
			}

			#pragma omp simd
			for(int x=0; x<img_x; x++) {
				// Get logical coordinates in the 3D map
				xp[x] = (s_eulers[0] * x + s_eulers[1] * y ) * padding_factor;
				yp[x] = (s_eulers[3] * x + s_eulers[4] * y ) * padding_factor;
				zp[x] = (s_eulers[6] * x + s_eulers[7] * y ) * padding_factor;

				// Use Fweight to discard pixels that project outside the sphere
				// (pixels with Fweight <= 0 will be skipped further down)
				//     --JZ, Nov. 26th 2018
				if (xp[x]*xp[x] + yp[x]*yp[x] + zp[x]*zp[x] > max_r2_vol)
				{
					Fweight[x] = (XFLOAT) 0.0;
				}

				// Only asymmetric half is stored
				if (xp[x] < (XFLOAT) 0.0) {
					// Get complex conjugated hermitian symmetry pair
					xp[x]   = -xp[x];
					yp[x]   = -yp[x];
					zp[x]   = -zp[x];
					imag[x] = -imag[x];
				}
			}  // for x direction

			for(int x=0; x<img_x; x++){
				if (Fweight[x] <= (XFLOAT) 0.0)
					continue;

				int x0 = floorf(xp[x]);
				XFLOAT fx = xp[x] - x0;

				int y0 = floorf(yp[x]);
				XFLOAT fy = yp[x] - y0;
				y0 -= mdl_inity;

				int z0 = floorf(zp[x]);
				XFLOAT fz = zp[x] - z0;
				z0 -= mdl_initz;

				XFLOAT mfx = (XFLOAT)1.0 - fx;
				XFLOAT mfy = (XFLOAT)1.0 - fy;
				XFLOAT mfz = (XFLOAT)1.0 - fz;

				XFLOAT mfz_mfy = mfz * mfy;
				size_t z0_mdl_x_mdl_y = (size_t)z0 * mdl_x_mdl_y;
				size_t y0_mdl_x = (size_t)y0 * (size_t)mdl_x;
				size_t idx_tmp;

				XFLOAT dd000 = mfz_mfy * mfx; // mfz *  mfy *  mfx
				XFLOAT dd001 = mfz_mfy - dd000; // mfz *  mfy *  fx

				{
					tbb::spin_mutex::scoped_lock lock(mutexes[z0 * mdl_y + y0]);

					idx_tmp = z0_mdl_x_mdl_y + y0_mdl_x + (size_t)x0; // z0 * mdl_x * mdl_y + y0 * mdl_x + x0;
					g_model_real  [idx_tmp]+=dd000 * real[x];
					g_model_imag  [idx_tmp]+=dd000 * imag[x];
					g_model_weight[idx_tmp]+=dd000 * Fweight[x];

					idx_tmp = idx_tmp + 1; // z0 * mdl_x * mdl_y + y0 * mdl_x + x1;
					g_model_real  [idx_tmp]+=dd001 * real[x];
					g_model_imag  [idx_tmp]+=dd001 * imag[x];
					g_model_weight[idx_tmp]+=dd001 * Fweight[x];
				}

				XFLOAT dd010 = (mfz - mfz_mfy) * mfx; // mfz *  fy *  mfx
				XFLOAT dd011 = (mfz - mfz_mfy) - dd010; // mfz *  fy *  fx

				{
					tbb::spin_mutex::scoped_lock lock(mutexes[z0 * mdl_y + y0 + 1]);

					idx_tmp = z0_mdl_x_mdl_y + y0_mdl_x + (size_t)mdl_x + (size_t)x0; // z0 * mdl_x * mdl_y + y1 * mdl_x + x0;
					g_model_real  [idx_tmp]+=dd010 * real[x];
					g_model_imag  [idx_tmp]+=dd010 * imag[x];
					g_model_weight[idx_tmp]+=dd010 * Fweight[x];

					idx_tmp = idx_tmp + 1; // z0 * mdl_x * mdl_y + y1 * mdl_x + x1;
					g_model_real  [idx_tmp]+=dd011 * real[x];
					g_model_imag  [idx_tmp]+=dd011 * imag[x];
					g_model_weight[idx_tmp]+=dd011 * Fweight[x];
				}

				XFLOAT dd100 = (mfy - mfz_mfy) * mfx; // fz *  mfy *  mfx
				XFLOAT dd101 = (mfy - mfz_mfy) - dd100; // fz *  mfy *  fx
				int z1 = z0 + 1;

				{
					tbb::spin_mutex::scoped_lock lock(mutexes[z1 * mdl_y + y0]);

					idx_tmp = z0_mdl_x_mdl_y + mdl_x_mdl_y + y0_mdl_x + (size_t)x0; // z1 * mdl_x * mdl_y + y0 * mdl_x + x0;
					g_model_real  [idx_tmp]+=dd100 * real[x];
					g_model_imag  [idx_tmp]+=dd100 * imag[x];
					g_model_weight[idx_tmp]+=dd100 * Fweight[x];

					idx_tmp = idx_tmp + 1; // z1 * mdl_x * mdl_y + y0 * mdl_x + x1;
					g_model_real  [idx_tmp]+=dd101 * real[x];
					g_model_imag  [idx_tmp]+=dd101 * imag[x];
					g_model_weight[idx_tmp]+=dd101 * Fweight[x];

				}

				XFLOAT dd110 = (1 - mfz - mfy + mfz_mfy) * mfx; // fz *  fy *  mfx
				XFLOAT dd111 = (1 - mfz - mfy + mfz_mfy) - dd110; // fz *  fy *  fx

				{
					tbb::spin_mutex::scoped_lock lock(mutexes[z1 * mdl_y + y0 + 1]);
					idx_tmp = z0_mdl_x_mdl_y + mdl_x_mdl_y + y0_mdl_x + (size_t)mdl_x + (size_t)x0; // z1 * mdl_x * mdl_y + y1 * mdl_x + x0;
					g_model_real  [idx_tmp]+=dd110 * real[x];
					g_model_imag  [idx_tmp]+=dd110 * imag[x];
					g_model_weight[idx_tmp]+=dd110 * Fweight[x];

					idx_tmp = idx_tmp + 1; // z1 * mdl_x * mdl_y + y1 * mdl_x + x1;
					g_model_real  [idx_tmp]+=dd111 * real[x];
					g_model_imag  [idx_tmp]+=dd111 * imag[x];
					g_model_weight[idx_tmp]+=dd111 * Fweight[x];
				}
			}  // for x direction

			pixel += (size_t)img_x;
		} // for y direction
	} // for img
}
#endif

template < bool DATA3D, bool CTF_PREMULTIPLIED >
void backprojectSGD(
		unsigned long imageCount,
		int     block_size,
		AccProjectorKernel projector,
		XFLOAT *g_img_real,
		XFLOAT *g_img_imag,
		XFLOAT *g_trans_x,
		XFLOAT *g_trans_y,
		XFLOAT *g_trans_z,
		XFLOAT* g_weights,
		XFLOAT* g_Minvsigma2s,
		XFLOAT* g_ctfs,
		unsigned long translation_num,
		XFLOAT significant_weight,
		XFLOAT weight_norm,
		XFLOAT *g_eulers,
		XFLOAT *g_model_real,
		XFLOAT *g_model_imag,
		XFLOAT *g_model_weight,
		int max_r,
		int max_r2,
		XFLOAT padding_factor,
		unsigned img_x,
		unsigned img_y,
		unsigned img_z,
		size_t   img_xyz,
		unsigned mdl_x,
		unsigned mdl_y,
		int mdl_inity,
		int mdl_initz,
		tbb::spin_mutex *mutexes)
{
	int img_y_half = img_y / 2;
	int img_z_half = img_z / 2;

	int max_r2_vol = max_r2 * padding_factor * padding_factor;

	for (unsigned long img=0; img<imageCount; img++) {
		XFLOAT s_eulers[9];

		for (int i = 0; i < 9; i++)
			s_eulers[i] = g_eulers[img*9+i];

		XFLOAT weight_norm_inverse = (XFLOAT) 1.0 / weight_norm;

		int pixel_pass_num(0);
		if(DATA3D)
			pixel_pass_num = (ceilf((float)img_xyz/(float)block_size));
		else
			pixel_pass_num = (ceilf((float)img_xyz/(float)block_size));

		// TODO - does this really help with the call to the projector in here?
		//
		// We collect block_size number of values before storing the results to
		// help vectorization and control memory accesses
		XFLOAT real[block_size], imag[block_size], Fweight[block_size];
		XFLOAT ref_real[block_size], ref_imag[block_size];
		XFLOAT xp[block_size], yp[block_size], zp[block_size];

		for (unsigned pass = 0; pass < pixel_pass_num; pass++)   {
			memset(Fweight,0,sizeof(XFLOAT)*block_size);
//			#pragma omp simd
			for(int tid=0; tid<block_size; tid++) {
				int ok_for_next(1);  // This flag avoids continues, helping the vectorizer

				size_t pixel(0);
				if(DATA3D)
					pixel = ((size_t)pass * (size_t)block_size) + (size_t)tid;
				else
					pixel = ((size_t)pass * (size_t)block_size) + (size_t)tid;

				if (pixel >= img_xyz)
					continue;  // just doesn't make sense to proceed in this case

				int x,y,z,xy;
				XFLOAT minvsigma2, ctf, img_real, img_imag, weight;

				if(DATA3D)
				{
					z =  CpuKernels::floorfracf(pixel, (size_t)((size_t)img_x*(size_t)img_y));
					xy = pixel % (img_x*img_y);
					x =             xy  % img_x;
					y = CpuKernels::floorfracf( xy,   (size_t)img_x);

					if (z > img_z_half)
					{
						z = z - img_z;

						if(x==0)
							ok_for_next=0;
					}
				}
				else
				{
					x =             pixel % img_x;
					y = CpuKernels::floorfracf( pixel , (size_t)img_x);
				}
				if (y > img_y_half)
				{
					y = y - img_y;
				}
				// Get logical coordinates in the 3D map

				if(DATA3D)
				{
					xp[tid] = (s_eulers[0] * x + s_eulers[1] * y + s_eulers[2] * z) * padding_factor;
					yp[tid] = (s_eulers[3] * x + s_eulers[4] * y + s_eulers[5] * z) * padding_factor;
					zp[tid] = (s_eulers[6] * x + s_eulers[7] * y + s_eulers[8] * z) * padding_factor;
				}
				else
				{
					xp[tid] = (s_eulers[0] * x + s_eulers[1] * y ) * padding_factor;
					yp[tid] = (s_eulers[3] * x + s_eulers[4] * y ) * padding_factor;
					zp[tid] = (s_eulers[6] * x + s_eulers[7] * y ) * padding_factor;
				}

				if (xp[tid]*xp[tid] + yp[tid]*yp[tid] + zp[tid]*zp[tid] > max_r2_vol)
				{
					ok_for_next = 0;
				}

				if(ok_for_next)
				{
					ref_real[tid] = (XFLOAT) 0.0;
					ref_imag[tid] = (XFLOAT) 0.0;

					if(DATA3D)
						projector.project3Dmodel(
							x,y,z,
							s_eulers[0], s_eulers[1], s_eulers[2],
							s_eulers[3], s_eulers[4], s_eulers[5],
							s_eulers[6], s_eulers[7], s_eulers[8],
							ref_real[tid], ref_imag[tid]);
					else
						projector.project3Dmodel(
							x,y,
							s_eulers[0], s_eulers[1],
							s_eulers[3], s_eulers[4],
							s_eulers[6], s_eulers[7],
							ref_real[tid], ref_imag[tid]);

					//WAVG
					minvsigma2 = g_Minvsigma2s[pixel];
					ctf = g_ctfs[pixel];
					img_real = g_img_real[pixel];
					img_imag = g_img_imag[pixel];
					Fweight[tid] = (XFLOAT) 0.0;
					real[tid] = (XFLOAT) 0.0;
					imag[tid] = (XFLOAT) 0.0;
					ref_real[tid] *= ctf;
					ref_imag[tid] *= ctf;
					XFLOAT inv_minsigma_ctf;

					if(CTF_PREMULTIPLIED)
						inv_minsigma_ctf = weight_norm_inverse * minvsigma2;
					else
						inv_minsigma_ctf = weight_norm_inverse * ctf * minvsigma2;

					XFLOAT temp_real, temp_imag;

					for (unsigned long itrans = 0; itrans < translation_num; itrans++)
					{
						weight = g_weights[img * translation_num + itrans];

						if (weight >= significant_weight)
						{
							weight = weight * inv_minsigma_ctf;
							if(CTF_PREMULTIPLIED)
								Fweight[tid] += weight * ctf * ctf;
							else
								Fweight[tid] += weight * ctf;

							if(DATA3D)
								CpuKernels::translatePixel(x, y, z, g_trans_x[itrans], g_trans_y[itrans], g_trans_z[itrans], img_real, img_imag, temp_real, temp_imag);
							else
								CpuKernels::translatePixel(x, y,    g_trans_x[itrans], g_trans_y[itrans],                    img_real, img_imag, temp_real, temp_imag);

							real[tid] += (temp_real-ref_real[tid]) * weight;
							imag[tid] += (temp_imag-ref_imag[tid]) * weight;
						}
					}

					//BP
					if (Fweight[tid] > (XFLOAT) 0.0)
					{
						// Only asymmetric half is stored
						if (xp[tid] < (XFLOAT) 0.0)
						{
							// Get complex conjugated hermitian symmetry pair
							xp[tid] = -xp[tid];
							yp[tid] = -yp[tid];
							zp[tid] = -zp[tid];
							imag[tid] = -imag[tid];
						}
					} // if (Fweight[tid] > (XFLOAT) 0.0)
				} //if(ok_for_next)
			} // for tid

			for(int tid=0; tid<block_size; tid++)
			{
				if (Fweight[tid] > (XFLOAT) 0.0)
				{
					int x0 = floorf(xp[tid]);
					XFLOAT fx = xp[tid] - x0;
					int x1 = x0 + 1;

					int y0 = floorf(yp[tid]);
					XFLOAT fy = yp[tid] - y0;
					y0 -= mdl_inity;
					int y1 = y0 + 1;

					int z0 = floorf(zp[tid]);
					XFLOAT fz = zp[tid] - z0;
					z0 -= mdl_initz;
					int z1 = z0 + 1;

					XFLOAT mfx = (XFLOAT)1.0 - fx;
					XFLOAT mfy = (XFLOAT)1.0 - fy;
					XFLOAT mfz = (XFLOAT)1.0 - fz;

					XFLOAT dd000 = mfz * mfy * mfx;
					XFLOAT dd001 = mfz * mfy *  fx;

					size_t z0MdlxMdly = (size_t)z0 * (size_t)mdl_x * (size_t)mdl_y;

					{
						tbb::spin_mutex::scoped_lock lock(mutexes[z0 * mdl_y + y0]);

						g_model_real  [z0MdlxMdly + y0 * mdl_x + x0] += dd000 * real[tid];
						g_model_imag  [z0MdlxMdly + y0 * mdl_x + x0] += dd000 * imag[tid];
						g_model_weight[z0MdlxMdly + y0 * mdl_x + x0] += dd000 * Fweight[tid];

						g_model_real  [z0MdlxMdly + y0 * mdl_x + x1] += dd001 * real[tid];
						g_model_imag  [z0MdlxMdly + y0 * mdl_x + x1] += dd001 * imag[tid];
						g_model_weight[z0MdlxMdly + y0 * mdl_x + x1] += dd001 * Fweight[tid];
					}

					XFLOAT dd010 = mfz *  fy * mfx;
					XFLOAT dd011 = mfz *  fy *  fx;

					{
						tbb::spin_mutex::scoped_lock lock(mutexes[z0 * mdl_y + y1]);

						g_model_real  [z0MdlxMdly + y1 * mdl_x + x0] += dd010 * real[tid];
						g_model_imag  [z0MdlxMdly + y1 * mdl_x + x0] += dd010 * imag[tid];
						g_model_weight[z0MdlxMdly + y1 * mdl_x + x0] += dd010 * Fweight[tid];

						g_model_real  [z0MdlxMdly + y1 * mdl_x + x1] += dd011 * real[tid];
						g_model_imag  [z0MdlxMdly + y1 * mdl_x + x1] += dd011 * imag[tid];
						g_model_weight[z0MdlxMdly + y1 * mdl_x + x1] += dd011 * Fweight[tid];
					}

					XFLOAT dd100 =  fz * mfy * mfx;
					XFLOAT dd101 =  fz * mfy *  fx;

					size_t z1MdlxMdly = (size_t)z1 * (size_t)mdl_x * (size_t)mdl_y;

					{
						tbb::spin_mutex::scoped_lock lock(mutexes[z1 * mdl_y + y0]);

						g_model_real  [z1MdlxMdly + y0 * mdl_x + x0] += dd100 * real[tid];
						g_model_imag  [z1MdlxMdly + y0 * mdl_x + x0] += dd100 * imag[tid];
						g_model_weight[z1MdlxMdly + y0 * mdl_x + x0] += dd100 * Fweight[tid];

						g_model_real  [z1MdlxMdly + y0 * mdl_x + x1] += dd101 * real[tid];
						g_model_imag  [z1MdlxMdly + y0 * mdl_x + x1] += dd101 * imag[tid];
						g_model_weight[z1MdlxMdly + y0 * mdl_x + x1] += dd101 * Fweight[tid];

					}

					XFLOAT dd110 =  fz *  fy * mfx;
					XFLOAT dd111 =  fz *  fy *  fx;

					{
						tbb::spin_mutex::scoped_lock lock(mutexes[z1 * mdl_y + y1]);

						g_model_real  [z1MdlxMdly + y1 * mdl_x + x0] += dd110 * real[tid];
						g_model_imag  [z1MdlxMdly + y1 * mdl_x + x0] += dd110 * imag[tid];
						g_model_weight[z1MdlxMdly + y1 * mdl_x + x0] += dd110 * Fweight[tid];

						g_model_real  [z1MdlxMdly + y1 * mdl_x + x1] += dd111 * real[tid];
						g_model_imag  [z1MdlxMdly + y1 * mdl_x + x1] += dd111 * imag[tid];
						g_model_weight[z1MdlxMdly + y1 * mdl_x + x1] += dd111 * Fweight[tid];
					}

				} // Fweight[tid] > (RFLOAT) 0.0
			} // for tid
		} // for pass
	} // for img
}

} // namespace
