#include "microscope.h"

#include "rts/cuda/error.h"
#include "rts/tools/progressbar.h"
#include "rts/cuda/timer.h"
#include "dataTypes.h"
#include "rts/visualization/colormap.h"

#include <QImage>

__global__ void bandpass(bsComplex* U, int uR, int vR, ptype du, ptype dv, ptype NAin, ptype NAout, ptype lambda)
{

	//get the current coordinate in the plane slice
	int iu = blockIdx.x * blockDim.x + threadIdx.x;
	int iv = blockIdx.y * blockDim.y + threadIdx.y;

	//make sure that the thread indices are in-bounds
	if(iu >= uR || iv >= vR) return;

	//compute the index (easier access to the scalar field array)
	int i = iv*uR + iu;

	ptype u, v;

	if(iu <= uR / 2)
        u = (ptype)iu * du;
    else
        u = -(ptype)(uR - 1 - iu) * du;

    if(iv <= vR / 2)
        v = (ptype)iv * dv;
    else
        v = -(ptype)(vR - 1 - iv) * dv;

    ptype fmag = sqrt(u*u + v*v);

    if(fmag < NAin / lambda || fmag > NAout / lambda)
        U[i] = 0;
	//U[i] = U[i];
}

microscopeStruct::microscopeStruct()
{
	scalarSim = true;
	D = NULL;
	Di = NULL;
}

void microscopeStruct::init()
{

    nf.scalarSim = scalarSim;
	//Ud.scalarField = scalarSim;
	//Ufd.scalarField = scalarSim;

	//Ud.init_gpu();
	//Ufd.init_gpu();

	//initialize the near field
	nf.init();

	//allocate space for the detector
	D = new scalarslice(Ud.R[0] / ss, Ud.R[1] / ss);
	Di = new scalarslice(Ud.R[0] / ss, Ud.R[1] / ss);

    //clear the detector
    clearDetector();

}

void microscopeStruct::destroy()
{
	delete D;
	D = NULL;

	delete Di;
	Di = NULL;

	Ud.kill_gpu();
	Ufd.kill_gpu();

	//destroy the near field
	nf.destroy();

}

void microscopeStruct::applyBandpass()
{
    //This function applies the objective bandpass to the near field
    //The near field structure stores the results, in order to save memory

    //first convert the near field to an angular spectrum (FFT)
    nf.U.toAngularSpectrum();

    //create one thread for each pixel of the field slice
	dim3 dimBlock(SQRT_BLOCK, SQRT_BLOCK);
	dim3 dimGrid((nf.U.R[0] + SQRT_BLOCK -1)/SQRT_BLOCK, (nf.U.R[1] + SQRT_BLOCK - 1)/SQRT_BLOCK);

	//compute the step size in the frequency domain
	ptype du = 1.0 / (nf.pos.X.len());
	ptype dv = 1.0 / (nf.pos.Y.len());

    //apply the objective band-pass filter
	bandpass<<<dimGrid, dimBlock>>>(nf.U.x_hat, nf.U.R[0], nf.U.R[1], du, dv, objective[0], objective[1], nf.lambda);

    //convert the near field image back to the spatial domain
    //  (this is the field at the detector)
	nf.U.fromAngularSpectrum();
}

void microscopeStruct::getFarField()
{
    //Compute the Far Field image of the focal plane

    //clear the memory from previous detector fields
    //Ud.kill_gpu();
    //Ufd.kill_gpu();

	//first crop the filtered near-field image of the source and scattered fields
	Ud = nf.U.crop(padding * Ud.R[0], padding * Ud.R[1], Ud.R[0], Ud.R[1]);
	Ufd = nf.Uf.crop(padding * Ufd.R[0], padding * Ufd.R[1], Ufd.R[0], Ufd.R[1]);

}

void microscopeStruct::integrateDetector()
{
	Ud.IntegrateAndResample(D, ss);
	Ufd.IntegrateAndResample(Di, ss);
}

void microscopeStruct::clearDetector()
{
	//zero-out the detector
	D->clear();
	Di->clear();
}

//flag for a vector simulation
void microscopeStruct::setPos(bsPoint pMin, bsPoint pMax, bsVector normal)
{
	pos = rts::quad<ptype, 3>(pMin, pMax, normal);
}

void microscopeStruct::setRes(int x_res, int y_res, int pad, int supersampling)
{
	padding = pad;
	ss = supersampling;

	Ufd.R[0] = Ud.R[0] = x_res * ss;
	Ufd.R[1] = Ud.R[1] = y_res * ss;
}

void microscopeStruct::setNearfield()
{
	//sets the values for the near field in order to create the specified detector image

	//compute the size of the near-field slice necessary to create the detector image
	nf.pos = pos * (padding * 2 + 1);

	//compute the resolution of the near-field slice necessary to create the detector image
	nf.setRes(Ud.R[0] * (padding * 2 + 1), Ud.R[1] * (padding * 2 + 1));

}

__global__ void calc_absorbance(ptype* A, ptype* D, ptype* Di, int N)
{
    //compute the index for this thread
	int i = blockIdx.x * blockDim.x + threadIdx.x;
	if(i >= N) return;

	A[i] = -log10(D[i] / Di[i]);

}

scalarslice microscopeStruct::getAbsorbance()
{
	//compute the magnitude of the field at each rtsPoint in the slice

	//create a scalar slice at the same resolution as the field
	scalarslice* A = new scalarslice(D->R[0], D->R[1]);

	//compute the total number of values in the slice
	unsigned int N = D->R[0] * D->R[1];
	int gridDim = (N+BLOCK-1)/BLOCK;

	calc_absorbance<<<gridDim, BLOCK>>>(A->S, D->S, Di->S, N);

	return *A;
}

__global__ void calc_transmittance(ptype* A, ptype* D, ptype* Di, int N)
{
    //compute the index for this thread
	int i = blockIdx.x * blockDim.x + threadIdx.x;
	if(i >= N) return;

	A[i] = D[i] / Di[i];

}

scalarslice microscopeStruct::getTransmittance()
{
	//compute the magnitude of the field at each rtsPoint in the slice

	//create a scalar slice at the same resolution as the field
	scalarslice* T = new scalarslice(D->R[0], D->R[1]);

	//compute the total number of values in the slice
	unsigned int N = D->R[0] * D->R[1];
	int gridDim = (N+BLOCK-1)/BLOCK;

	calc_transmittance<<<gridDim, BLOCK>>>(T->S, D->S, Di->S, N);

	return *T;
}

scalarslice microscopeStruct::getIntensity()
{
    //create a scalar slice at the same resolution as the field
	scalarslice* I = new scalarslice(D->R[0], D->R[1]);

	HANDLE_ERROR(cudaMemcpy(I->S, D->S, sizeof(ptype) * D->R[0] * D->R[1], cudaMemcpyDeviceToDevice));

	return *I;

}

void microscopeStruct::SimulateScattering()
{
	nf.Simulate();
}

void microscopeStruct::SimulateImaging()
{
	applyBandpass();
    getFarField();
    integrateDetector();
}

void microscopeStruct::Simulate()
{
    SimulateScattering();
	SimulateImaging();
}

void microscopeStruct::SimulateExtendedSource()
{

	clearDetector();

    //for each source in the source list
	int npts = focalPoints.size();
	float t=0;
    for(int i = 0; i<npts; i++)
	{
		nf.focus = focalPoints[i].f;
		nf.A = focalPoints[i].A;

        gpuStartTimer();
		Simulate();

		t += gpuStopTimer();

		rtsProgressBar((double)(i+1)/(double)npts * 100);
		//unsigned char c;
		//cin>>c;
	}
	if(verbose)
	{
        cout<<endl;
        cout<<"Time per source: "<<t/npts<<"ms"<<endl;
    }

}

void microscopeStruct::LoadExtendedSource(std::string filename)
{
    //this function loads an image of an extended source and creates a list of corresponding point sources

    QImage sourceImage(filename.c_str());

    //get the resolution of the image (the boundary is scaled to match the detector)
    int Rx = sourceImage.width();
    int Ry = sourceImage.height();

    //for each pixel
    int x, y;
    float u, v;
    for(x=0; x<Rx; x++)
        for(y=0; y<Ry; y++)
        {
            //create a new point source
            sourcePoint p;

            //compute the coordinate of the focal point
            u = (ptype)x / (ptype)Rx;
            v = (ptype)y / (ptype)Ry;

            p.f = pos(u, v);

            //get the amplitude of the focal point
            QRgb rgb = sourceImage.pixel(x, y);
            //float A = qGray(rgb);
			if(qGray(rgb) != 0)
			{
				p.A = (ptype) qGray(rgb) / 255;

				//insert the point source into the list
				focalPoints.push_back(p);
			}
        }
}

std::string microscopeStruct::toStr()
{
	stringstream ss;
	ss<<nf.toStr();

	ss<<"----------Optics--------------"<<endl<<endl;
	ss<<"Objective NA: "<<objective[0]<<" to "<<objective[1]<<endl;
	return ss.str();


}