#include "peg_pic.h"

void reinj(int *, int, int *);
void moveLastParticle( int isp, int trg, int also_c_omt );
void copyParticleSpecies(int isp, int trg);

/****************************************************************/
/* Routine to adjust (initialize, re-pack and inject) particles */

void adjust() {
	static int init_flag=1;
	int i, ii, isp, k, s;
	static float **extra;       
	float dum, vx_t, vy_t, vz_t, sigmd; 
	int realNsp;

#ifdef REFLECTED_OWN_SPECIES
	/* real number of species (we have double the number because of reflection storage) */
	realNsp = (int) nsp / 2;
#else
	realNsp = nsp;
#endif

	if (init_flag)  {
		i=0;

		for (isp=0; isp<nsp; isp++) {
			if (inject[isp]) {
				/* check if we have to do some injection */
				i +=1;
			}

			/* define initial values for grid reflection on left and right side */
			gridref_l[isp] = gridref_l0[isp];
			gridref_r[isp] = gridref_r0[isp];
		}

		if (i) {
			/* yes, we have to do some injection; allocate space */
			extra= (float **)malloc(nsp*sizeof(float *));

			for (isp=0; isp<nsp; isp++) {
				extra[isp]= (float *)malloc(2*sizeof(float));
			}
		}

		for (isp=0; isp<nsp; isp++) {
			if (inject[isp]) {
				if (fabs(enter[isp][0]) > 0.) {
					extra[isp][0] = 0.5123123;
				} else {
					extra[isp][0] = 0.;
				}

				if (fabs(enter[isp][1]) > 0.) {
					extra[isp][1] = 0.5123123;
				} else {
					extra[isp][1] = 0.;
				}
			}

			if (sourcD) {
				Stemp[isp]=0.;                  /* for source */
			}
		}
		init_flag = 0;
	}

	for (isp=0; isp<nsp; isp++) {
		if (dinj[isp]) {
			jwall[isp] =0.0;   

			if (slow) {          /* initialize f() diagnostic arrays */
				for (s=0; s<NDP; s++) {
					for (ii=0; ii<nbin[isp]; ii++) {
						fe_mid[isp][s][ii] = 0; 
					}
					
					for (ii=0; ii<nbin_fvx[isp]; ii++) {
						f_vpar[isp][s][ii] =0;
					}
				} 
			}
		}

		/* calculate current values for grid reflection on left and right side */
		gridref_l[isp] = gridref_l0[isp] + gridref_lsramp[isp] * t;
		gridref_r[isp] = gridref_r0[isp] + gridref_rsramp[isp] * t;
	}

	/*
	 * Check boundaries for reflection and absorption of particles.
	 */
	for (isp=0; isp<nsp; isp++) {
		if ((np[isp] > 0) && (dinj[isp])) {
			i = 0;

			do {
				/* eliminate "outsiders" and inject secondaries */
				if (x[isp][i] > xnc) {        /* particle crossed the RHS boundary */
					if (i_t[isp]==2) {                  /* inclined mag field: outsiders */
						if ((vx[isp][i]*cs+vz[isp][i]*sn)>0.) {
							reinj(&i, isp, &np[isp]);     
						} else {
							if (slow) {
								dum = Ep(isp, i);
								walld(isp, dum, 0., 1);           /* wall diagnostics */
							} else if (Ef_flag) {
								F_tempRHS[isp] +=1;
								Efrhs_temp[isp] += Ep(isp, i);
							} else {
								F_tempRHS[isp] +=1;
							}

							/* copy last particle into current cell and do the checking again */
							moveLastParticle(isp, i, 1);
						}
						/* end of if (i_t[isp]==2) */
					} else {
						if (slow) {
							dum = Ep(isp, i);

							if (srf[isp][1]) {                        /* surface physics */
								sigmd = sqrt(dum)/(vx[isp][i]+DBL_MIN); /* = 1/cosQ */
								surfphys(isp, sigmd, dum, 1);
							}

							walld(isp, dum, 0., 1);           /* wall diagnostics */

							/*  end of if (slow...) */
						} else if (Ef_flag) {
							F_tempRHS[isp] +=1;
							dum = Ep(isp, i);
							Efrhs_temp[isp] += dum;

							if (srf[isp][1]) {                         /* surface physics */
								sigmd = sqrt(dum)/(vx[isp][i]+DBL_MIN);
								surfphys(isp, sigmd, dum, 1);
							}
						} else if (srf[isp][1]) {                 /* surface physics */
							F_tempRHS[isp] +=1;
							dum = Ep(isp, i);
							sigmd = sqrt(dum)/(vx[isp][i]+DBL_MIN);
							surfphys(isp, sigmd, dum, 1);
						} else  {
							F_tempRHS[isp] +=1;
						}

						if (reflux < 2) {                     /* removing of particle */
							/* before removing, check if we bounce back the particle (depends on reflection coefficient) */
							if (frand() > gridref_r[isp]) {
								/* on RIGHT boundary: particle is absorbed */
								/* copy last particle into current cell and do the checking again */
								moveLastParticle(isp, i, 0);
							} else {
								/* particle is reflected back from boundary */
								x[isp][i] = 2*xnc - x[isp][i];
								vx[isp][i] = -vx[isp][i];

#ifdef REFLECTED_OWN_SPECIES
								if( isp < realNsp) {
									/* do this moving of species index only for "real" species (with indexes under realNsp) */
									/* store particle to different index (because it was reflected); all other parameters (m, q, v, x) are the same */
									copyParticleSpecies(isp, i);
									/* copy last particle into current cell and do the checking again */
									moveLastParticle(isp, i, 0);
								}
#endif
							}

						} else if (reflux==2) {                         /* reflection */
							x[isp][i] = 2*xnc - x[isp][i];
							vx[isp][i] = -vx[isp][i];

#ifdef REFLECTED_OWN_SPECIES
							if( isp < realNsp) {
								/* do this moving of species index only for "real" species (with indexes under realNsp) */
								/* store particle to different index (because it was reflected); all other parameters (m, q, v, x) are the same */
								copyParticleSpecies(isp, i);
								/* copy last particle into current cell and do the checking again */
								moveLastParticle(isp, i, 0);
							}
#endif

						} else if (reflux==3) {     /* reinjection from other boundary */
							x[isp][i] -= xnc;
						}

					}   /*  end of else */
					/* end of if (x > xnc) */

				} else if (x[isp][i] < 0) {  /* particle crossed the LHS boundary */
					if (slow) {
						dum = Ep(isp, i);
						sigmd = sqrt(dum)/(-vx[isp][i]+DBL_MIN);

						if (srf[isp][0])  {                        /* surface physics */
							surfphys(isp, sigmd, dum, 0);
						}

						walld(isp, dum, sigmd, 0);            /* wall diagnostics */

						/* end of slow */
					} else if (Ef_flag) {
						F_temp[isp] +=1;
						dum = Ep(isp, i);
						Eflhs_temp[isp] += dum;

						if (srf[isp][0]) {                        /* surface physics */
							sigmd = sqrt(dum)/(-vx[isp][i]+DBL_MIN);
							surfphys(isp, sigmd, dum, 0);
						}
					} else if (srf[isp][0]) {                   /* surface physics */
						F_temp[isp] +=1;
						dum = Ep(isp, i);
						sigmd = sqrt(dum)/(-vx[isp][i]+DBL_MIN);
						surfphys(isp, sigmd, dum, 0);
					} else  {
						F_temp[isp] +=1;
					}

					if (!reflux) {                      /* removing of particle */  
						/* before removing, check if we bounce back the particle (depends on reflection coefficient) */
						if (frand() > gridref_l[isp]) {
							/* on LEFT boundary: particle is absorbed */
							jwall[isp] += q[isp];
							/* copy last particle into current cell and do the checking again */
							moveLastParticle(isp, i, 0);

						} else {
							/* particle is reflected back from boundary */
							x[isp][i] = -x[isp][i];
							vx[isp][i] = -vx[isp][i];

#ifdef REFLECTED_OWN_SPECIES
							if( isp < realNsp) {
								/* do this moving of species index only for "real" species (with indexes under realNsp) */
								/* store particle to different index (because it was reflected); all other parameters (m, q, v, x) are the same */
								copyParticleSpecies(isp, i);
								/* copy last particle into current cell and do the checking again */
								moveLastParticle(isp, i, 0);
							}
#endif
						}
					} else if (reflux < 3) {                      /* reflection  */
						x[isp][i] = -x[isp][i];
						vx[isp][i] = -vx[isp][i];

#ifdef REFLECTED_OWN_SPECIES
						if( isp < realNsp) {
							/* do this moving of species index only for "real" species (with indexes under realNsp) */
							/* store particle to different index (because it was reflected); all other parameters (m, q, v, x) are the same */
							copyParticleSpecies(isp, i);
							/* copy last particle into current cell and do the checking again */
							moveLastParticle(isp, i, 0);
						}
#endif

					} else if (reflux == 3)  { /* reinjection from other boundary */
						x[isp][i] += xnc;
					}

					/* end of if (x < 0) */	
				} else {
					if (slow) {  /* "Middle" diagnostics */ 
						for (k=0; k<NDP; k++) {
							if (xs_mid[k] <= x[isp][i] && x[isp][i] <= xf_mid[k]) {
								// Energy distribution
								s = (Ep(isp, i) - emin[isp])/de[isp]; 

								if (0 <= s && s< nbin[isp]) {
									fe_mid[isp][k][s] += 1;
								}

								// Velocity distribution
								if (bmag > 0.0) {
									dum = cs*vx[isp][i] + sn*vz[isp][i];
									s = (dum - vx_min[isp])/de_fvx[isp];
								} else {
									s = (vx[isp][i] - vx_min[isp])/de_fvx[isp];
								}

								if (0 <= s && s< nbin_fvx[isp]) {
									f_vpar[isp][k][s] += 1;
								}

								break;
							}
						} /* end of the for (s=0; s<NDP; s++) loop */
					} /* end of slow */

					if (i_t[isp]==2) {               /* particle came back */
						c_omt[isp][i]=0;
					}

					if (ht_sp[isp]>0.) {  /* heating of particles in the source region */ 
						if ((x[isp][i]>=nc2)&&(x[isp][i]<=nc3)) {
							if (!ht1) {
								if ((frand()>heat[isp])) {
									k= 1;

									if (bmag > 0.0) {
										if (vx[isp][i]*cs + vz[isp][i]*sn < 0.0)  {
											k= -1;
										}
									} else if (vx[isp][i] < 0.0) {
										k= -1;
									}

									heatup(isp, i, k); 
								}
							} else {
								if ((x[isp][i]>=nc21)&&(x[isp][i]<=nc31)) {
									if ((frand()>heat1[isp])) {
										k= 1;

										if (bmag > 0.0) {
											if (vx[isp][i]*cs + vz[isp][i]*sn < 0.0) {
												k= -1;
											}
										} else if (vx[isp][i] < 0.0) {
											k= -1;
										}

										heatup1(isp, i, k);
									}
								} else if ((frand()>heat[isp])) {
									k= 1;

									if (bmag > 0.0) {
										if (vx[isp][i]*cs + vz[isp][i]*sn < 0.0){
											k= -1;
										}
									} else if (vx[isp][i] < 0.0) {
										k= -1;
									}

									heatup(isp, i, k);
								}
							}
						}  
					}   // end of if (ht_sp[isp]>0.) loop

					/* advance to next particle in current species */
					i++;

				}  /* end of else */
			} while( i < np[isp] );

		}   /* end of "if(np[isp]>0...)" loop */
	}     /* end of "for( isp..)"  loop */

	/*
	 * Now INJECT new particles at walls.
	 */
	for (isp=0; isp<realNsp; isp++) {
//printf("inject = %d dinj = %d  i_t=%d", inject[isp], dinj[isp], i_t[isp]);

		if ((inject[isp]) && (dinj[isp])) {   
			for(k=0; k<2; k++) {
				extra[isp][k] += enter[isp][k]; 

//printf("  insert %d = %lf ", k, extra[isp][k]);

				while (extra[isp][k] >= 1.0) {
					if (k==0) {
						vmaxwell0(&vx_t, &vy_t, &vz_t, isp);
					} else {
						vmaxwell1(&vx_t, &vy_t, &vz_t, isp);
					}

//printf("---> new particle: (%f, %f, %f)\n", vx_t, vy_t, vz_t);

					if((i_t[isp]) && (k) && ((vx_t*cs+vz_t*sn)>0.)) {
						extra[isp][k] -= 1.0;  /* removing particles with v||<0  */
					} else {
//printf( "old len = %d .... ", np[isp]);

						ii = np[isp];
						np[isp]++;

						vx[isp][ii] = vx_t;
						vy[isp][ii] = vy_t;
						vz[isp][ii] = vz_t;
						x[isp][ii] = k*nc + frand()*nstep[isp]*vx_t;
//printf("k=%d, nc=%d, isp=%d, nstep[isp]=%d, vx_t=%lf --- x[isp][ii]=%lg\n", k, nc, isp, nstep[isp], vx_t, x[isp][ii]);
//						x[isp][ii] = length/10000;
//printf("  %d. x=%lf vx=%lf, vy=%lf, vz=%lf\n", ii, x[isp][ii], vx[isp][ii], vy[isp][ii], vz[isp][ii]);

//						printf("new len=%d, (%f, %f, %f)\n", ii, vx[isp][ii], vy[isp][ii], vz[isp][ii]);

						if (!k) {
							jwall[isp] -= q[isp];
						}

						extra[isp][k] -= 1.0;
					}
				}   /* end "while extra..)" loop */
			}   /* end "for k= 0,1" loop */
		}   /* end "if( inject..)"  loop */
	}  /* end of particle INJECTION at walls */

	for (isp=0; isp<realNsp; isp++) {
		if (dinj[isp]) {
			jwall[isp] /= (nstep[isp]*dt);
		}
	}

	if (sourcD) {
		psrc();                    /* First particle source */     
	}

	if (sourcD1) {
		psrc1();                   /* Second particle source */
	}

}    /* END of ADJUST */


/*************************************************************/
/****  Inclined mag field: outsiders that are not removed ****/

void reinj(int *j, int iisp, int *inp) {
	int i, nnp;
	float dum;

	i= *j;
	nnp= *inp -1;
	c_omt[iisp][i]++;

	if (c_omt[iisp][i] > count_max[iisp]) {                             /* if t_out>T_cyclotron, remove particle */
		x[iisp][i] = x[iisp][nnp];
		vx[iisp][i] = vx[iisp][nnp];
		vy[iisp][i] = vy[iisp][nnp];
		vz[iisp][i] = vz[iisp][nnp];
		c_omt[iisp][i]=c_omt[iisp][nnp];
		*inp -= 1;
	} else {
		if (c_omt[iisp][i]==1) {            /* if particle crossed the RHB */
			if (slow) {
				dum = Ep(iisp, i);
				walld(iisp, dum, 0., 1);
			} else if (Ef_flag) {
				F_temp[iisp] +=1;
				Eflhs_temp[iisp] += Ep(iisp, i);
			} else  {
				F_temp[iisp] +=1;
			}
		}

		*j += 1;
	}
}     /* END of reinj() */

/*************************************************************/

/*
 * Movie LAST particle for given species to target ("trg") index.
 * Also reduce number of particles of given species.
 */
void moveLastParticle( int isp, int trg, int also_c_omt ) {
	int src = np[isp]-1;
	x[isp][trg] = x[isp][src];
	vx[isp][trg] = vx[isp][src];
	vy[isp][trg] = vy[isp][src];
	vz[isp][trg] = vz[isp][src];
	np[isp] -= 1;

	if( also_c_omt ) {
		c_omt[isp][trg] = c_omt[isp][src];
	}

}

/*
 * Move particle in species "isp" and with index "src" to "new" species to the end of array.
 * Do NOTHING with moved particle.
 */
void copyParticleSpecies(int isp, int src) {
	int realNsp = (int) nsp / 2;
	int new_isp = isp+realNsp;
	int newPos = np[new_isp]; 	/* current number of particles in target species */

	/* "copy" particle in new species index and increase number of particles */
	x[new_isp][newPos] = x[isp][src];
	vx[new_isp][newPos] = vx[isp][src];
	vy[new_isp][newPos] = vy[isp][src];
	vz[new_isp][newPos] = vz[isp][src];
//	c_omt[new_isp][newPos] = c_omt[isp][src];
	np[new_isp] ++;
}
