package fusionOptics.optimisationMulti;

import java.util.ArrayList;
import javax.print.attribute.standard.Media;

import oneLiners.OneLiners;
import otherSupport.ColorMaps;

import binaryMatrixFile.BinaryMatrixWriter;

import seed.digeom.FunctionND;
import seed.digeom.IDomain;
import seed.digeom.RectangularDomain;
import seed.optimization.HookeAndJeeves;
import seed.optimization.Optimizer;

import fusionOptics.Util;
import fusionOptics.collection.IntersectionProcessor;
import fusionOptics.drawing.RayDrawer;
import fusionOptics.drawing.SVGRayDrawing;
import fusionOptics.drawing.VRMLDrawer;
import fusionOptics.optimisation.OptimiseOptic;
import fusionOptics.surfaces.Dish;
import fusionOptics.surfaces.Plane;
import fusionOptics.types.Element;
import fusionOptics.types.Intersection;
import fusionOptics.types.Material;
import fusionOptics.types.Optic;
import fusionOptics.types.Surface;

/** New optimisation system for optimising lots of parameters of lots of optics against lots of costs. */
public class OptimiseMulti extends FunctionND {

	private ArrayList<Parameter> parameters = new ArrayList<Parameter>();
	private ArrayList<RayBundle> rayBundles = new ArrayList<RayBundle>();
	private Element allElements;
	
	
	/** Debugging output */
	private String hitsOutputFile = null;
	private boolean hitsWriteEnable = false;
	private BinaryMatrixWriter hitsOutWriter = null;	
	/** Current iteration number */
	protected int iterationNo = -1;
	/** Current eval params */
	protected double currentEvalParams[];
	
	/** Graphics output */
	private int drawingIterationPeriod;
	private int drawingSkipRays;
	private String drawingFileNamePrefix;
	private Class drawingType = null;
	private String drawingExtraVRML = null;
	
	/** The optimiser */
	private Optimizer opt;
	
	/** Number of rays that hit their image planes on the last evaluation */
	public int nRaysInImageLast;
	
	public OptimiseMulti() {
		opt = new HookeAndJeeves(this); //default, not great at more than a few dimensions
		
		//HookeAndJeeves opt = new HookeAndJeeves(this);
		//SuperGeneticMk2 opt = new SuperGeneticMk2(4, 1, true);
		
		/*MSHGPPConfig cfg = new MSHGPPConfig();
		cfg.populationSize = 20;
		cfg.numChildren = 2;
		cfg.initFromCurrent = true;
		MegaSuperHyperGeneticProblemPacifier opt = new MegaSuperHyperGeneticProblemPacifier(cfg);
		//*/
		
		/*IStoppingCondition lineSearchStop = new MaxIterCondition(50);
		GoldenSection lineOptimizer = new GoldenSection(lineSearchStop);
		lineOptimizer.setInitialBracketMethod(new BracketingByParameterSpace());
		Optimizer opt = new LineSearchOptimizer(null, new ConjugateGradientDirectionFR(), lineOptimizer);
		//*/
	}
	
	/** Sets the MegaOptimisation optimiser (seed.optimiser...) */
	public void setOptimiser(Optimizer opt) {
		this.opt = opt;
	}
	
	public double[] optimise(int nIters){
		return optimise(getParams(), nIters);
	}
	
	public double[] optimise(double pInit[], int nIters){
		if(allElements == null) throw new IllegalArgumentException("No elements to trace, please call OptimiseMulti.setTracingElements() first");
		if(rayBundles.size() <= 0) throw new IllegalArgumentException("No rays, please call OptimiseMulti.addRayBundle() at least once ");
		if(parameters.size() <= 0) throw new IllegalArgumentException("No parameters, please call OptimiseMulti.addParameter() at least once ");
		
		if(opt == null){
			opt = new HookeAndJeeves(this); //simple default, otherwise use some of the following from outside:
		}
		
		opt.setObjectiveFunction(this);
		
		opt.init(pInit);
		
		//do the optimiser init before the SVG setup, so it doesnt draw the init exploration
		if(hitsOutputFile != null)
			hitsOutWriter = new BinaryMatrixWriter(hitsOutputFile, pInit.length + 5);
				
		iterationNo = -1;
		hitsWriteEnable = true;
		double initCost = (drawingType == null) ? eval(pInit) : evalAndDraw(pInit);
		hitsWriteEnable = false; //only enable hits writing on evals we do out here, not that the optimiser does, other we get all the 'attempts', no matter how bad
		
		for(iterationNo=0; iterationNo < nIters; iterationNo++){
				
			opt.refine();
			double p[] = opt.getCurrentPos();
			
			//re-run the eval, because the optimisers don't guarantee the optimal was eval'd last
			double cost;
			
			hitsWriteEnable = true;
			if(drawingType != null && (iterationNo % drawingIterationPeriod) == 0){
				cost = evalAndDraw(p);				
			}else
				cost = eval(p); 
			hitsWriteEnable = false;
				
			System.out.print("i="+ iterationNo + "\tcost=" + cost + "\tnRaysInImageLast=" + nRaysInImageLast + "\tp=");
			for(int j=0; j < p.length; j++){
				System.out.print("\t" + p[j]);
			}
			System.out.println();
			
			dumpRayStats();			
			opt.dumpStatus();
		}
		
		double p[] = opt.getCurrentPos();
		setParams(p);
		
		System.out.print("Init: cost=" + initCost + ", p = "); OneLiners.dumpArray(pInit);
		System.out.print("End:  cost=" + eval(p) + ", p = "); OneLiners.dumpArray(p);

		return p;
	}

	private void dumpRayStats() {
		double sumFWHM=0, sumTargDist=0, sumIntensityLoss=0;
		int nBundles = rayBundles.size();
		for(RayBundle rayBundle : rayBundles){
			sumFWHM += rayBundle.getLastEvalFWHM();
			sumTargDist += rayBundle.getLastEvalTargetDist();
			sumIntensityLoss += rayBundle.getLastEvalIntensityLoss();
		}
		System.out.println("i=" + iterationNo 
				+ " <fwhm>=" + (sumFWHM/nBundles)
				+ " <dist>=" + (sumTargDist/nBundles)
				+ " <Iloss>=" + (sumIntensityLoss/nBundles));
	}

	private void setParams(double[] p) {
		int n = parameters.size();
		if(p.length != n)
			throw new IllegalArgumentException("Params are the wrong length");
			
		int i=0;
		for(Parameter parameter : parameters){
			parameter.set(p[i]);
			i++;
		}
	}

	private double[] getParams() {
		int n = parameters.size();
		double p[] = new double[n];
			
		int i=0;
		for(Parameter parameter : parameters){
			p[i] = parameter.get();
			i++;
		}
		return p;
	}

	@Override
	public double eval(double[] p) {
		currentEvalParams = p;
		setParams(p);
		
		double cost = 0;		
		for(RayBundle rayBundle : rayBundles){
			rayBundle.optimiser = this;
			
			cost += rayBundle.traceAndEvaluate(null);
		}
		
		return cost;
	}
	
	public double evalAndDraw(double p[]){
		RayDrawer rayDrawer;
		if(drawingType == SVGRayDrawing.class){
			rayDrawer = new SVGRayDrawing(drawingFileNamePrefix + "-iter_" + iterationNo + ".svg", Util.getBoundingBox(allElements), true);			
			
		}else if(drawingType == VRMLDrawer.class){
			rayDrawer = new VRMLDrawer(drawingFileNamePrefix + "-iter_" + iterationNo + ".vrml");
			if(drawingExtraVRML != null){
				((VRMLDrawer)rayDrawer).addVRML(drawingExtraVRML);
			}
		}else 
			throw new IllegalArgumentException("Unknown drawing type");
		rayDrawer.setSkipRays(drawingSkipRays);

		double colorMap[][] = ColorMaps.jet(rayBundles.size());
		
		if(rayDrawer instanceof SVGRayDrawing){
			((SVGRayDrawing)rayDrawer).generateLineStyles(colorMap, allElements.getBoundarySphereRadius() / 5000);
			((SVGRayDrawing)rayDrawer).setRayColour(0);
		}
	
		rayDrawer.drawElement(allElements);
				
		double cost = 0;		
		int bundleNum = 0;
		for(RayBundle rayBundle : rayBundles){
			rayBundle.optimiser = this;
		
			int colNum = bundleNum;
			
			if(rayDrawer instanceof SVGRayDrawing){					
				((SVGRayDrawing)rayDrawer).setRayColour(colNum);
		
			}else if(rayDrawer instanceof VRMLDrawer){
				((VRMLDrawer)rayDrawer).setRayColour(colorMap[colNum]);
				
			}
			
			bundleNum++;
			
			cost += rayBundle.traceAndEvaluate(rayDrawer);
		}
	
		if(drawingExtraVRML != null){ //erm, hack
			((VRMLDrawer)rayDrawer).addVRML("}");
		}
		
		rayDrawer.destroy();
		rayDrawer = null;

		return cost;
	}
	
	@Override
	public IDomain getDomain() { 
		int n = parameters.size();
		double min[] = new double[n];
		double max[] = new double[n];
			
		int i=0;
		for(Parameter parameter : parameters){
			min[i] = parameter.min();
			max[i] = parameter.max();
			i++;
		}
		return new RectangularDomain(min, max);
	}
	
	/** If set, an svg files are created containing the tracing for all evaluations */  
	public void setRegularDrawing(String fileNamePrefix, Class type, int nIterations, int nSkipRays){
		this.drawingFileNamePrefix = fileNamePrefix;
		this.drawingType = type;
		this.drawingIterationPeriod = nIterations;
		this.drawingSkipRays = nSkipRays;
	}
	
	public void setExtraVRML(String drawingExtraVRML){
		this.drawingExtraVRML = drawingExtraVRML;
	}
	
	
	public void singleParamScan(int pIdx, double p0, double p1, int nSteps) {
		double p[] = getParams();
		
		double v0 = eval(p);
		double dp = (p1 - p0) / (nSteps - 1);
		System.out.println("v0=\t" + v0);
		for(int i=0; i < nSteps; i++){
			p[pIdx] = p0 + i * dp;
			double v = eval(p);
			System.out.println(i + "\t" + p[pIdx] + "\t" + v);
		}
		
	}

	/** Sets fileName to write the optimisation progress to, the file format will be:
	 * { iteration No, evaluation No, ray bundle No, ray No, imgX, imgY, {params} }
	 * @param hitsOutputFile		Filename, or null to not write debugging info	 */
	public void setHitsOutputFile(String hitsOutputFile){ this.hitsOutputFile = hitsOutputFile;	}
	public BinaryMatrixWriter getHitsOutWriter() { return hitsWriteEnable ? hitsOutWriter : null; }
	
	public int getIterationNo() { return iterationNo; }
	public Element getAllElements() { return allElements; }

	public void addParameter(Parameter parameter){ parameters.add(parameter);	}
	public void addRayBundle(RayBundle rayBundle){ rayBundles.add(rayBundle);	}
	public void setTracingElements(Element allElements) { this.allElements = allElements;	}
	public void addRayBundles(RayBundle rayBundleArray[]){ 
		for(RayBundle rayBundle : rayBundleArray)
			rayBundles.add(rayBundle);
	}

	public void dumpParams() {
		for(Parameter parameter : parameters){
			System.out.println(parameter.toString());
		}
	}
	
	public void dumpRayBundles() {
		for(RayBundle rayBundle : rayBundles){
			System.out.println(rayBundle.toString());
		}
	}

	public Object getCurrentEvalParams() { return currentEvalParams;	}

	public void eval() { eval(getParams()); }

}
