package imseProc.proc.pitch;

import jafama.FastMath;

import java.nio.ByteOrder;
import oneLiners.OneLiners;
import org.eclipse.swt.widgets.Composite;

import binaryMatrixFile.BinaryMatrixFile;
import algorithmrepository.LinearInterpolation1D;
import algorithmrepository.LinearInterpolation2D;
import imseProc.core.ByteBufferImage;
import imseProc.core.IMSEProc;
import imseProc.core.ImagePipeController;
import imseProc.core.Img;
import imseProc.core.ImgProcPipe;
import imseProc.core.ImgSource;
import imseProc.proc.transform.FeatureTransform;
import imseProc.proc.transform.FeatureTransformCubic;

/** Polarisation to pitch processor.
 * 
 * Uses the ray-tracer to calculate a very approximate pitch angle from the input polarisation images. 
 * Might even do current one day from really naive j = dBz/dR 
 * 
 * This works on the images in (R,Z) space, so needs to be in the chain after TransformProcessor.
 * It anyway needs the geometry transform data stored in the metadata by TransformProcessor
 */
public class PitchProcessor extends ImgProcPipe {
		
	private PitchSWTController controller;
	
	private FeatureTransformCubic xform;
			
	/** Beams indices 0-3 for Q1-Q4, 1 for lat/lon, -2 for auto */
	public static final int BEAMSEL_AUTO = -1;
	public int beamSelection = 1;
	
	/** If set, calcs dBz/dR which should be something like current */ 
	private boolean calcCurrent = true;
	
	private double BPhiAt1m = 2.0; //just scales the current, we still get the correct 1/R 
	private double imgR[];
	
	//grid for ray-trcaing and calculation of angles
	private int nGX = 30, nGY = 30;
		
	private LinearInterpolation2D interpPol0[];
	private LinearInterpolation2D interpDPolDPitch[];
	
	public PitchProcessor() { 
		super(ByteBufferImage.class);
	}
	
	public PitchProcessor(ImgSource source, int selectedIndex) {
		super(ByteBufferImage.class, source, selectedIndex);		
		
	}
	
	@Override
	protected int[] sourceIndices(int outIdx) {		
		return new int[]{ outIdx }; //always 1:1
	}
	
	@Override
	protected void preCalc(boolean settingsHadChanged) {
		super.preCalc(settingsHadChanged);
		nGX = 30; nGY = 30;
		//first, we need all the geometry info and other bits
		double imgPhi[][][] = (double [][][])connectedSource.getSeriesMetaData("Transform/outputPhiClosest");
		double viewPos[] = (double[])connectedSource.getSeriesMetaData("Transform/viewPosition");
		imgR = (double[])connectedSource.getSeriesMetaData("/Transform/imageOutR");
		double imgZ[] = (double[])connectedSource.getSeriesMetaData("/Transform/imageOutZ");
		//beamIdx we can look up per-image
		
		if(imgPhi == null || viewPos == null || imgR == null || imgZ == null)
			throw new RuntimeException("Incomplete transform data.");
		
		//everything per-beam
		//double rayStartPos[][][][] = new double[4][][][];
		double pol0[][][] = new double[4][][];
		double dPolDPitch[][][] = new double[4][][];
		interpPol0 = new LinearInterpolation2D[4];
		interpDPolDPitch = new LinearInterpolation2D[4];
		
		//calc R,Z on calc grid
		LinearInterpolation1D interpR = new LinearInterpolation1D(OneLiners.linSpace(0.0, 1.0, inWidth), imgR); 
		LinearInterpolation1D interpZ = new LinearInterpolation1D(OneLiners.linSpace(0.0, 1.0, inHeight), imgZ);
		double gX[] = OneLiners.linSpace(0.0, 1.0, nGX);
		double gY[] = OneLiners.linSpace(0.0, 1.0, nGY);
		double gR[] = interpR.eval(gX);
		double gZ[] = interpZ.eval(gY);		
		
		PitchProcRaytracer rayTracer = new PitchProcRaytracer(IMSEProc.getMetaPulse(connectedSource), viewPos);
		
		double dPitch = 6.0 * Math.PI / 180;
		
		for(int iB=0; iB < 4; iB++){
			if(beamSelection != iB && beamSelection != BEAMSEL_AUTO)
				continue;
			
			int nRaysOK = 0;
			 
			//rayStartPos[iB] = new double[nGY][nGX][3];
			pol0[iB] = new double[nGY][nGX];
			dPolDPitch[iB] = new double[nGY][nGX];
			
			//LinearInterpolation2D interpPhi = new LinearInterpolation2D(imgR, imgZ, imgPhi[iB]); imgR,imgZ are not always in the correct order
			//BinaryMatrixFile.mustWrite("/tmp/phi", imgR, imgZ, imgPhi[iB], false);
			for(int iY=0; iY < nGY; iY++){				
				for(int iX=0; iX < nGX; iX++){
					
					//linear interpolate the phi of the 'near point to beam'... ish
					double fx = (gR[iX] - imgR[0]) / (imgR[1] - imgR[0]);
					double fy = (gZ[iY] - imgZ[0]) / (imgZ[1] - imgZ[0]);
					int ifx = Math.min((int)fx, imgR.length-2);
					int ify = Math.min((int)fy, imgZ.length-2);
					fx -= ifx; fy -= ify;
					double phi;
					try{
					phi= (1.0 - fx) * (1.0 - fy) * imgPhi[iB][ify][ifx] +
									fx * (1.0 - fy) * imgPhi[iB][ify+1][ifx] +
									(1.0 - fx) * fy * imgPhi[iB][ify][ifx+1] +
									fx * fy * imgPhi[iB][ify+1][ifx+1];
					}catch(RuntimeException err){
						err.printStackTrace();
						throw err;
					}
					//double phi = interpPhi.eval(gR[iX], gZ[iY]);
					if(Double.isNaN(phi)){
						pol0[iB][iY][iX] = Double.NaN;
						dPolDPitch[iB][iY][iX] = Double.NaN;
						continue;
					}
					
					double rayStartPos[] = new double[]{							
						gR[iX] * FastMath.cos(phi),
						gR[iX] * FastMath.sin(phi),
						gZ[iY] 
					};
					
					double ret[] = rayTracer.getPolPitchInfo(iB, rayStartPos, dPitch);
					pol0[iB][iY][iX] = ret[0];
					dPolDPitch[iB][iY][iX] = ret[1];
					
					if(!Double.isNaN(ret[0]))
						nRaysOK++;
				}	
			}
			
			interpPol0[iB] = new LinearInterpolation2D(gX, gY, pol0[iB]);
			interpDPolDPitch[iB] = new LinearInterpolation2D(gX, gY, dPolDPitch[iB]);
			
			System.out.println("PitchProcessor: " + nRaysOK + " of " + (nGX*nGY) + " rays hit for beam " + iB);
		}
		
		//someone might want these
		setSeriesMetaData("Pitch/pol0", pol0);
		setSeriesMetaData("Pitch/dPolDPitch", dPolDPitch);
		
		rayTracer.destroy();
	}
	
	
	/** Image allocation */
	protected boolean checkOutputSet(int nImagesIn){		
		outWidth = inWidth;
		outHeight = inHeight;
	    int nImagesOut = nImagesIn;
		
        //allocate or re-use
        ByteBufferImage newOutSet[] = ByteBufferImage.checkBulkAllocation(this, (ByteBufferImage[])imagesOut, outWidth, outHeight, ByteBufferImage.DEPTH_DOUBLE, nImagesOut, ByteOrder.LITTLE_ENDIAN);
        if(newOutSet != imagesOut){
        	imagesOut = newOutSet;
        	return true;
        }
        
        return false;
	}
	
	@Override
	protected boolean doCalc(Img imageOutG, Img[] sourceSet, boolean settingsHadChanged) throws InterruptedException {
		Img imageIn = sourceSet[0];		
		ByteBufferImage imageOut = (ByteBufferImage)imageOutG;
		
		
		int beamIdx;
		if(beamSelection == BEAMSEL_AUTO){
			beamIdx = (Integer)getImageMetaData("beamIndex", imageIn.getSourceIndex());
			
		}else{
			beamIdx = beamSelection;
		}
				
		imageIn.startReading();									
		try{
			for(int iY=0; iY < inHeight; iY++) {
				for(int iX=0; iX < inWidth; iX++) {
					double val = imageIn.getPixelValue(iX, iY) * Math.PI / 180;
					
					val = -19.4* Math.PI / 180 - val; //convert from thetaA (µ corrected measured rel Savart) to thetaF (full rel to IMSE frame)
					
					if(!Double.isNaN(val)){
						//System.out.println("dbg sucks");
					}
					
					double x = iX / (inWidth - 1.0);
					double y = iY / (inHeight - 1.0);
					
					double pol0 = interpPol0[beamIdx].eval(x, y);
					double dPolDPitch = interpDPolDPitch[beamIdx].eval(x, y);
					
					val = (val - pol0) / dPolDPitch; //convert to Bz/Bphi
					
					double Bphi = BPhiAt1m / imgR[iX];					
					val = val * Bphi; //covnert to Bzish
					
					imageOut.setPixelValue(iX, iY, val);
				}
				//differential and calc the current
				if(calcCurrent){
					for(int iX=0; iX < (inWidth-1); iX++) {
						double mu0 = 4*Math.PI*1e-7;
						double bz0 = imageOut.getPixelValue(iX, iY);
						double bz1 = imageOut.getPixelValue(iX+1, iY);
						double j = (bz1 - bz0) / (imgR[iX+1] - imgR[iX]) / mu0; 
						imageOut.setPixelValue(iX, iY, j / 1e3); //in kA
					}
					imageOut.setPixelValue(inWidth-1, iY, 0);
					int n=15;
					double row[] = new double[inWidth];
					for(int iX=n; iX < (inWidth-n-2); iX++) {						
						for(int dX=-n; dX <= n; dX++){
							row[iX] += imageOut.getPixelValue(iX + dX, iY) / (2*n+1);							
						}
					}
					for(int iX=0; iX < inWidth; iX++) {
						imageOut.setPixelValue(iX, iY, row[iX]);
					}
				}
			}
		}finally{
			imageIn.endReading();
		}
		
		return true;
	}
	
	@Override
	public void notifySourceChanged() {
		super.notifySourceChanged();
		if(autoCalc)
			calc();
	}

	@Override
	public void imageChanged(int idx) { 
 		if(autoCalc)
			calc();
 	}
	
	@Override
	public int getNumImages() {
		return (imagesOut != null) ? imagesOut.length : 0;
	}

	@Override
	public Img getImage(int imgIdx) {
		return (imgIdx >= 0 && imgIdx < imagesOut.length) ? imagesOut[imgIdx] : null;
	}
	
	@Override
	public Img[] getImageSet() { return imagesOut; }

	@Override
	public PitchProcessor clone() { return new PitchProcessor(connectedSource, getSelectedSourceIndex());	}

	@Override
	/** For the sink side */
	public ImagePipeController createPipeController(Class interfacingClass, Object args[], boolean asSink) {
		if(interfacingClass == Composite.class){
			controller = new PitchSWTController((Composite)args[0], (Integer)args[1], this, asSink);
			controllers.add(controller);
			return controller;
		}
		return null;
	}
	
	public void setAutoCalc(boolean autoCalc) {
		if(!this.autoCalc && autoCalc){
			this.autoCalc = true;
			updateAllControllers();
			calc();			
		}else{
			this.autoCalc = autoCalc;
			updateAllControllers();
		}
	}
	public boolean getAutoCalc() { return this.autoCalc; }

	public int getCalcProgress() { return calcProgress; }
	
	@Override
	public void setSource(ImgSource source) {
		super.setSource(source);
		if(autoCalc)
			calc();
	}
	
	@Override
	public void destroy() {		
		super.destroy();
	}
	
	public FeatureTransform getTransform(){ return xform; }
	public void pointsMapModified(){
		settingsChanged = true;
		updateAllControllers();
		if(autoCalc)
			calc();
	}

	public void invalidate() {
		settingsChanged = true;
	}

	public void setBeamSelection(int beamSelection) {
		this.beamSelection = beamSelection;
		settingsChanged = true;
		if(autoCalc)
			calc();
	}

	public int getCalcNX() { return this.nGX; }
	public int getCalcNY() { return this.nGY; }
	public void setCalcGrid(int nX, int nY) {
		nGX = nX;
		nGY = nY;		
		settingsChanged = true;
		if(autoCalc)
			calc();		
	}

	public boolean getCalcCurrent() { return calcCurrent;	}
	public void setCalcCurrent(boolean enable) {
		this.calcCurrent = enable;
		settingsChanged = true;
		if(autoCalc)
			calc();	
	}

}