// Calculation of LVIS scattering due to up to 2 counterpropagating pairs of perpendicular molasses beams
// C++ ROOT script
// Author: Andres Cimmarusti and David Norris
//	   Joint Quantum Institute
//	   University of Maryland
//	   College Park, MD 20742
// Last update: 08/17/2010

/*
TODO:

- I have included the calculation of the trajectories and velocities of the atoms prior to reaching the molasses beams. These trajectories assume no collisions between atoms or any interaction whatsoever... although it doesn't seem to be very important

- So far gravity neglected, though it should be fairly easy to implement.

- I have assumed the intersection of the molasses beams is a sphere. This is not accurate because it should be the intersection of two cylinders, but it seems like a fair approximation ("..assume the cow is a sphere..").

- Molasses beams are assumed to be perfect plane waves. This is not accurate. They all have a gaussian beam intensity profile that should be taken into account. This should allow to implement a laser lens to provide a position dependent force to collimate the LVIS.

- Potential performance improvements: The velocities of the atoms only vary in the molasses region. Perhaps there is an easier way to keeping the velocity constant in the vector container without having to push back all the time.

- From several tests I've run it is quite clear that the molasses stage creates the most elements (dt shortest). Improvements in other stages are of little help.

*/

#include <iostream>
#include <vector>
#include <ctime>
#include "definitions.h"


void lvis_opt_mol()
{
	//------ROOT libraries------//

	gSystem->Load( "libPhysics" );
	gSystem->Load( "libMathCore" );


	//------Calculations and conversions------//

	d0 = 2 * winrad;							// Distance to LVIS hole center(FIXME: should vary with angle)

	tube = molchamrad - d0;							// In vacuo tube length 

	dvt0 = hole / ( tube + motchamsize / 2 ) * v0; 				// LVIS initial transverse speed width

	matom = nucleon_num * m_nuc;						// Mass of an atom

	thetam = atan( molchamrad / winrad );					// Polar angle corresponding to chamber window aperture


	//------Recoil velocity parameters------//

	vrec = h / ( matom * lam );						// Magnitude

	vrec_abs.SetTheta( pi / 2 );						// Absorption recoil velocity

	vrec_abs.SetMag ( vrec );

	vrec_spon.SetMag ( vrec );						// Spontaneous emission recoil velocity


	//------Initial velocity spread parameters------//

	v0dev = dv0 / ( 2 * sqrt( 2 * log( 2 ) ) );				// Standard deviation LVIS initial speed

	dtheta0 = ( dvt0 - dv0 * sin( theta0 ) ) / ( v0 * cos( theta0 ) );	// LVIS initial polar angle spread

	theta0dev = dtheta0 / ( 2 * sqrt( 2 * log( 2 ) ) );			// Standard deviation LVIS initial polar angle


	//------Position of LVIS hole center------//

	r0.SetTheta( pi - theta0 );

	r0.SetPhi( pi + phi0 );

	r0.SetMag( d0 );


	//------Initialization and inputs------//

	vector<Double_t> time_atom;

	vector<Double_t> x;

	vector<Double_t> y;

	vector<Double_t> z;

	vector<Double_t> vx;

	vector<Double_t> vy;

	vector<Double_t> vz;

	Double_t *w = new Double_t [nbeams];		// Beam spot size

	Double_t *w0 = new Double_t [nbeams];		// Beam waist

	Double_t *zr = new Double_t [nbeams];		// Rayleigh length

	Double_t *phasecur = new Double_t [nbeams];	// Phase curvature

	Double_t *guoy = new Double_t [nbeams];		// Guoy phase shift

	Double_t *detun = new Double_t [nbeams];	// Detunings

	Double_t *gs0 = new Double_t [nbeams];		// Gaussian beam on-resonance saturation parameter

	Double_t *s0 = new Double_t [nbeams];		// On-resonance saturation parameter ( I / Isat ) at waist

	Double_t *s = new Double_t [nbeams];		// Saturation parameter

	Double_t *sr = new Double_t [nbeams];		// Scattering rates


	//------Wavevector and other beam parameters------//

	if ( nbeams == 0 ) cout << "No molasses beams" << endl;

	for ( i = 0 ; i < nbeams ; i ++ )
	{
		k[i].SetX( 1 );

		k[i].SetMag( 1 / lam );

		k[i].SetTheta( pi / 2 );

		if ( i < 2 )
		{
			k[i].SetPhi( i * pi );

			pol[i].SetXYZ( 0 , 1 , 0 );
		}
		else
		{
			k[i].SetPhi( ( 2 * i - 3 ) * pi / 2 );

			pol[i].SetXYZ( 1 , 0 , 0 );
		}

		w0pos[i] = - k[i].Unit();

		w0pos[i].SetMag( 0.05 );

		rotz[i].RotateZ( k[i].Phi() + 3 * pi / 2 ); 		// Preserve order, first rotate around z then around x

		rotx[i].RotateX( k[i].Theta() );

		//cout << "Enter on-resonance saturation parameter ( I / Isat ) for beam " << i + 1 << " : ";
		//cin >> s0[i];

		s0[i] = 3;

		//cout << "Enter detuning (number of linewidths) for beam " << i + 1 << " : ";
		//cin >> detun[i];

		detun[i] = 2;

		detun[i] *= gam;

		w0[i] = 2e-6;

		zr[i] = pi * pow ( w0[i] , 2 ) / lam;

		width += w0[i] * sqrt( 1 + pow( w0pos[i].Mag() / zr[i] , 2 ) ) / nbeams;
/*
		cout << "molasses width of beam " << i + 1 << " : " << w0[i] * sqrt( 1 + pow( w0pos[i].Mag() / zr[i] , 2 ) ) << endl;

		cout << "polarization beam " << i + 1 << " : ( " << pol[i].X() << " " << pol[i].Y() << " " << pol[i].Z() << " )" << endl;

		cout << "wavevector beam " << i + 1 << " : ( " << k[i].X() << " " << k[i].Y() << " " << k[i].Z() << " )" << endl;

		cout << "waist position beam " << i + 1 << " : ( " << w0pos[i].X() << " " << w0pos[i].Y() << " " << w0pos[i].Z() << " )" << endl;*/
	}

	cout << "avg molasses beam spot size : " << width << endl;

	time ( & start ) ;


	//------Individual trajectory per atom------//

	for ( j = 0 ; j < natoms ; j ++ )
	{
		t = 0;

		randcont1 = randnum1.Rndm();

		randcont2 = randnum2.Rndm();


		//------Initial position spread in tube reference frame------//

		rini.SetPerp( randcont2 * hole / 2 );

		rini.SetPhi( randcont1 * 2 * pi );

		rini.SetZ( 0 );


		//------Transformation to center chamber reference frame------//

		rini.RotateY( - theta0 );

		rini.RotateZ( - phi0 );

		r = r0 + rini;


		//------Settings to avoid molasses if nbeams = 0------//

		if ( nbeams == 0 )
		{
			atom_out = 1;

			phi0 = randcont1 * 2 * pi;
		}
		else

			atom_out = 0;


		//------Atom initial velocity------//

		v.SetPhi( phi0 );

		v.SetTheta( randnum1.Gaus( theta0 , theta0dev ) );

		v.SetMag( randnum2.Gaus( v0 , v0dev ) );

		rmol = r - cmol;					// Change to molasses frame of reference

		dtp[0] = TMath::Abs( ( rmol.Mag() - width ) / ( v.Dot( rmol.Unit() ) ) / 100;

		if ( j < 5 ) cout << "coarse time step (premol) : " << dtp[0] << endl;


		//------LVIS propagation prior to molasses beams------//

		while ( rmol.Mag() > width && atom_out == 0 )
		{
			dist_atom_mol = rmol.Mag() - width;

			if ( j < nplot )
			{
				time_atom.push_back ( t );

				x.push_back ( r( 0 ) );

				y.push_back ( r( 1 ) );

				z.push_back ( r( 2 ) );

				vx.push_back ( v( 0 ) );

				vy.push_back ( v( 1 ) );

				vz.push_back ( v( 2 ) );
			}

			t += dtp[0];

			r += v * dtp[0];

			rmol = r - cmol;

			if ( rmol.Mag() - width > dist_atom_mol )

				atom_out = 1;
		}


		//------LVIS propagation through molasses beams------//

		while ( rmol.Mag() < width && atom_out == 0 )
		{
			tot_sr = 0;

			tot_s = 0;


			//------Random numbers: 1 next scattering time, 1 absorption probability, 2 reemission direction------//

			randnum1.RndmArray( 4 , prob );

			for ( i = 0 ; i < nbeams ; i ++ )
			{

				//------Gaussian beam correction to on-resonance saturation parameter------//

				rbeam = rmol - w0pos[i];

				rbeam.Transform( rotz[i] );

				rbeam.Transform( rotx[i] );

				w[i] = w0[i] * sqrt( 1 + pow( rbeam.Z() / zr[i] , 2 ) );

				gs0[i] = s0[i] * pow( w0[i] / w[i] , 2 ) * exp( - 2 * pow( rbeam.Perp() / w[i] , 2 ) );

				/* FIXME: this part seems insufficient, but I don't know how else to do it. I'm basically taking the intensity of each beam 					* independently. There is also the fact that the counterpropagating beams will set up a standing wave pattern that will
				* affect the overall intensity of the molasses, but if I do it as a whole, then I'm not able to tell which beam was 	
				* responsible for the "atom kick"
				*/


				//------Saturation parameter per beam------//

				s[i] = gs0[i] / ( 1 + pow( 2 * ( detun[i] + k[i].Dot( v ) ) / gam  , 2 ) );

				tot_s += s[i];
			}


			//------Scattering rate per beam------//

			for ( i = 0 ; i < nbeams ; i ++ )
			{
				sr[i] = 2 * pi * ( gam / 2 ) * s[i] / ( 1 + tot_s );

				tot_sr += sr[i];
			}


			//------Time to next scattering: Waiting time distribution------//

			dt = - log ( prob[0] ) / tot_sr;

			t += dt;


			//------From which molasses beam did the atom get a kick?------//

			if ( nbeams == 2 )
			{
				if ( prob[1] < sr[0] / tot_sr )

					vrec_abs.SetPhi( 0 );

				else

					vrec_abs.SetPhi( pi );
			}
			else
			{
				if ( prob[1] < sr[0] / tot_sr )

					vrec_abs.SetPhi( 0 );

				else
				{
					if ( prob[1] < ( sr[0] + sr[1] ) / tot_sr )

						vrec_abs.SetPhi( pi );

					else
					{
						if ( prob[1] < ( sr[0] + sr[1] + sr[2] ) / tot_sr )

							vrec_abs.SetPhi( pi / 2 );

						else

							vrec_abs.SetPhi( 3 * pi / 2 );
					}
				}
			}

			//------Velocity changes due to absorption and spontaneous emission------//

			vrec_spon.SetTheta( prob[2] * pi );

			vrec_spon.SetPhi( prob[3] * 2 * pi );

			v += vrec_abs + vrec_spon;


			//------Position changes due to change in velocity------//

			r += v * dt;

			rmol = r - cmol;


			//------Update position and velocity------//

			if ( j < nplot )
			{
				time_atom.push_back ( t );

				x.push_back ( r( 0 ) );

				y.push_back ( r( 1 ) );

				z.push_back ( r( 2 ) );

				vx.push_back ( v( 0 ) );

				vy.push_back ( v( 1 ) );

				vz.push_back ( v( 2 ) );
			}
		}

		if ( j < 5 ) cout << "# elements (postmol) : " << time_atom.size() << endl;

		transvel->Fill( v.Perp() );


		//------LVIS propagation to cavity------//

		dtp[0] = TMath::Abs( ( zcav - r.Z() ) / v.Z() ) / 100 ;

		dtp[1] = rcavmod / TMath::Abs( v.Z() ) / 4 ;

		if ( j < 5 )
		{
			cout << "coarse time step (postmol) : " << dtp[0] << endl;

			cout << "fine time step (postmol) : " << dtp[1] << endl;
		}

		prec = 0;

		while ( r.Z() < zcav + mirdiam / 2 )
		{
			if ( j < nplot )
			{
				time_atom.push_back ( t );

				x.push_back ( r( 0 ) );

				y.push_back ( r( 1 ) );

				z.push_back ( r( 2 ) );

				vx.push_back ( v( 0 ) );

				vy.push_back ( v( 1 ) );

				vz.push_back ( v( 2 ) );
			}

			t += dtp[prec];

			r += v * dtp[prec];

			if ( r.Mag() > molchamrad || nbeams == 0 )
			{
				radxz = sqrt( pow( r.Z() - zcav , 2 ) + pow( r.X() , 2 ) );

				prec = 1;

				if ( r.Theta() > thetam && nbeams != 0 )

					break;

				else

					if ( r.Y() > - cavspa / 2 && r.Y() < cavspa / 2 && radxz < rcavmod )
					{
						vzmodsum += v.Z();

						natoms_mod ++;

						break;
					}
			}
		}

		if ( j < 5 ) cout << "# elements (end) : " << time_atom.size() << endl;

		if ( j < nplot )
		{

			//------Graph position vs time per atom------//

			xvst = new TGraph( time_atom.size() , &(time_atom[0]) , &(x[0]) );

			yvst = new TGraph( time_atom.size() , &(time_atom[0]) , &(y[0]) );

			zvst = new TGraph( time_atom.size() , &(time_atom[0]) , &(z[0]) );

			xvst->SetLineColor( randcont1 * 40 );

			yvst->SetLineColor( randcont1 * 40 );

			zvst->SetLineColor( randcont1 * 40 );

			x_vs_t->Add( xvst );

			y_vs_t->Add( yvst );

			z_vs_t->Add( zvst );


			//------Graph velocity vs time per atom------//

			vxvst = new TGraph( time_atom.size() , &(time_atom[0]) , &(vx[0]) );

			vyvst = new TGraph( time_atom.size() , &(time_atom[0]) , &(vy[0]) );

			vzvst = new TGraph( time_atom.size() , &(time_atom[0]) , &(vz[0]) );

			vxvst->SetLineColor( randcont1 * 40 );

			vyvst->SetLineColor( randcont1 * 40 );

			vzvst->SetLineColor( randcont1 * 40 );

			vx_vs_t->Add( vxvst );

			vy_vs_t->Add( vyvst );

			vz_vs_t->Add( vzvst );


			//------Clear vector containers for next atom------//

			time_atom.clear();

			x.clear();

			y.clear();

			z.clear();

			vx.clear();

			vy.clear();

			vz.clear();
		}

	}

	delete[] detun;

	delete[] s0;

	delete[] gs0;

	delete[] s;

	delete[] sr;

	delete[] w;

	delete[] w0;

	delete[] phasecur;

	delete[] guoy;

	delete[] zr;


	//------Timing------//

	time ( & end );

	timing = difftime ( end , start );

	cout.precision( 5 );

	cout << "In vacuo tube length required : " << tube << endl;

	cout << "Atoms that go through : " << (Double_t) natoms_mod / natoms * 100 << "%" << endl;

	if ( natoms_mod != 0 )

		cout << "Average Vz of atoms going through : " << (Double_t) vzmodsum / natoms_mod << " m/s" << endl;

	cout << "Processing time " << timing << " seconds" << endl;
}