#include <math.h>
#include <complex>
#include <iostream>
#include <fstream>
#include "globals.h"
#include <stdlib.h>
//#include "cufft.h"
using namespace std;

#define pi 3.14159

typedef complex<double> scComplex;

extern int cbessjyva(double v,complex<double> z,double &vm,complex<double>*cjv,
    complex<double>*cyv,complex<double>*cjvp,complex<double>*cyvp);
extern int bessjyv(double v,double x,double &vm,double *jv,double *yv,
    double *djv,double *dyv);

complex<double> Jl_neg(complex<double> x)
{
	//this function computes the bessel function of the first kind Jl(x) for l = -0.5
	return ( sqrt(2.0/pi) * cos(x) )/sqrt(x);
}

double Jl_neg(double x)
{
	//this function computes the bessel function of the first kind Jl(x) for l = -0.5
	return ( sqrt(2.0/pi) * cos(x) )/sqrt(x);
}

double Yl_neg(double x)
{
	//this function computes the bessel function of the second kind Yl(x) for l = -0.5;
	return ( sqrt(2.0/pi) * sin(x) )/sqrt(x);
}

void computeB(complex<double>* B, double radius, complex<double> refIndex, double lambda, int Nl)
{
	double k = (2*pi)/lambda;
	int b = 2;

	//allocate space for the real bessel functions
	double* jv = (double*)malloc(sizeof(double)*(Nl+b));
	double* yv = (double*)malloc(sizeof(double)*(Nl+b));
	double* jvp = (double*)malloc(sizeof(double)*(Nl+b));
	double* yvp = (double*)malloc(sizeof(double)*(Nl+b));

	//allocate space for the complex bessel functions
	complex<double>* cjv = (complex<double>*)malloc(sizeof(complex<double>)*(Nl+b));
	complex<double>* cyv = (complex<double>*)malloc(sizeof(complex<double>)*(Nl+b));
	complex<double>* cjvp = (complex<double>*)malloc(sizeof(complex<double>)*(Nl+b));
	complex<double>* cyvp = (complex<double>*)malloc(sizeof(complex<double>)*(Nl+b));

	double kr = k*radius;
	complex<double> knr = k*refIndex*(double)radius;
	complex<double> n = refIndex;

	//compute the bessel functions for k*r
	double vm;// = Nl - 1;
	bessjyv((Nl)+0.5, kr, vm, jv, yv, jvp, yvp);
	//cout<<"Nl: "<<Nl<<"  vm: "<<vm<<endl;
	//printf("Nl: %f, vm: %f\n", (float)Nl, (float)vm);

	//compute the bessel functions for k*n*r
	cbessjyva((Nl)+0.5, knr, vm, cjv, cyv, cjvp, cyvp);

	//scale factor for spherical bessel functions
	double scale_kr = sqrt(pi/(2.0*kr));
	complex<double> scale_knr = sqrt(pi/(2.0*knr));

	complex<double> numer, denom;
	double j_kr;
	double y_kr;
	complex<double> j_knr;
	complex<double> j_d_knr;
	double j_d_kr;
	complex<double> h_kr;
	complex<double> h_d_kr;
	complex<double> h_neg;
	complex<double> h_pos;

	//cout<<"B coefficients:"<<endl;
	for(int l=0; l<Nl; l++)
	{
		//compute the spherical bessel functions
		j_kr = jv[l] * scale_kr;
		y_kr = yv[l] * scale_kr;
		j_knr = cjv[l] * scale_knr;

		//compute the Hankel function
		h_kr = complex<double>(j_kr, y_kr);

		//compute the derivatives
		if(l == 0)
		{
			//spherical bessel functions for l=0
			j_d_kr = scale_kr * (Jl_neg(kr) - (jv[l] + kr*jv[l+1])/kr )/2.0;
			j_d_knr = scale_knr * ( Jl_neg(knr) - (cjv[l] + knr*cjv[l+1])/knr )/2.0;
			h_neg = complex<double>(scale_kr*Jl_neg(kr), scale_kr*Yl_neg(kr));
			h_pos = complex<double>(scale_kr*jv[l+1], scale_kr*yv[l+1]);
			h_d_kr = (h_neg - (h_kr + kr*h_pos)/kr)/2.0;
		}
		else
		{
			//spherical bessel functions
			j_d_kr = scale_kr * (jv[l-1] - (jv[l] + kr*jv[l+1])/kr )/2.0;
			j_d_knr = scale_knr * ( cjv[l-1] - (cjv[l] + knr*cjv[l+1])/knr )/2.0;
			h_neg = complex<double>(scale_kr*jv[l-1], scale_kr*yv[l-1]);
			h_pos = complex<double>(scale_kr*jv[l+1], scale_kr*yv[l+1]);
			h_d_kr = (h_neg - (h_kr + kr*h_pos)/kr)/2.0;
		}

		numer = j_kr*j_d_knr*n - j_knr*j_d_kr;
		denom = j_knr*h_d_kr - h_kr*j_d_knr*n;
		B[l] =  numer/denom;

		//B[l] = scComplex(temp.real(), temp.imag());
		//cout<<B[l]<<endl;
	}

	free(jv);
	free(yv);
	free(jvp);
	free(yvp);
	free(cjv);
	free(cyv);
	free(cjvp);
	free(cyvp);
}

void Legendre(double* P, double x, int Nl)
{
	//computes the legendre polynomials from orders 0 to Nl-1
	P[0] = 1;
	if(Nl == 1) return;
	P[1] = x;
	for(int l = 2; l < Nl; l++)
	{
		P[l] = ((2*l - 1)*x*P[l-1] - (l - 1)*P[l-2])/l;
	}	

}

complex<double> integrateUi(double cAngleI, double cAngleO, double oAngleI, double oAngleO, double M = 2*pi)
{
	/*This function integrates the incident field of magnitude M in the far zone
	in order to evaluate the field at the central pixel of a detector.
	cNAi = condenser inner angle
	cNAo = condenser outer angle
	oNAi = objective inner angle
	oNAo = objective outer angle
	M = field magnitude*/

	double alphaIn = max(cAngleI, oAngleI);
	double alphaOut = min(cAngleO,oAngleO);

	complex<double> Ui;
	if(alphaIn > alphaOut)
		Ui = complex<double>(0.0, 0.0);
	else
		Ui = complex<double>(M * 2 * pi * (cos(alphaIn) - cos(alphaOut)), 0.0f);

	return Ui;

}

void computeCondenserAlpha(double* alpha, int Nl, double cAngleI, double cAngleO)
{
	/*This function computes the condenser integral in order to build the field of incident light
	alpha = list of Nl floating point values representing the condenser alpha as a function of l
	Nl = number of orders in the incident field
	cAngleI, cAngleO = inner and outer condenser angles (inner and outer NA)*/

	//compute the Legendre polynomials for the condenser aperature
	double* PcNAo = (double*)malloc(sizeof(double)*(Nl+1));
	Legendre(PcNAo, cos(cAngleO), Nl+1);
	double* PcNAi = (double*)malloc(sizeof(double)*(Nl+1));
	Legendre(PcNAi, cos(cAngleI), Nl+1);

	for(int l=0; l<Nl; l++)
	{
		//integration term
		if(l == 0)
			alpha[l] = -PcNAo[l+1] + PcNAo[0] + PcNAi[l+1] - PcNAi[0];
		else
			alpha[l] = -PcNAo[l+1] + PcNAo[l-1] + PcNAi[l+1] - PcNAi[l-1];

		alpha[l] *= 2 * pi;
	}

}

complex<double> integrateUs(double r, double lambda, complex<double> eta, 
						   double cAngleI, double cAngleO, double oAngleI, double oAngleO, double M = 2*pi)
{
	/*This function integrates the incident field of magnitude M in the far zone
	in order to evaluate the field at the central pixel of a detector.
	r = sphere radius
	lambda = wavelength
	eta = index of refraction
	cNAi = condenser inner NA
	cNAo = condenser outer NA
	oNAi = objective inner NA
	oNAo = objective outer NA
	M = field magnitude*/

	//compute the required number of orders
	double k = 2*pi/lambda;
	int Nl = (int)ceil( k + 4 * exp(log(k*r)/3) + 3 );

	//compute the material coefficients B
	complex<double>* B = (complex<double>*)malloc(sizeof(complex<double>)*Nl);
	//compute the Legendre polynomials for the condenser and objective aperatures
	double* PcNAo = (double*)malloc(sizeof(double)*(Nl+1));
	Legendre(PcNAo, cos(cAngleO), Nl+1);
	double* PcNAi = (double*)malloc(sizeof(double)*(Nl+1));
	Legendre(PcNAi, cos(cAngleI), Nl+1);

	double* PoNAo = (double*)malloc(sizeof(double)*(Nl+1));
	Legendre(PoNAo, cos(oAngleO), Nl+1);
	double* PoNAi = (double*)malloc(sizeof(double)*(Nl+1));
	Legendre(PoNAi, cos(oAngleI), Nl+1);

	//store the index of refraction;
	complex<double> IR(eta.real(), eta.imag());
	
	//compute the scattering coefficients
	computeB(B, r, IR, lambda, Nl);

	//aperature terms for the condenser (alpha) and objective (beta)
	double alpha;
	double beta;
	double c;
	complex<double> Us(0.0, 0.0);

	for(int l=0; l<Nl; l++)
	{
		//integration term
		if(l == 0)
		{
			alpha = -PcNAo[l+1] + PcNAo[0] + PcNAi[l+1] - PcNAi[0];
			beta = -PoNAo[l+1] + PoNAo[0] + PoNAi[l+1] - PoNAi[0];
		}
		else
		{
			alpha = -PcNAo[l+1] + PcNAo[l-1] + PcNAi[l+1] - PcNAi[l-1];
			beta = -PoNAo[l+1] + PoNAo[l-1] + PoNAi[l+1] - PoNAi[l-1];
		}
		c = (2*pi)/(2.0 * l + 1.0);
		Us += c * alpha * beta * B[l] * M;
		
			
	}
	free(PcNAo);
	free(PcNAi);
	free(PoNAo);
	free(PoNAi);
	free(B);

	return Us;

}

void pointSpectrum()
{
	PD.StartTimer(SIMULATE_SPECTRUM);
	//clear the previous spectrum
	SimSpectrum.clear();

	double dNu = 2.0f;
	double lambda;
	
	//compute the angles based on NA
	double cAngleI = asin(cNAi);
	double cAngleO = asin(cNAo);
	double oAngleI = asin(oNAi);
	double oAngleO = asin(oNAo);

	//implement a reflection-mode system if necessary
	if(opticsMode == ReflectionOpticsType){
		
		//set the condenser to match the objective
		cAngleI = oAngleI;
		cAngleO = oAngleO;

		//invert the objective
		oAngleO = pi - cAngleI;
		oAngleI = pi - cAngleO;
	}

	//integrate the incident field at the detector position
	complex<double> Ui = integrateUi(cAngleI, cAngleO, oAngleI, oAngleO, 2*pi);
	double I0 = Ui.real() * Ui.real() + Ui.imag() * Ui.imag();
	I0 *= scaleI0;

	

	//double I;
	SpecPair temp;
	double nu;
	complex<double> eta;
	complex<double> Us, U;

	double vecLen = 0.0;
	for(unsigned int i=0; i<EtaK.size(); i++)
	{
		nu = EtaK[i].nu;
		lambda = 10000.0f/nu;
		if(applyMaterial)
			eta = complex<double>(EtaN[i].A, EtaK[i].A);
		else
			eta = complex<double>(baseIR, 0.0);


		//integrate the scattered field at the detector position
		Us = integrateUs(radius, lambda, eta, cAngleI, cAngleO, oAngleI, oAngleO, 2*pi);
		U = Us + Ui;
		double I = U.real() * U.real() + U.imag() * U.imag();
		
		temp.nu = nu;

		//set the spectrum value based on the current display type
		if(dispSimType == AbsorbanceSpecType)
			temp.A = -log10(I/I0);
		else
			temp.A = I;

		if(dispNormalize)
			vecLen += temp.A * temp.A;
		
		SimSpectrum.push_back(temp);		
	}
	vecLen = sqrt(vecLen);

	if(dispNormalize)
		for(unsigned int i=0; i<SimSpectrum.size(); i++)
			SimSpectrum[i].A = (SimSpectrum[i].A / vecLen) * dispNormFactor;

	PD.EndTimer(SIMULATE_SPECTRUM);
}

void updateSpectrum(double* I, double I0, int n)
{
	SimSpectrum.clear();
	SpecPair temp;

	//update the displayed spectrum based on the computed intensity I
	for(int i=0; i<n; i++)
	{
		temp.nu = EtaK[i].nu;

		//set the spectrum value based on the current display type
		if(dispSimType == AbsorbanceSpecType)
			temp.A = -log10(I[i]/I0);
		else
			temp.A = I[i];

		SimSpectrum.push_back(temp);
	}
}

void computeCassegrainAngles(double& cAngleI, double& cAngleO, double& oAngleI, double& oAngleO)
{
	//compute the angles based on NA
	cAngleI = asin(cNAi);
	cAngleO = asin(cNAo);
	oAngleI = asin(oNAi);
	oAngleO = asin(oNAo);

	//implement a reflection-mode system if necessary
	if(opticsMode == ReflectionOpticsType){
		
		//set the condenser to match the objective
		cAngleI = oAngleI;
		cAngleO = oAngleO;

		//invert the objective
		oAngleO = pi - cAngleI;
		oAngleI = pi - cAngleO;
	}


}

int computeNl()
{
	double maxNu = EtaK.back().nu;
	double maxLambda = 10000.0f/maxNu;
	double k = 2*pi/maxLambda;
	int Nl = (int)ceil( k + 4 * exp(log(k*radius)/3) + 3 );

	return Nl;
}

void computeBArray(complex<double>* B, int Nl, int nLambda)
{
	double nu;
	complex<double> eta;
	double* Lambda = (double*)malloc(sizeof(double) * nLambda);

	//for each wavenumber nu
	for(unsigned int i=0; i<EtaK.size(); i++)
	{
		//compute information based on wavelength and material
		nu = EtaK[i].nu;
		Lambda[i] = 10000.0f/nu;
		if(applyMaterial)
			eta = complex<double>(EtaN[i].A, EtaK[i].A);
		else
			eta = complex<double>(baseIR, 0.0);

		//allocate memory for the scattering coefficients
		//complex<float>* B = (complex<float>*)malloc(sizeof(complex<float>)*Nl);		

		complex<double> IR(eta.real(), eta.imag());
		computeB(&B[i * Nl], radius, IR, Lambda[i], Nl);
	}
}

void computeOpticalParameters(double& cAngleI, double& cAngleO, double& oAngleI, double& oAngleO, double& I0, double* alpha, int Nl)
{
	computeCassegrainAngles(cAngleI, cAngleO, oAngleI, oAngleO);	

	//evaluate the incident field intensity
	I0 = 0.0;
	complex<double> Ui;

	Ui = integrateUi(cAngleI, cAngleO, oAngleI, oAngleO, 2*pi);
	I0 = Ui.real()*2*pi;

	//compute alpha (condenser integral)
	computeCondenserAlpha(alpha, Nl, cAngleI, cAngleO);
}

void gpuDetectorSpectrum(int numSamples)
{
	//integrate across the objective aperature and calculate the resulting intensity on a detector
	PD.StartTimer(SIMULATE_SPECTRUM);
	//clear the previous spectrum
	SimSpectrum.clear();	

	//compute Nl (maximum order of the spectrum)
	int Nl = computeNl();

	double* alpha = (double*)malloc(sizeof(double)*(Nl + 1));
	double cAngleI, cAngleO, oAngleI, oAngleO, I0;
	computeOpticalParameters(cAngleI, cAngleO, oAngleI, oAngleO, I0, alpha, Nl);

	//allocate space for a list of wavelengths
	int nLambda = EtaK.size();
	
	//allocate space for the 2D array (Nl x nu) of scattering coefficients
	complex<double>* B = (complex<double>*)malloc(sizeof(complex<double>) * Nl * nLambda);
	computeBArray(B, Nl, nLambda);
	

	//allocate temporary space for the spectrum
	double* I = (double*)malloc(sizeof(double) * EtaK.size());

	//compute the spectrum on the GPU
	PD.StartTimer(SIMULATE_GPU);
	cudaComputeSpectrum(I, (double*)B, alpha, Nl, nLambda, oAngleI, oAngleO, cAngleI, cAngleO, numSamples);
	PD.EndTimer(SIMULATE_GPU);

	updateSpectrum(I, I0, nLambda);

	PD.EndTimer(SIMULATE_SPECTRUM);

}

void SimulateSpectrum()
{
	if(pointDetector)
		pointSpectrum();
	else
		gpuDetectorSpectrum(objectiveSamples);
		//detectorSpectrum(objectiveSamples);
}

double absorbanceDistortion(){

	//compute the mean of the spectrum
	double sumSim = 0.0;
	for(unsigned int i=0; i<SimSpectrum.size(); i++)
	{
		sumSim += SimSpectrum[i].A;
	}
	double meanSim = sumSim/SimSpectrum.size();

	//compute the distortion (MSE from the mean)
	double sumSE = 0.0;
	for(unsigned int i=0; i<SimSpectrum.size(); i++)
	{
		sumSE += pow(SimSpectrum[i].A - meanSim, 2);
	}
	double MSE = sumSE / SimSpectrum.size();

	return MSE;
}

double intensityDistortion(){

	//compute the magnitude of the spectrum
	double sumSim = 0.0;
	for(unsigned int i=0; i<SimSpectrum.size(); i++)
	{
		sumSim += SimSpectrum[i].A * SimSpectrum[i].A;
	}
	double magSim = sqrt(sumSim);

	//compute the distortion (MSE from the mean)
	double sumSE = 0.0;
	for(unsigned int i=0; i<SimSpectrum.size(); i++)
	{
		sumSE += (SimSpectrum[i].A/magSim) * (1.0/SimSpectrum.size());
	}
	double MSE = sumSE;

	return MSE;
}

void MinimizeDistortion(){
	ofstream outFile("distortion.txt");

	//set the parameters for the distortion simulation
	double step = 0.001;

	oNAi = 0.2;
	oNAo = 0.5;

	//compute the optical parameters
	//compute Nl (maximum order of the spectrum)
	int Nl = computeNl();

	double* alpha = (double*)malloc(sizeof(double)*(Nl + 1));
	double cAngleI, cAngleO, oAngleI, oAngleO, I0;
	
	//allocate space for a list of wavelengths
	int nLambda = EtaK.size();

	//allocate temporary space for the spectrum
	double* I = (double*)malloc(sizeof(double) * EtaK.size());

	//calculate the material parameters
	//allocate space for the 2D array (Nl x nu) of scattering coefficients
	complex<double>* B = (complex<double>*)malloc(sizeof(complex<double>) * Nl * nLambda);
	computeBArray(B, Nl, nLambda);



	double D;
	double e = 0.001;
	for(double i=0.0; i<=oNAo-step; i+=step)
	{
		
		for(double o=oNAi+step; o<=1.0; o+=step)
		{
			
			
			//set the current optical parameters
			cNAi = i;
			cNAo = o;

			//compute the optical parameters
			computeOpticalParameters(cAngleI, cAngleO, oAngleI, oAngleO, I0, alpha, Nl);

			//simulate the spectrum
			cudaComputeSpectrum(I, (double*)B, alpha, Nl, nLambda, oAngleI, oAngleO, cAngleI, cAngleO, objectiveSamples);
			updateSpectrum(I, I0, nLambda);

			if(dispSimType == AbsorbanceSpecType)
			{
				if(i + e >= o || i + e >= oNAo || oNAi + e >= o || oNAi + e >= oNAo)
					D = 0.0;
				else
					D = absorbanceDistortion();
			}
			else
			{
				if(i >= o || oNAi >= oNAo)
					D=0;
				else
					D = intensityDistortion();
			}
			outFile<<D<<"       ";
		}
		outFile<<endl;
		cout<<i<<endl;
	}
	outFile.close();
}