#ifndef __DRESIDUALFINDER_CU__
#define __DRESIDUALFINDER_CU__

#include <GPUDistToRT.cu>
#include <GPUResidualFinder.h>
#include <GPUMagneticFieldStepper.h>

//----------------
// GPU_FindClosestTrajectoryPoint
//----------------
__global__ void GPU_FindClosestTrajectoryPoint(int N, int M, trajectory_t *trajectories, int *wire_hit_d, int *hit_index_d)
{
	//int idx_i = threadIdx.x;  // hit (N)
	//int idx_j = blockIdx.x;  // trajectory (M)
	int idx_i = blockIdx.x; //N
	int idx_j = threadIdx.x; //M

	if(idx_i<N && idx_j<M){

		// Get pointers to trajectory and wire structures
		// this thread is to find closest step for
		trajectory_t *trajectory = &trajectories[idx_j];
		int wireIndex = wire_hit_d[idx_i];
		Wire *wire = &wires_d[wireIndex]; // get pointer to wire in device constant memory

		
		float vals[2] = {99999.0, 99999.0};
		int ptr[2] = {0, 0}; //index of trajectory points
		step_t *step = trajectory->steps;
		for(int i = 0; i < trajectory->Nsteps; i++, step++)
		{
			// Distance from point to wire is magnitude of
			// cross product of vector pointing from wire
			// center to point and the normalized vector 
			// pointing in direction of wire.
			float x = step->pos.x - wire->x; 
			float y = step->pos.y - wire->y;
			float z = step->pos.z - wire->z;
			float3 b = make_float3(x, y, z);
			float3 &wire_dir = wire->udir;
			
			float3 cross_prod = cross(b, wire_dir);

			// Only care about minimum magnitude which is
			// also the minimum magnitude squared
			float val2 = cross_prod.x*cross_prod.x + cross_prod.y*cross_prod.y + cross_prod.z*cross_prod.z;

			// Use boolean result of less than operator as
			// an integer to index which element of the two
			// element arrays to store this in. This way, the
			// minimum is always put into element 1.
			int idx = val2 < vals[1]; //1 or 0 based on min or not

			vals[idx] = val2; //Set new min
			ptr[idx] = i;
		} //min distance point should be in ptr[1]

		int finalIndex = N*idx_j + idx_i;
//ptr[1]=(int)(sqrt(vals[1])*1000.0);
//ptr[1] = finalIndex;
		hit_index_d[finalIndex] = ptr[1];

	}
}

//----------------
// GPU_CalcResiduals
//----------------
__global__ void GPU_CalcResiduals(int N, int M, trajectory_t *trajectories, int *wire_hit_d, int *hit_index_d, float *hit_dist_d, float *hit_resi_d)
{
   //int idx_i = threadIdx.x;  // hit (N)
   //int idx_j = blockIdx.x;  // trajectory (M)
	int idx_i = blockIdx.x; //N
	int idx_j = threadIdx.x; //M
	
	if(idx_i<N && idx_j<M){

		// Get pointers to trajectory and wire structures
		// this thread is to find closest step for
		trajectory_t *trajectory = &trajectories[idx_j];
		int wireIndex = wire_hit_d[idx_i];
		Wire *wire = &wires_d[wireIndex]; // get pointer to wire in device constant memory

		int idx = N*idx_j + idx_i;
		int step_idx = hit_index_d[idx];
		step_t *step = &trajectory->steps[step_idx];

		hit_dist_d[idx] = GPU_DistToRT(wire, step);
		hit_resi_d[idx] = hit_dist_d[idx];
	}
}

//----------------
// GPU_CalcMatrix_Binv
//----------------
__global__ void GPU_CalcMatrix_Binv(int N, int M, trajectory_t *trajectories, int *wire_hit_d, float *hit_resi_d, float *Binv_matrix_d)
{
	/// Finding the change in the state vector requires calculating
	/// the inverse of the "B" matrix and then inverting it to get "B".
	/// The "Binv" matrix is made from the "F" matrix. "F" is just the
	/// Jacobian dR_i/dS_j where R_i is the change in the residual of
	/// the i-th hit and S_j is the corresponding change of the j-th state
	/// parameter. The uncertainties are included through an inverted
	/// covariance matrix "Vinv" (see below). The formula for "Binv" is:
	///
	///     Binv = Ft*Vinv*F
	///
	/// Each thread running this routine will calculate a single
	/// element of the Binv matrix. In reality, there should be 
	/// 25 threads per Binv and one Binv matrix for every mass
	/// hypothesis.
	///
	/// The covariance "V" is the sum of the measurement uncertainties
	/// and the multiple scattering covariance. The measurement
	/// uncertainties are a diagonal matrix and the multiple
	/// scattering covariance can contain off-diagnal elements. In
	/// the current version, the multiple scattering is not included.
	/// This means that V is a diagonal matrix. Further, we do not
	/// distinguish between CDC and FDC and assume all measurements
	/// have the same uncertainty. This makes the inversion of the 
	/// V matrix trivial. It is just an identity matrix scaled by
	/// 1/sigma^2. If this were not the case, then this kernel should
	/// be broken up into 2 so that the different dimensioned matrix
	/// multiplications could be done in two steps. As it is now
	/// though, this reduces to Ft*F / sigma^2
	
	float err = 0.1443375673; // FDC sigma = 0.5/sqrt(12)
	float err2 = err*err;

	int idx_hypoth = blockIdx.x;
	int idx_i = threadIdx.x;
	int idx_j = threadIdx.y;
	if( idx_i<5 && idx_j<5 && idx_hypoth<N_MASS_HYPOTHESES){
	
		// Need the offset pointing to the start of the trajectories
		// used for the +ds and -ds parameter tweaks.
		
		// offset to the block of trajectories for this mass hypothesis
		int idx_hypoth_off = idx_hypoth*N_TRAJECTORIES_PER_HYPOTHESIS*N;
		
		// offset to trajectory for state parameter i. The
		// "+1" is because the first trajectory corresponds
		// to the non-tweaked state parameters
		int idx_traj_i_plus = idx_hypoth_off + (1 + 2*idx_i)*N;
		int idx_traj_i_minus = idx_traj_i_plus + N; // -ds always follows +ds trajectory
		
		// Similarly for state parameter j
		int idx_traj_j_plus = idx_hypoth_off + (1 + 2*idx_j)*N;
		int idx_traj_j_minus = idx_traj_j_plus + N; // -ds always follows +ds trajectory

		// Calculate single element in Ft*F 
		float sum = 0.0;
		for(int i=0; i<N; i++){
			float diff_i = hit_resi_d[idx_traj_i_plus + i] - hit_resi_d[idx_traj_i_minus + i];
			float diff_j = hit_resi_d[idx_traj_j_plus + i] - hit_resi_d[idx_traj_j_minus + i];

			diff_i /= STATE_DS[idx_i];
			diff_j /= STATE_DS[idx_j];

			sum += diff_i*diff_j;
		}

		// Copy value into Binv memory and scale by 1/sigma^2 for "Vinv"
		int idx = idx_i + 5*idx_j;
		Binv_matrix_d[idx] = sum/err2;
	}
}

//----------------
// GPU_Calc_dSdR
//----------------
__global__ void GPU_Calc_dSdR(int N, int M, float *hit_resi_d, float *B_matrix_d, float *dSdR_matrix_d)
{
	/// The change in the state parameters is calculated by the following:
	///
	///    B*Ft*Vinv*R
	///
	///  where "B", "Ft", and "Vinv" are all as described above in 
	/// GPU_CalcMatrix_Binv(). The "R" vector is just the vector of
	/// residuals from the un-tweaked parameters. Therefore, the matrix
	///
	///    dSdR = B*Ft*Vinv
	///
	/// can be thought of as the change in state parameter per unit
	/// residual. We use the same shortcut for Vinv as described above
	/// so that this kernel need only calculate an element of the matrix
	/// B*Ft and then scale it by 1/sigma^2 to make it dSdR. A separate
	/// kernel will be used for the final multiplication leading to a
	/// 5 element array containing the change in state parameters.

	float err = 0.1443375673; // FDC sigma = 0.5/sqrt(12)
	float err2 = err*err;

	int idx_hypoth = blockIdx.x;
	int idx_i = threadIdx.x; // N
	int idx_j = threadIdx.y; // n
	if( idx_i<N && idx_j<5 && idx_hypoth<N_MASS_HYPOTHESES){
	
		// Need the offset pointing to the start of the trajectories
		// used for the +ds and -ds parameter tweaks.
		
		// offset to the block of trajectories for this mass hypothesis
		int idx_hypoth_off = idx_hypoth*N_TRAJECTORIES_PER_HYPOTHESIS*N;
		
		// offset to trajectory for state parameter i. The
		// "+1" is because the first trajectory corresponds
		// to the non-tweaked state parameters
		int idx_traj_j_plus = idx_hypoth_off + (1)*N + idx_i;
		int idx_traj_j_minus = idx_traj_j_plus + N; // -ds always follows +ds trajectory

		// Calculate single element in B*Ft
		float sum = 0.0;
		for(int i=0; i<5; i++){
			float diff_j = hit_resi_d[idx_traj_j_plus + 2*i*N] - hit_resi_d[idx_traj_j_minus + 2*i*N];

			diff_j /= STATE_DS[i];   // Ft

			sum += diff_j*B_matrix_d[idx_j*5 + i];
		}

		// Copy value into Binv memory and scale by 1/sigma^2 for "Vinv"
		int idx = idx_i + N*idx_j;
		dSdR_matrix_d[idx] = sum/err2;
	}
	

}

//----------------
// GPU_Calc_dS
//----------------
__global__ void GPU_Calc_dS(int N, int M, float *hit_resi_d, float *dSdR_matrix_d, float *dS_matrix_d)
{
	int idx_hypoth = blockIdx.x;
	int idx_j = threadIdx.y; // n
	if( idx_j<5 && idx_hypoth<N_MASS_HYPOTHESES){
	
		// offset to the block of trajectories for this mass hypothesis
		// Since we need the residuals from the un-tweaked trajectory
		// and that happens to be te first one, then this is also the 
		// offset to the trajectory of interest.
		int idx_hypoth_off = idx_hypoth*N_TRAJECTORIES_PER_HYPOTHESIS*N;

		int idx_traj_j_plus = idx_hypoth_off + (0)*N; // need untweaked residual vector here
	
		float sum = 0.0;
		for(int i=0; i<N; i++){
			sum += hit_resi_d[idx_traj_j_plus + i] * dSdR_matrix_d[i + N*idx_j];
		}

		dS_matrix_d[idx_j] = sum;
	}
}
/*
//----------------
// GPU_CalcChiSqPerNdof
//----------------
__global__ void GPU_CalcChiSqPerNdof(int N, int M, float *hit_resi_d, float *sum_resi_d)
{
	int idx_j = threadIdx.x;
	if(idx_j < M){
		float wireSum = 0;
		for(int i = 0; i < N; i++) //loop through hits
		{
			// Add residual for hit i to trajectory j
			float v = hit_resi_d[idx_j*N + i];
			float err = 0.1443375673; // FDC sigma = 0.5/sqrt(12)
			v /= err;
			wireSum += v*v;
		} 
		
		// Number of degrees of freedom is number of hits
		// minus number of fit parameters
		double Ndof = (double)N - 5.0; 
		
		sum_resi_d[idx_j] = wireSum/Ndof; // record chi-sq/Ndof
	}
}
*/


#endif  // __DRESIDUALFINDER_CU__