package fusionOptics.optimisationMulti;

import com.sun.org.apache.xml.internal.security.algorithms.Algorithm;

import jafama.FastMath;
import otherSupport.RandomManager;
import algorithmrepository.Algorithms;
import binaryMatrixFile.BinaryMatrixWriter;
import fusionOptics.Util;
import fusionOptics.collection.IntersectionProcessor;
import fusionOptics.drawing.RayDrawer;
import fusionOptics.drawing.SVGRayDrawing;
import fusionOptics.drawing.VRMLDrawer;
import fusionOptics.surfaces.Plane;
import fusionOptics.tracer.Tracer;
import fusionOptics.types.Element;
import fusionOptics.types.Intersection;
import fusionOptics.types.Pol;
import fusionOptics.types.RaySegment;


/** A Bundle of rays fired in the optimiser */
public class RayBundle implements IntersectionProcessor {
	private final int minimumHitsForStatistics = 10;
	 
	
	/** The optimiser main class, for us to get info during run */
	protected OptimiseMulti optimiser;
	
	/** Plane on which to assess focus of the ray bundle */
	private Plane imagePlane;	
	
	/** All rays are fired at the same wavelength */
	private double wavelength;
	
	/** Initial rays used for evaluation - the same ones are used for each eval, to remove random noise */
	protected double rayStarts[][];
	protected double rayDirs[][];	
	protected int nRays;
	
	/** Where they want the rays to hit the image plane (R,U) */
	private double targetPos[] = null;
	
	/** Sharpness cost weighting for this bundle */
	private double sharpnessWeight = 1.0; 
	/** Intensity cost weighting for this bundle */	
	private double intensityWeight = 0.3;
	/** Target cost weighting for this bundle */	
	private double targetWeight = 0.0;
		
	/** Our bundle ID */
	private int bundleID;
	
	/** Sums of hits and positions and squares on image plane (R,U) for last evaluation*/
	private int evalNHits;	
	private double lastEvalSumR, lastEvalSumU;
	private double lastEvalSumR2, lastEvalSumU2;
	
	/** Spot statistics for last evaluation */
	private double lastEvalFWHM;		
	private double lastEvalTargetDist; 
	private double lastEvalIntensityLoss;

	public double getLastEvalFWHM() {	return lastEvalFWHM;	}
	public double getLastEvalTargetDist() { return lastEvalTargetDist;	}
	public double getLastEvalIntensityLoss() { return lastEvalIntensityLoss;	}

	public RayBundle(int bundleID) {
		this.bundleID = bundleID;
	}
	
	/** Initialise as rays diverging from a series of points across the given vector
	 * @param rayTarget		Element to fire rays at
	 * @param nPoints		Number of points from which to fire rays
	 * @param centralPoint	The central imaging point
	 * @param pointsDelta	Direction and spacing of points 
	 * @param wavelength	Wavelength of rays
	 * @param nRaysPerPoint			Number of rays from each point			
	 */
	public static RayBundle[] initRaysImagingAlongLine(Element rayTarget, Plane imagePlane, double centralPoint[], double pointsDelta[], 
													int nPoints, double wavelength, int nRaysPerPoint, int bundleID0){
		RayBundle[] rayBundles = new RayBundle[nPoints];		
		for(int iP=0; iP < nPoints; iP++){
			double objectPoint[] = new double[]{
				centralPoint[0] + (iP - nPoints/2.0) * pointsDelta[0],
				centralPoint[1] + (iP - nPoints/2.0) * pointsDelta[1],
				centralPoint[2] + (iP - nPoints/2.0) * pointsDelta[2] 
			};
			rayBundles[iP] = new RayBundle(bundleID0 + iP);
			rayBundles[iP].initRaysImaging(rayTarget, objectPoint, wavelength, nRaysPerPoint);
			rayBundles[iP].setImagePlane(imagePlane);
		}
		return rayBundles;
	}
	
	/** Initialise the evaluation ray tracing using rays diverging from a point 
	 * @param rayTarget		Element to fire rays at
	 * @param objectPoints	Points from which to fire rays - double[posIdx][x/y/z]
	 * @param wavelength	Wavelength of rays
	 * @param nRaysPerPoint			Number of rays from each point			
	 */
	public void initRaysImaging(Element rayTarget, double sourcePoint[], double wavelength, int nRays){
		this.wavelength = wavelength;
		this.nRays = nRays;
		
		rayStarts = new double[nRays][];
		rayDirs = new double[nRays][];
		for(int iR = 0; iR < nRays; iR++){
			rayStarts[iR] = sourcePoint.clone();
			rayDirs[iR] = Tracer.generateRandomRayTowardSurface(sourcePoint, rayTarget);
		}
	}
	
	/** Initialise the evaluation ray tracing using sets of parallel rays at various angles
	 * @param rayTarget		Element to fire rays at.
	 * @param rayDir0		Direction of rays at 0 angle - usually the optic axis.
	 * @param maxAngle		Maximum angle away from rayDir0
	 * @param wavelength	Wavelength of rays
	 * @param nSets			Number of ray bunches, at different angles
	 * @param nRaysPerSet	Number of rays at each angle
	 */
	public static RayBundle[] initRaysParallelMultiAngle(Element rayTarget, Plane imagePlane, double rayDir0[], double maxAngle, double rayLength, 
													double wavelength, int nBundles, int nRaysPerSet, int bundleID0){

		double cosMaxTheta = Math.cos(maxAngle);
		double aU[] = Util.createPerp(rayDir0);
		double bU[] = Util.cross(rayDir0, aU);		
		Util.reNorm(aU);
		Util.reNorm(bU);
		
		RayBundle rayBundles[] = new RayBundle[nBundles];
		for(int iS=0; iS < nBundles; iS++){

			//double cosTheta=1.0;
			//if(nSets > 1) //generate an incoming direction, uniformly distributed over the specified angle range
				//cosTheta = 1 - RandomManager.instance().nextUniform(0, 1) * (1 - cosMaxTheta);
			
			//double theta = -maxAngle + iS * 2 * maxAngle / (nSets - 1);
			double theta = iS * maxAngle / (nBundles - 1);
			 			
			//double sinTheta = FastMath.sqrt(1 - cosTheta*cosTheta);
			double cosTheta = FastMath.cos(theta);
			double sinTheta = FastMath.sin(theta);
			
			double phi = 0;//RandomManager.instance().nextUniform(0, 1) * 2 * Math.PI;
			
			//generate in coord sys (a,b,c) with c as axis toward target 
			double a = sinTheta * FastMath.cos(phi);
			double b = sinTheta * FastMath.sin(phi);
			double c = cosTheta;			
			 
			double rayDir[] = new double[]{
				c * rayDir0[0] + a * aU[0] + b * bU[0],
				c * rayDir0[1] + a * aU[1] + b * bU[1],
				c * rayDir0[2] + a * aU[2] + b * bU[2]
			};
			
			rayBundles[iS] = new RayBundle(bundleID0 + iS);
			rayBundles[iS].initRaysParallel(rayTarget, rayDir, aU, bU, rayLength, wavelength, nRaysPerSet);
			rayBundles[iS].setImagePlane(imagePlane);
		}
		return rayBundles;
	}
	
	/** Initialise the rays at a fixed angle
	 * @param rayTarget		Element to fire rays at.
	 * @param rayDir0		Direction of rays at 0 angle - usually the optic axis.
	 * @param maxAngle		Maximum angle away from rayDir0
	 * @param wavelength	Wavelength of rays
	 * @param nSets			Number of ray bunches, at different angles
	 * @param nRaysPerSet	Number of rays at each angle
	 * @param aU, bU		Unit vectors defining a plane perp to the central ray angle (not necc. this one) (see initRaysParallelMultiAngle() )
	 */
	public void initRaysParallel(Element rayTarget, double rayDir[], double aU[], double bU[], double rayLength, double wavelength, int nRays){
		this.wavelength = wavelength;
		this.nRays = nRays;
		
		rayStarts = new double[nRays][];
		rayDirs = new double[nRays][];
		
		//generate multiple rays in this direction, distributed randomly accross the bouding
		//sphere of the target element
		for(int iR=0; iR < nRays; iR++){
			rayDirs[iR] = rayDir.clone();
			
			//generate a random point on the circle made by the projection of the target bounding sphere
			// on to the perpendicular plane to rayDir0
			
			double r = rayTarget.getBoundarySphereRadius() * FastMath.sqrt(RandomManager.instance().nextUniform(0, 1));
			double ang = RandomManager.instance().nextUniform(0, 1) * 2*Math.PI;
			double lA = r * FastMath.sin(ang);
			double lB = r * FastMath.cos(ang);
			double cc[] = rayTarget.getBoundarySphereCentre();
			
			rayStarts[iR] = new double[]{
					cc[0] + lA * aU[0] + lB * bU[0] - rayLength * rayDir[0],
					cc[1] + lA * aU[1] + lB * bU[1] - rayLength * rayDir[1],
					cc[2] + lA * aU[2] + lB * bU[2] - rayLength * rayDir[2],
			};
		}		
	}	
	
	/** Trace this ray bundle and evaluate the cost function */
	public double traceAndEvaluate(RayDrawer rayDrawer) {
		
		if(imagePlane == null)
			throw new IllegalArgumentException("imagePlane not set for ray bundle '"+bundleID+"'");
		
		evalNHits = 0;
		lastEvalSumR = 0;
		lastEvalSumR2 = 0;
		lastEvalSumU = 0;
		lastEvalSumU2 = 0;
		
		if(rayDrawer != null)
			rayDrawer.startGroup("bundle_" + bundleID);				 
		
		for(int rayNo=0; rayNo < rayStarts.length; rayNo++){ //for each ray
			RaySegment ray = new RaySegment();
			
			ray.startPos = rayStarts[rayNo];
			ray.dir = rayDirs[rayNo];

			ray.E0 = new double[][]{{ 1, 0, 0, 0 }};
			ray.length = Double.POSITIVE_INFINITY;
			ray.up = Util.createPerp(ray.dir);
			
			ray.wavelength = wavelength;
			
			Tracer.trace(optimiser.getAllElements(), ray, Integer.MAX_VALUE, 1e-3, false);
			
			int n = ray.processIntersections(imagePlane, this);
			
			if(n > 0 && rayDrawer != null)
				rayDrawer.drawRay(ray);
			
			Pol.recoverAll();
		}
		
		if(rayDrawer != null)
			rayDrawer.endGroup();
		
		// process stats		
		double meanR = lastEvalSumR / evalNHits;
		double meanU = lastEvalSumU / evalNHits;
		double stdR2 = (lastEvalSumR2 - lastEvalSumR*lastEvalSumR/evalNHits) / evalNHits;
		double stdU2 = (lastEvalSumU2 - lastEvalSumU*lastEvalSumU/evalNHits) / evalNHits;
		
		lastEvalFWHM = 2.35 * FastMath.sqrt(stdR2+stdU2);		
		lastEvalTargetDist = (targetPos != null) ? 
				FastMath.sqrt(FastMath.pow2(meanR - targetPos[0]) + FastMath.pow2(meanU - targetPos[1])) : 0;
		lastEvalIntensityLoss = ((double)nRays - evalNHits) / nRays;
		
		double cost = sharpnessWeight * lastEvalFWHM + targetWeight * lastEvalTargetDist + intensityWeight * lastEvalIntensityLoss;
				
		if(evalNHits <= minimumHitsForStatistics){
			System.err.println("WARNING: Ray bundle "+bundleID+" had too few hits ("+evalNHits+"), not adding to statistics.");
			cost += 100.0 * FastMath.exp(- evalNHits / 10.0);
		}
		

		return cost;
	}

	@Override
	public void nextIntersection(Intersection hit) {
		
		double posRU[] = imagePlane.posXYZToPlaneRU(hit.pos);
		
		
		BinaryMatrixWriter hitsOut = optimiser.getHitsOutWriter();
		if(hitsOut != null){
			int iterNo = optimiser.getIterationNo();
			if(hitsOut != null)
				hitsOut.writeRow(iterNo, bundleID, evalNHits, posRU, optimiser.getCurrentEvalParams());
		}
		
		evalNHits++;
		lastEvalSumR += posRU[0];
		lastEvalSumU += posRU[1];
		lastEvalSumR2 += posRU[0]*posRU[0];
		lastEvalSumU2 += posRU[1]*posRU[1];
	}
	
	
	@Override
	public String toString() {
		
		double meanR = lastEvalSumR / evalNHits;
		double meanU = lastEvalSumU / evalNHits;
		
		return "RayBundle[ID="+bundleID+
					"]: start[0]={" + rayStarts[0][0] + ", " +rayStarts[0][1] + ", " +rayStarts[0][2] + 
					" }, lastEval: n="+evalNHits + ", mean={" + meanR + ", " + meanU + "}, fwhm = " + lastEvalFWHM + ", Iloss = " + lastEvalIntensityLoss + ", targDist = " + lastEvalTargetDist;   
	}
	
	/** Sets the sharpness weight. 
	 * The FWHM of the spot is multiplied by this and added to the cost */
	public void setSharpnessWeight(double sharpnessWeight) { this.sharpnessWeight = sharpnessWeight;	}
	
	/** Sets the intensity weight. 
	 * The number of rays that do NOT hit the target is multiplied by this and added to the cost */
	public void setIntensityWeight(double intensityWeight) { this.intensityWeight = intensityWeight;	}
	
	/** Sets the target weight. 
	 * The distance of the mean hit from the target position is multiplied by this and added to the cost */
	public void setTargetWeight(double targetWeight) { this.targetWeight = targetWeight;	}
	
	/** Sets the plane on which to test the imaging/intesity/target */
	public void setImagePlane(Plane imagePlane) { this.imagePlane = imagePlane; }
	
	/** Sets the target position on the image plane (R,U) */
	public void setTargetPos(double posRU[]){ this.targetPos = posRU; }

	/** Returns the source position of the first ray */
	public double[] getSourcePos() { return rayStarts[0]; }

	
}
