#include "focus3d.h"

//---------------------
bool focus3d::test_parameters(focus3d_parameters q)
	{
	return q.test();
	}
//---------------------
void focus3d::load_parameters(focus3d_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());
	}
//---------------------
focus3d::focus3d(focus3d_parameters q)
	{
	I = complex<double>(0.0,1.0);
	reset_parameters(q);
	}
//---------------------
void focus3d::reset_parameters(focus3d_parameters q)
	{
	if(test_parameters(q) == true)
		{
		load_parameters(q);
		initialize();
		}
	else
		{
		errorstream.str(string());
		errorstream << "Input error: Check your parameter input!\n";
		show_error(errorstream.str());
		}
	}
//----------------------
focus3d::~focus3d()
	{
	x.free();
	y.free();
	z.free();
	theta.free();
	phi.free();
	simpson_13_weights.free();
	angular_weights.free();
	theta2.free();
	phi2.free();
	E_inc_x.free();
	E_inc_y.free();
	E_inc_z.free();
	integrand_amplitude_x.free();
	integrand_amplitude_y.free();
	integrand_amplitude_z.free();
	exp_arg1.free();
	exp_arg2.free();
	exp_arg3.free();
	integrand_exp_term.free();
	E_focus_x.free();
	E_focus_y.free();
	E_focus_z.free();
	}
//----------------------
void focus3d::populate_x()
	{
	x.resize(param.Nx);
	x  = param.xmin + tensor::i*Dx;
	}
//----------------------
void focus3d::populate_y()
	{
	y.resize(param.Ny);
	y  = param.ymin + tensor::i*Dy;
	}
//----------------------
void focus3d::populate_z()
	{
	z.resize(param.Nz);
	z  = param.zmin + tensor::i*Dz;
	}
//----------------------
void focus3d::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 focus3d::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 focus3d::calculate_integrand_amplitudes()
	{
	// !!!!!!! call this function only after populating theta2, phi2, and E_inc_x,y,z
	integrand_amplitude_x.resize(param.Ntheta,param.Nphi);
	integrand_amplitude_y.resize(param.Ntheta,param.Nphi);
	integrand_amplitude_z.resize(param.Ntheta,param.Nphi);

	integrand_amplitude_x =
	  sin(theta2) * sqrt(cos(theta2))
	  * ( + E_inc_x * ( sqr(cos(phi2))*cos(theta2) + sqr(sin(phi2)) )
	      + E_inc_y * ( -(1-cos(theta2))*sin(phi2)*cos(phi2) )
	      + E_inc_z * ( -cos(phi2)*sin(theta2) )
	    );
	integrand_amplitude_y =
	  sin(theta2) * sqrt(cos(theta2))
	  * ( + E_inc_x * ( -(1-cos(theta2))*sin(phi2)*cos(phi2) )
	      + E_inc_y * ( sqr(sin(phi2))*cos(theta2) + sqr(cos(phi2)) )
	      + E_inc_z * ( -sin(phi2)*sin(theta2) )
	    );
	integrand_amplitude_z =
	  sin(theta2) * sqrt(cos(theta2))
	  * ( + E_inc_x * ( sin(theta2)*cos(phi2) )
	      + E_inc_y * ( sin(theta2)*sin(phi2) )
	      + E_inc_z * ( cos(theta2) )
	    );
	}
//----------------------
void focus3d::calculate_integrand_exp_args()
	{
	// !!!!!!! call this function only after populating theta2 and phi2

	exp_arg1.resize(param.Ntheta,param.Nphi);
	exp_arg2.resize(param.Ntheta,param.Nphi);
	exp_arg3.resize(param.Ntheta,param.Nphi);

	exp_arg1 = I * (2.0*M_PI/param.lambda) * sin(theta2) * cos(phi2); 	// i.k.sin(t).cos(f)
	exp_arg2 = I * (2.0*M_PI/param.lambda) * sin(theta2) * sin(phi2);	// i.k.sin(t).sin(f)
	exp_arg3 = I * (2.0*M_PI/param.lambda) * cos(theta2);						// i.k.cos(t)

	prefactor = 0;
	prefactor = I * (1.0/param.lambda) * param.f * exp(I * param.f * (2.0*M_PI/param.lambda));
	}
//----------------------
void focus3d::calculate_simpson_13_weights()
	{
	simpson_13_weights.resize(param.Ntheta,param.Nphi);
	double a,b;
	(Dtheta == 0) ? (a = 3.0) : (a = Dtheta); // Dtheta =0 only when Ntheta = 1
	(Dphi == 0) ? (b = 3.0) : (b = Dphi); // Dphi =0 only when Nphi = 1
	//cout << "a = " << a << "\tb = " << b <<  "\n";
	simpson_13_weights = ::calculate_simpson_13_weights(param.Ntheta,param.Nphi,a,b);
	}
//----------------------
void focus3d::set_angular_weights(Array<complex<double>,2> aweights)
	{
	initialize_angular_weights();
	if( (aweights.extent(0) == param.Ntheta) && (aweights.extent(1) == param.Nphi) )
		{
		angular_weights = aweights;
		}
	else
		{
		errorstream.str(string());
		errorstream << "Input error: Input matrix should be of size " << param.Ntheta << " x " << param.Nphi << "\n";
		show_error(errorstream.str());
		}
	}
//----------------------
void focus3d::set_angular_weights(string g_of_theta_file)
	{
	initialize_angular_weights();
	Array<complex<double>,1> interpolated_g_of_theta(param.Ntheta);
	interpolate_data(g_of_theta_file, theta, interpolated_g_of_theta);
	cout << interpolated_g_of_theta << "\n";
	for (int l =0; l<param.Ntheta;l++)
		{
		angular_weights(l,Range::all()) = interpolated_g_of_theta(l);
		}
	}
//----------------------
void focus3d::set_angular_weights(string ampl_coeff_file, string phase_coeff_file)
	{
	initialize_angular_weights();

	int temp1 = determine_poly_fit_order(ampl_coeff_file);
	int temp2 = determine_poly_fit_order(phase_coeff_file);

	if(temp1 != temp2)
		{
		errorstream.str(string());
		errorstream << "Input error: Polynomial coefficient files are not of equal lengths!\n";
		show_error(errorstream.str());
		}

	int poly_fit_order = temp1;

	Array<double,1> poly_coeffs_ampl(poly_fit_order+1), poly_coeffs_phase(poly_fit_order+1);

	poly_coeffs_ampl 	= load_polynomial_coeffs(ampl_coeff_file,poly_fit_order); // load the ampl polynomial coeffs file
	poly_coeffs_phase	= load_polynomial_coeffs(phase_coeff_file,poly_fit_order); // load the phase polynomial coeffs file
	cout << "Amplitude coeffs:\n" << poly_coeffs_ampl << endl
	<< "Phase coeffs:\n" << poly_coeffs_phase << endl;

	double temp_ampl;
	double temp_phase;

	for (int l =0; l<param.Ntheta;l++)
		{
		temp_ampl = 0;
		temp_phase = 0;
		for(int i =0; i< poly_fit_order+1;i++)
			{
			temp_ampl += poly_coeffs_ampl(i)*pow(theta(l), poly_fit_order-i);
			temp_phase += poly_coeffs_phase(i)*pow(theta(l), poly_fit_order-i);
			}
		angular_weights(l,Range::all()) = temp_ampl*exp(I*temp_phase);
		}
	}
//----------------------
void focus3d::initialize_angular_weights()
	{
	angular_weights.resize(param.Ntheta,param.Nphi);
	angular_weights = 1.0; 		// by default rays at all angles contribute equally to the focus
	}
//----------------------
void focus3d::initialize_E_incs()
	{
	E_inc_x.resize(param.Ntheta,param.Nphi); E_inc_x = 1.0; // x-component - default input polarization
	E_inc_y.resize(param.Ntheta,param.Nphi); E_inc_y = 0.0; // y-component
	E_inc_z.resize(param.Ntheta,param.Nphi); E_inc_y = 0.0; // z-component
	}
//----------------------
void focus3d::set_incident_field(Array<complex<double>,2> E_x,
                                 Array<complex<double>,2> E_y,
                                 Array<complex<double>,2> E_z)
	{
	if( (E_x.extent(0) == param.Ntheta) && (E_x.extent(1) == param.Nphi)
	        && (E_y.extent(0) == param.Ntheta) && (E_y.extent(1) == param.Nphi)
	        && (E_z.extent(0) == param.Ntheta) && (E_z.extent(1) == param.Nphi) )
		{
		E_inc_x = E_x;
		E_inc_y = E_y;
		E_inc_z = E_z;
		}
	else
		{
		errorstream.str(string());
		errorstream << "Input error: Input matrices should be of size " << param.Ntheta << " x " << param.Nphi << "\n";
		show_error(errorstream.str());
		}
	}
//----------------------
void focus3d::set_incident_field(complex<double> E_x,
                                 complex<double> E_y,
                                 complex<double> E_z)
	{
	E_inc_x = E_x;
	E_inc_y = E_y;
	E_inc_z = E_z;
	}
//----------------------
void focus3d::set_incident_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>,2> xcomp;
	fin >> xcomp;
	fin.close();
	// check the sanity of the data
	if( (xcomp.extent(0) != param.Ntheta) && (xcomp.extent(1) != param.Nphi)  )
		{
		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>,2> ycomp;
	fin >> ycomp;
	fin.close();
	// check the sanity of the data
	if( (ycomp.extent(0) != param.Ntheta) && (ycomp.extent(1) != param.Nphi)  )
		{
		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>,2> zcomp;
	fin >> zcomp;
	fin.close();
	// check the sanity of the data
	if( (zcomp.extent(0) != param.Ntheta) && (zcomp.extent(1) != param.Nphi)  )
		{
		errorstream.str(string());
		errorstream << "Dimensions of data in file " << filename_zcomp << " do not match!" << endl;
		show_error(errorstream.str());
		}

	E_inc_x = xcomp;
	E_inc_y = ycomp;
	E_inc_z = zcomp;
	}

//----------------------
void focus3d::initialize()
	{
	populate_theta();
	populate_phi();
	populate_x();
	populate_y();
	populate_z();
	initialize_E_incs(); 				//default input polarization is along x-axis
	initialize_angular_weights();
	calculate_simpson_13_weights();
	calculate_integrand_exp_args();
	}
//----------------------
void focus3d::calculate_focus()
	{
	calculate_integrand_amplitudes();

	E_focus_x.resize(param.Nx,param.Ny,param.Nz);
	E_focus_y.resize(param.Nx,param.Ny,param.Nz);
	E_focus_z.resize(param.Nx,param.Ny,param.Nz);
	integrand_exp_term.resize(param.Ntheta,param.Nphi);

	E_focus_x = 0;
	E_focus_y = 0;
	E_focus_z = 0;

	int loop_var = 1;
	for(int i=0; i<param.Nx; i++)
		{
		cout << round(double(loop_var)*100.00/double(param.Nx*param.Ny*param.Nz)) << "%\t";
		cout.flush();
		for(int j=0; j<param.Ny; j++)
			{
			for(int k=0; k<param.Nz; k++)
				{
				integrand_exp_term = exp( + exp_arg1*x(i)
				                          + exp_arg2*y(j)
				                          + exp_arg3*z(k) );
				E_focus_x(i,j,k) = sum( angular_weights
				                        * integrand_amplitude_x
				                        * integrand_exp_term
				                        * simpson_13_weights );
				E_focus_y(i,j,k) = sum( angular_weights
				                        * integrand_amplitude_y
				                        * integrand_exp_term
				                        * simpson_13_weights );
				E_focus_z(i,j,k) = sum( angular_weights
				                        * integrand_amplitude_z
				                        * integrand_exp_term
				                        * simpson_13_weights );
				E_focus_x(i,j,k) = prefactor*E_focus_x(i,j,k);
				E_focus_y(i,j,k) = prefactor*E_focus_y(i,j,k);
				E_focus_z(i,j,k) = prefactor*E_focus_z(i,j,k);
				//cout << E_focus_x(i,j,k) << "\t" << E_focus_y(i,j,k) << "\t" << E_focus_z(i,j,k) << "\n";
				loop_var++;
				}
			}
		}
	cout << "Done\n";
	cout.flush();
	loop_var = 0;
	}
//----------------------
void focus3d::calculate_focus(Array<complex<double>,3>& x_comp,
                              Array<complex<double>,3>& y_comp,
                              Array<complex<double>,3>& z_comp)
	{
	calculate_focus();
	x_comp = E_focus_x;
	y_comp = E_focus_y;
	z_comp = E_focus_z;
	}
//----------------------
void focus3d::save_results(string idir)
	{
	namestream.str(string());
	namestream << idir << "_focus_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_focus_x.bar";
	fout.open(outfilename.str().c_str(), ios::out);
	if(!fout)
		{
		show_outfile_error(outfilename.str());
		}
	fout << E_focus_x;
	fout.close();

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

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

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

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

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

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

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

