package parody;

import org.ejml.data.DMatrixRMaj;
import org.ejml.dense.row.CommonOps_DDRM;
import org.ejml.dense.row.NormOps_DDRM;
import org.ejml.simple.SimpleMatrix;

import com.sun.org.apache.bcel.internal.generic.NEWARRAY;

import parody.objective.ObjectiveFunction;
import parody.optim.BFGSOptimizer;
import parody.optim.CommonsGradientOptimizer;
import parody.optim.DDogLegBFGSOptimizer;
import parody.optim.DDogLegGradientFreeOptimizer;
import parody.optim.NewtonsMethodOptimizer;
import parody.optim.Optimizer;
import parody.optim.SimpleHessianObjectiveFunction;
import parody.optim.SimpleObjectiveFunction;
import parody.optim.SteepestDescentOptimizer;
import parody.util.Constants;
import parody.util.MatrixUtils;
import parody.util.MultivariateNormal;
import parody.util.SymmetricMatrix;
import parody.util.SymmetricMatrixArray;

public class KernelObjectiveFunction extends SimpleObjectiveFunction {

	private ObjectiveFunction f;
	private MultivariateNormal distr;
	private KernelObjectiveFunctionOfData kernelFOfData;
	private Optimizer optim;
	private DMatrixRMaj hessian;
	private DMatrixRMaj hessianInv;
	
	public KernelObjectiveFunction(MultivariateNormal kernel, ObjectiveFunction f){
		super(f.paramDim());
		int dataDim = f.dataDim();
		int paramDim = f.paramDim();
		kernelFOfData = new KernelObjectiveFunctionOfData(f, kernel);
		setKernel(kernel);
		this.f = f;
		hessian = new DMatrixRMaj(dataDim, dataDim);
		hessianInv = new DMatrixRMaj(dataDim, dataDim);
		dataGradientParGradient = new DMatrixRMaj(dataDim, paramDim);
		parGradientDataHessian = new SymmetricMatrixArray(paramDim, dataDim);
		zhat = new double[dataDim];
		zhatParJacobian = new DMatrixRMaj(dataDim, paramDim);
		optim = new NewtonsMethodOptimizer();
		optim.init(kernelFOfData);
	}
	
	public void setKernel(MultivariateNormal kernel){
		distr = kernel;
		kernelFOfData.setKernel(kernel);
		setDirty();
		
	}
	
	private void setDirty(){
		zhatDirty = true;
		hessianCholDirty = true;
		hessianInvDirty = true;
		zhatParJacobianDirty = true;
		parGradientDataHessianDirty = true;
	}
	
	private boolean zhatDirty;
	private double[] zhat;
	
	private boolean zhatConverged;
	
	private void calcZhat(){
		if(zhatDirty){
			f.setParams(getCurrentState());
			System.arraycopy(distr.getMean(), 0, zhat, 0, zhat.length);
			boolean res = optim.optimize(zhat);
			if(!res){
				DMatrixRMaj grad = new DMatrixRMaj(f.dataDim(), 1);
				kernelFOfData.computeGradient(grad.data);
				double gradNorm = NormOps_DDRM.normF(grad);
				zhatConverged = gradNorm < 1e-5;
			} else {
				zhatConverged = true;
			}
			zhatDirty = false;
		}
	}
	
	public double[] getZhat(){
		calcZhat();
		return zhat;
	}
	
	private boolean hessianCholDirty;
	private DMatrixRMaj hessianChol;
	private void calcHessianChol(){
		calcZhat();
		if(hessianCholDirty){
			kernelFOfData.setInput(zhat);
			kernelFOfData.computeHessian(hessian);
			hessianChol = SymmetricMatrix.getChol(hessian);
			hessianCholDirty = false;		
		}
	}
	
	private boolean hessianInvDirty;
	private void calcHessianInv(){
		if(hessianInvDirty){
			calcHessianChol();
			SymmetricMatrix.invFromChol(hessianChol, hessianInv);
			hessianInvDirty = false;
		}
	}
	
	@Override
	public double computeFunctionAlways() {
		calcHessianChol();
		if(!zhatConverged){
			return Double.NaN;
		}
		double hessianDet = SymmetricMatrix.detFromChol(hessianChol);
		if(hessianDet <= 0){
			return Double.NaN;
		} else {
			return 0.5*Math.log(hessianDet)+kernelFOfData.computeFunction();
		}
		
	}
	
	public SymmetricMatrixArray parGradientDataHessian;
	public boolean parGradientDataHessianDirty;
	
	public void calcParGradientDataHessian(){
		if(parGradientDataHessianDirty){
			f.paramGradientOfDataHessian(parGradientDataHessian);
			parGradientDataHessianDirty = false;
		}
	}

	@Override
	public void computeGradientAlways(double[] out) {
		calcZhatParJacobian(); //Includes inverse Hessian
		if(!zhatConverged){
			out[0] = Double.NaN;
			return;
		}
		double hessianDet = SymmetricMatrix.detFromChol(hessianChol);
		if(hessianDet <= 0){
			out[0] = Double.NaN;
			return;
		}
		calcParGradientDataHessian();
		int dataDim = f.dataDim();
		int paramDim = f.paramDim();
		
		SymmetricMatrixArray dataThirdDerivatives = new SymmetricMatrixArray(dataDim, dataDim);
		f.dataThirdDerivative(dataThirdDerivatives);
		
		DMatrixRMaj logHessDetDataGradient = new DMatrixRMaj(dataDim, 1);
		DMatrixRMaj logHessDetParamGradient = new DMatrixRMaj(paramDim, 1);
		
		for(int iDat = 0; iDat < dataDim; iDat++){
			logHessDetDataGradient.set(iDat, 0, MatrixUtils.aTransProductTrace(hessianInv, dataThirdDerivatives.getMatrix(iDat)));
		}
		
		for(int iPar = 0; iPar < f.paramDim(); iPar++){
			logHessDetParamGradient.set(iPar, 0, MatrixUtils.aTransProductTrace(hessianInv, parGradientDataHessian.getMatrix(iPar)));
		}
		
		DMatrixRMaj densDataGradient = new DMatrixRMaj(dataDim, 1);
		DMatrixRMaj integrandDataGradient = new DMatrixRMaj(dataDim, 1);
		
		f.dataGradient(integrandDataGradient.data);
		distr.dataGradientLogLikelihood(densDataGradient.data);
		
		CommonOps_DDRM.subtract(integrandDataGradient, densDataGradient, integrandDataGradient);
		
		DMatrixRMaj integrandParamGradient = new DMatrixRMaj(paramDim, 1);
		f.paramGradient(integrandParamGradient.data);
		
		DMatrixRMaj dataGradient = new DMatrixRMaj(dataDim, 1);
		DMatrixRMaj paramGradient = new DMatrixRMaj(paramDim, 1);
		paramGradient.setData(out);
		
		CommonOps_DDRM.add(0.5, logHessDetDataGradient, integrandDataGradient, dataGradient);
		CommonOps_DDRM.add(0.5, logHessDetParamGradient, integrandParamGradient, paramGradient);

		CommonOps_DDRM.multAddTransA(zhatParJacobian, dataGradient, paramGradient);
	}
	
	private DMatrixRMaj dataGradientParGradient;
	private DMatrixRMaj zhatParJacobian;
	private boolean zhatParJacobianDirty;
	
	private void calcZhatParJacobian(){
		if(zhatParJacobianDirty){
			calcHessianInv();
			f.paramGradientDataGradient(dataGradientParGradient);
			CommonOps_DDRM.mult(-1.0, hessianInv,
					dataGradientParGradient, zhatParJacobian);
			zhatParJacobianDirty = false;
		}
	}
	
	public DMatrixRMaj zhatParJacobian(){
		calcZhatParJacobian();
		return zhatParJacobian;
	}

	@Override
	public double minValue() {
		return Double.NEGATIVE_INFINITY;
	}
	
	public void setInput(double x) {
		super.setInput(x);
		setDirty();
	}

	@Override
	public void setInput(double[] x) {
		super.setInput(x);
		setDirty();
	}

}
