#include "dipole_radiation.h"

bool dipole_radiation::test_parameters(dipole_parameters q)
	{
	return q.test();
	}
//---------------------
void dipole_radiation::load_parameters(dipole_parameters  q)
	{
	param.Nx 					= q.Nx;
	param.Ny 					= q.Ny;
	param.Nz 					= q.Nz;
	param.Ntheta 			= q.Ntheta;
	param.Nphi 				= q.Nphi;
	param.xmin 				= q.xmin;
	param.xmax 				= q.xmax;
	param.ymin 				= q.ymin;
	param.ymax 				= q.ymax;
	param.zmin 				= q.zmin;
	param.zmax 				= q.zmax;
	param.theta_min 	= q.theta_min;
	param.theta_max 	= q.theta_max;
	param.phi_min 		= q.phi_min;
	param.phi_max 		= q.phi_max;
	param.f 					= q.f;
	param.lambda 			= q.lambda;

	if(param.Nx != 1)
		{
		Dx 	= (param.xmax-param.xmin)/double(param.Nx-1);
		}
	else
		{
		Dx = 0;
		}
	if(param.Ny != 1)
		{Dy = (param.ymax-param.ymin)/double(param.Ny-1);
		}
	else
		{
		Dy = 0;
		}
	if(param.Nz != 1)
		{
		Dz	= (param.zmax-param.zmin)/double(param.Nz-1);
		}
	else
		{
		Dz = 0;
		}
	if(param.Ntheta != 1)
		{
		Dtheta 	= (param.theta_max-param.theta_min)/double(param.Ntheta-1);
		}
	else
		{
		Dtheta = 0;
		}
	if(param.Nphi != 1)
		{
		Dphi 	= (param.phi_max-param.phi_min)/double(param.Nphi-1);
		}
	else
		{
		Dphi = 0;
		}

	outputstream.str(string());
	outputstream << "Input parameters:\n******************\nlambda       = " << param.lambda
	<< " µm\nfocal length = " << param.f
	<< " µm\nNx     = " << param.Nx 			<< "\txmin      = " << param.xmin 			<< " µm\txmax      = " << param.xmax      << " µm\tDx     = " << Dx
	<< " µm\nNy     = " << param.Ny 			<< "\tymin      = " << param.ymin 			<< " µm\txmax      = " << param.ymax      << " µm\tDy     = " << Dy
	<< " µm\nNz     = " << param.Nz 			<< "\tzmin      = " << param.zmin 			<< " µm\tzmax      = " << param.zmax      << " µm\tDz     = " << Dz
	<< " µm\nNtheta = " << param.Ntheta 	<< "\ttheta_min = " << param.theta_min << " rad\ttheta_max = " << param.theta_max << " rad\tDtheta = " << Dtheta
	<< " rad\nNphi   = " << param.Nphi 		<< "\tphi_min   = " << param.phi_min 		<< " rad\tphi_max   = " <<  param.phi_max 	<< " rad\tDphi   = " << Dphi << " rad\n\n";

	show_output(outputstream.str());
	}
//---------------------
void dipole_radiation::reset_parameters(dipole_parameters d)
	{
	if(test_parameters(d) == true)
		{
		load_parameters(d);
		initialize();
		}
	else
		{
		errorstream.str(string());
		errorstream << "Input error: Check your parameter input!\n";
		show_error(errorstream.str());
		}
	}
//---------------------
dipole_radiation::dipole_radiation(dipole_parameters d)
	{
	I = complex<double>(0.0,1.0);
	reset_parameters(d);
	}
//---------------------
dipole_radiation::~dipole_radiation()
	{
	x.free();
	y.free();
	z.free();
	theta.free();
	phi.free();
	x3.free();
	y3.free();
	z3.free();
	simpson_13_weights.free();
	E_far_x.free();
	E_far_y.free();
	E_far_z.free();
	E_refracted_x.free();
	E_refracted_y.free();
	E_refracted_z.free();
	E_dip_x.free();
	E_dip_y.free();
	E_dip_z.free();
	}
//---------------------
void dipole_radiation::populate_x()
	{
	x.resize(param.Nx);
	x  = param.xmin + tensor::i*Dx;
	x3.resize(param.Nx,param.Ny,param.Nz);
	for(int i =0; i<param.Nx;i++)
		{
		x3(i,Range::all(),Range::all()) = x(i);
		}
	}
//---------------------
void dipole_radiation::populate_y()
	{
	y.resize(param.Ny);
	y  = param.ymin + tensor::i*Dy;
	y3.resize(param.Nx,param.Ny,param.Nz);
	for(int j =0; j<param.Ny;j++)
		{
		y3(Range::all(),j,Range::all()) = y(j);
		}
	}
//----------------------
void dipole_radiation::populate_z()
	{
	z.resize(param.Nz);
	z = param.zmin + tensor::i*Dz;
	z3.resize(param.Nx,param.Ny,param.Nz);
	for (int k =0; k<param.Nz; k++)
		{
		z3(Range::all(),Range::all(),k) = z(k);
		}
	}
//-----------------------
void dipole_radiation::populate_theta()
	{
	theta.resize(param.Ntheta);
	theta = param.theta_min + tensor::i*Dtheta;
	theta2.resize(param.Ntheta,param.Nphi);
	for(int l = 0; l<param.Ntheta; l++)
		{
		theta2(l,Range::all()) = theta(l);
		}
	}
//-----------------------
void dipole_radiation::populate_phi()
	{
	phi.resize(param.Nphi);
	phi = param.phi_min + tensor::i*Dphi;
	phi2.resize(param.Ntheta,param.Nphi);
	for(int m = 0; m<param.Nphi; m++)
		{
		phi2(Range::all(),m) = phi(m);
		}
	}
//-----------------------
void dipole_radiation::populate_refracted_field()
	{
	E_refracted_x.resize(param.Ntheta,param.Nphi);
	E_refracted_y.resize(param.Ntheta,param.Nphi);
	E_refracted_z.resize(param.Ntheta,param.Nphi);
	E_refracted_x = 0.0;
	E_refracted_y = 0.0;
	E_refracted_z = 0.0;
	}
//-------------------------
void dipole_radiation::populate_far_field()
	{
	E_far_x.resize(param.Ntheta,param.Nphi);
	E_far_y.resize(param.Ntheta,param.Nphi);
	E_far_z.resize(param.Ntheta,param.Nphi);
	E_far_x = 0.0;
	E_far_y = 0.0;
	E_far_z = 0.0;
	}
//-------------------------
void dipole_radiation::calculate_simpson_13_weights()
	{
	simpson_13_weights.resize(param.Nx,param.Ny,param.Nz);
	double a,b,c;
	(Dx == 0) ? (a = 3.0) : (a = Dx); // Dx =0 only when Nx = 1
	(Dy == 0) ? (b = 3.0) : (b = Dy); // Dy =0 only when Ny = 1
	(Dz == 0) ? (c = 3.0) : (c = Dz); // Dz =0 only when Nz = 1
	//cout << "a = " << a << "\tb = " << b << "\tc = " << c << "\n";
	simpson_13_weights = ::calculate_simpson_13_weights(param.Nx,param.Ny,param.Nz,a,b,c);
	}
//------------------------
void dipole_radiation::initialize()
	{
	populate_x();
	populate_y();
	populate_z();
	populate_theta();
	populate_phi();
	populate_refracted_field();
	populate_far_field();
	initialize_E_dipole();
	calculate_simpson_13_weights();
	}
//---------------------------
void dipole_radiation::initialize_E_dipole()
	{
	E_dip_x.resize(param.Nx,param.Ny,param.Nz);
	E_dip_y.resize(param.Nx,param.Ny,param.Nz);
	E_dip_z.resize(param.Nx,param.Ny,param.Nz);
	E_dip_x = 1.0;
	E_dip_y = 0.0;
	E_dip_z = 0.0;
	}
//---------------------------
void dipole_radiation::set_dipole_field(Array<complex<double>,3> E_x,
                                        Array<complex<double>,3> E_y,
                                        Array<complex<double>,3> E_z)
	{
	E_dip_x = E_x;
	E_dip_y = E_y;
	E_dip_z = E_z;
	}
//--------------------------
void dipole_radiation::set_dipole_field(string filename_xcomp, string filename_ycomp, string filename_zcomp)
	{
	// -------------- load xcomp from data file
	fin.open(filename_xcomp.c_str());
	if(!fin)
		{
		errorstream.str(string());
		errorstream << "Unable to read file " << filename_xcomp << endl;
		show_error(errorstream.str());
		}
	Array<complex<double>,3> xcomp;
	fin >> xcomp;
	fin.close();
	// check the sanity of the data
	if( (xcomp.extent(0) != param.Nx) && (xcomp.extent(1) != param.Ny) && (xcomp.extent(2) != param.Nz)  )
		{
		errorstream.str(string());
		errorstream << "Dimensions of data in file " << filename_xcomp << " do not match!" << endl;
		show_error(errorstream.str());
		}

	// -------------- load ycomp from data file
	fin.open(filename_ycomp.c_str());
	if(!fin)
		{
		errorstream.str(string());
		errorstream << "Unable to read file " << filename_ycomp << endl;
		show_error(errorstream.str());
		}
	Array<complex<double>,3> ycomp;
	fin >> ycomp;
	fin.close();
	// check the sanity of the data
	if( (ycomp.extent(0) != param.Nx) && (ycomp.extent(1) != param.Ny) && (ycomp.extent(2) != param.Nz)  )
		{
		errorstream.str(string());
		errorstream << "Dimensions of data in file " << filename_ycomp << " do not match!" << endl;
		show_error(errorstream.str());
		}

	// -------------- load zcomp from data file
	fin.open(filename_zcomp.c_str());
	if(!fin)
		{
		errorstream.str(string());
		errorstream << "Unable to read file " << filename_zcomp << endl;
		show_error(errorstream.str());
		}
	Array<complex<double>,3> zcomp;
	fin >> zcomp;
	fin.close();
	// check the sanity of the data
	if( (zcomp.extent(0) != param.Nx) && (zcomp.extent(1) != param.Ny) && (zcomp.extent(2) != param.Nz)  )
		{
		errorstream.str(string());
		errorstream << "Dimensions of data in file " << filename_zcomp << " do not match!" << endl;
		show_error(errorstream.str());
		}

	E_dip_x = xcomp;
	E_dip_y = ycomp;
	E_dip_z = zcomp;
	}
//--------------------------
void dipole_radiation::calculate_radiation( Array<complex<double>, 2>& E_rad_x,
    Array<complex<double>, 2>& E_rad_y,
    Array<complex<double>, 2>& E_rad_z)
	{
	calculate_radiation();
	E_rad_x = E_far_x;
	E_rad_y = E_far_y;
	E_rad_z = E_far_z;
	}
//--------------------------

void dipole_radiation::calculate_radiation()
	{

	double X = 0;
	double Y = 0;
	double Z = 0;
	double k0 = 2.0*M_PI/param.lambda;

	Array<double,3> R(param.Nx,param.Ny,param.Nz);
	Array<double,3> phase_factor(param.Nx,param.Ny,param.Nz);
	Array<complex<double>,3> prefact(param.Nx,param.Ny,param.Nz);
	//cout << "x3  = " << x3 << endl << "y3 = " << y3 << endl << "z3 = " << z3 << endl;

	int loop_var = 1;
	for (int l = 0; l<param.Ntheta; l++)
		{
		cout << round(double(loop_var)*100.00/double(param.Ntheta*param.Nphi)) << "%\t";
		cout.flush();
		for(int m =0; m<param.Nphi;m++)
			{
			X = param.f*sin(theta(l))*cos(phi(m));
			Y = param.f*sin(theta(l))*sin(phi(m));
			Z = param.f*cos(theta(l));

			/*
			Far field at the curved surface of the collecting lens of focal length fp is simply related to the polarization at the focus through the relation E = -[exp(ik.R)/(4.pi.R^3)](R x (R x P) ) where R = r-r' where vector r is the vector joining the point of interest from the eorigin and r' is the vector from the origin to the source point in the focal volume. The above relationship can be expressed in a matrix form. Below r=r(X,Y,Z), r'=r'(x,y,z) and P=P(x,y,z). Note |r| = fp.	
			*/

			R = sqrt( sqr(X-x3) + sqr(Y-y3) + sqr(Z-z3) );
			phase_factor = k0*param.f - k0*( + sin(theta(l))*cos(phi(m))*x3
			                                 + sin(theta(l))*sin(phi(m))*y3
			                                 + cos(theta(l))*z3
			                               );
			prefact = -(1.0/(4.0*M_PI*R*R*R))*exp(I*phase_factor);

			/*
					integrating all the field contributions from individual (x,y,z) points to calculate the far field at E(theta,phi) on a sphere of radius fp (focal length of the collecting lens). For integration we use Simpson's 1/3 rule.
			*/

			E_far_x(l,m) = sum( 	( - E_dip_x * ( sqr(Y-y3) + sqr(Z-z3) )
			                       + E_dip_y * (X-x3)*(Y-y3)
			                       + E_dip_z * (X-x3)*(Z-z3)
			                     ) * prefact * simpson_13_weights
			                  );
			E_far_y(l,m) = sum(  ( + E_dip_x * (X-x3)*(Y-y3)
			                       - E_dip_y * ( sqr(X-x3) + sqr(Z-z3) )
			                       + E_dip_z * (Y-y3)*(Z-z3)
			                     ) * prefact * simpson_13_weights
			                  );
			E_far_z(l,m) = sum( 	( + E_dip_x * (X-x3)*(Z-z3)
			                       + E_dip_y * (Y-y3)*(Z-z3)
			                       - E_dip_z * ( sqr(X-x3) + sqr(Y-y3) )
			                     ) * prefact * simpson_13_weights
			                  );
			loop_var++;
			}
		}
	cout << "Done\n"; cout.flush();
	}

//--------------------------------
void dipole_radiation::calculate_refraction(Array<complex<double>, 2>& E_refr_x,
    Array<complex<double>, 2>& E_refr_y,
    Array<complex<double>, 2>& E_refr_z)
	{
	calculate_refraction();
	E_refr_x = E_refracted_x;
	E_refr_y = E_refracted_y;
	E_refr_z = E_refracted_z;
	}
//------------------------------------
void dipole_radiation::calculate_refraction()
	{
	Array<double,2> temp(param.Ntheta,param.Nphi);
	temp = 1.0/sqrt(cos(theta2));
	E_refracted_x =
	  ( + E_far_x * ( sqr(cos(phi2))*cos(theta2) + sqr(sin(phi2)) )
	    + E_far_y * ( -(1-cos(theta2))*sin(phi2)*cos(phi2) )
	    + E_far_z * ( -cos(phi2)*sin(theta2) )
	  )*temp;
	E_refracted_y =
	  ( + E_far_x * ( -(1-cos(theta2))*sin(phi2)*cos(phi2) )
	    + E_far_y * ( sqr(sin(phi2))*cos(theta2) + sqr(cos(phi2)) )
	    + E_far_z * ( -sin(phi2)*sin(theta2) )
	  )*temp;
	E_refracted_z =
	  ( + E_far_x * ( sin(theta2)*cos(phi2) )
	    + E_far_y * ( sin(theta2)*sin(phi2) )
	    + E_far_z * ( cos(theta2) )
	  )*temp;
	}
//--------------------------------
void dipole_radiation::calculate_refracted_radiation(Array<complex<double>, 2>& E_rad_x,
    Array<complex<double>, 2>& E_rad_y,
    Array<complex<double>, 2>& E_rad_z)
	{
	calculate_radiation();
	calculate_refraction(E_rad_x,E_rad_y,E_rad_z);
	}
//----------------------------------
void dipole_radiation::save_results(string idir)
	{
	namestream.str(string());
	namestream << idir << "_rad_lambda_" << param.lambda << "/";
	string outputdir = namestream.str();
	my_mkdir(outputdir);

	outputstream.str(string());
	outputstream << "\nSaving the results to " << outputdir << "\n";
	show_output(outputstream.str());

	outfilename.str(string());	outfilename << outputdir << "simulation_parameters.dat";
	param.save(outfilename.str());

	outfilename.str(string());	outfilename << outputdir << "x.bar";
	fout.open(outfilename.str().c_str(), ios::out);
	if(!fout)
		{
		show_outfile_error(outfilename.str());
		}
	fout << x;
	fout.close();

	outfilename.str(string());	outfilename << outputdir << "y.bar";
	fout.open(outfilename.str().c_str(), ios::out);
	if(!fout)
		{
		show_outfile_error(outfilename.str());
		}
	fout << y;
	fout.close();

	outfilename.str(string());	outfilename << outputdir << "z.bar";
	fout.open(outfilename.str().c_str(), ios::out);
	if(!fout)
		{
		show_outfile_error(outfilename.str());
		}
	fout << z;
	fout.close();

	outfilename.str(string());	outfilename << outputdir << "theta.bar";
	fout.open(outfilename.str().c_str(), ios::out);
	if(!fout)
		{
		show_outfile_error(outfilename.str());
		}
	fout << theta;
	fout.close();

	outfilename.str(string());	outfilename << outputdir << "phi.bar";
	fout.open(outfilename.str().c_str(), ios::out);
	if(!fout)
		{
		show_outfile_error(outfilename.str());
		}
	fout << phi;
	fout.close();

	outfilename.str(string());	outfilename << outputdir << "E_far_x.bar";
	fout.open(outfilename.str().c_str(), ios::out);
	if(!fout)
		{
		show_outfile_error(outfilename.str());
		}
	fout << E_far_x;
	fout.close();

	outfilename.str(string());	outfilename << outputdir << "E_far_y.bar";
	fout.open(outfilename.str().c_str(), ios::out);
	if(!fout)
		{
		show_outfile_error(outfilename.str());
		}
	fout << E_far_y;
	fout.close();

	outfilename.str(string());	outfilename << outputdir << "E_far_z.bar";
	fout.open(outfilename.str().c_str(), ios::out);
	if(!fout)
		{
		show_outfile_error(outfilename.str());
		}
	fout << E_far_z;
	fout.close();

	outfilename.str(string());	outfilename << outputdir << "E_refracted_x.bar";
	fout.open(outfilename.str().c_str(), ios::out);
	if(!fout)
		{
		show_outfile_error(outfilename.str());
		}
	fout << E_refracted_x;
	fout.close();

	outfilename.str(string());	outfilename << outputdir << "E_refracted_y.bar";
	fout.open(outfilename.str().c_str(), ios::out);
	if(!fout)
		{
		show_outfile_error(outfilename.str());
		}
	fout << E_refracted_y;
	fout.close();

	outfilename.str(string());	outfilename << outputdir << "E_refracted_z.bar";
	fout.open(outfilename.str().c_str(), ios::out);
	if(!fout)
		{
		show_outfile_error(outfilename.str());
		}
	fout << E_refracted_z;
	fout.close();

	outfilename.str(string());	outfilename << outputdir << "E_dipole_x.bar";
	fout.open(outfilename.str().c_str(), ios::out);
	if(!fout)
		{
		show_outfile_error(outfilename.str());
		}
	fout << E_dip_x;
	fout.close();

	outfilename.str(string());	outfilename << outputdir << "E_dipole_y.bar";
	fout.open(outfilename.str().c_str(), ios::out);
	if(!fout)
		{
		show_outfile_error(outfilename.str());
		}
	fout << E_dip_y;
	fout.close();

	outfilename.str(string());	outfilename << outputdir << "E_dipole_z.bar";
	fout.open(outfilename.str().c_str(), ios::out);
	if(!fout)
		{
		show_outfile_error(outfilename.str());
		}
	fout << E_dip_z;
	fout.close();

	cout << "Finished!\n";
	}
//------------------
void dipole_radiation::show_outfile_error(string filename)
	{
	errorstream.str(string());
	errorstream << "Unable to write file " << filename << "\n";
	show_error(errorstream.str());
	}
//-------------------
void dipole_radiation::show_error(string err_str)
	{
	cerr << "!!!! " << err_str;
	exit(-1);
	}
//-------------------
void dipole_radiation::show_output(string out_str)
	{
	cout << out_str;
	}

