/*
 *  Copyright 2007-2015 The OpenMx Project
 *
 *  Licensed under the Apache License, Version 2.0 (the "License");
 *  you may not use this file except in compliance with the License.
 *  You may obtain a copy of the License at
 *
 *       http://www.apache.org/licenses/LICENSE-2.0
 *
 *   Unless required by applicable law or agreed to in writing, software
 *   distributed under the License is distributed on an "AS IS" BASIS,
 *   WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 *  See the License for the specific language governing permissions and
 *  limitations under the License.
 *
 */

/***********************************************************
*
*  omxStateSpaceExpectation.c
*
*  Created: Michael D. Hunter 	Date: 2012-10-28 20:07:36
*
*  Contains code to calculate the objective function for a
*   state space model.  Currently, this is done with a
*   Kalman filter in separate Predict and Update steps.
*   Later this could be done with one of several Kalman
*   filter-smoothers (a forward-backward algorithm).
*
**********************************************************/

#include "omxExpectation.h"
#include "omxBLAS.h"
#include "omxFIMLFitFunction.h"
#include "omxStateSpaceExpectation.h"
#include <Eigen/Core>
#include <Eigen/Dense>
#include <Eigen/LU>
#include <unsupported/Eigen/MatrixFunctions>
#include <iostream>

void omxCallStateSpaceExpectation(omxExpectation* ox, const char *, const char *) {
    if(OMX_DEBUG) { mxLog("State Space Expectation Called."); }
	omxStateSpaceExpectation* ose = (omxStateSpaceExpectation*)(ox->argStruct);

	omxRecompute(ose->A, NULL);
	omxRecompute(ose->B, NULL);
	omxRecompute(ose->C, NULL);
	omxRecompute(ose->D, NULL);
	omxRecompute(ose->Q, NULL);
	omxRecompute(ose->R, NULL);

	// Probably should loop through all the data here!!!
	if(ose->t == NULL){
		omxKalmanPredict(ose);
	} else {
		omxKalmanBucyPredict(ose);
	}
	omxKalmanUpdate(ose);
}

void omxDestroyStateSpaceExpectation(omxExpectation* ox) {

	if(OMX_DEBUG) { mxLog("Destroying State Space Expectation."); }

	omxStateSpaceExpectation* argStruct = (omxStateSpaceExpectation*)(ox->argStruct);

	omxFreeMatrix(argStruct->r);
	omxFreeMatrix(argStruct->s);
	omxFreeMatrix(argStruct->z);
	omxFreeMatrix(argStruct->x);
	omxFreeMatrix(argStruct->y);
	omxFreeMatrix(argStruct->K);
	omxFreeMatrix(argStruct->P);
	omxFreeMatrix(argStruct->S);
	omxFreeMatrix(argStruct->Y);
	omxFreeMatrix(argStruct->Z);
	omxFreeMatrix(argStruct->det);
	omxFreeMatrix(argStruct->covInfo);
	omxFreeMatrix(argStruct->cov);
	omxFreeMatrix(argStruct->means);
	omxFreeMatrix(argStruct->smallC);
	omxFreeMatrix(argStruct->smallD);
	omxFreeMatrix(argStruct->smallR);
	omxFreeMatrix(argStruct->smallr);
	omxFreeMatrix(argStruct->smallK);
	omxFreeMatrix(argStruct->smallS);
	omxFreeMatrix(argStruct->smallY);

	delete argStruct;

}

void omxPopulateSSMAttributes(omxExpectation *ox, SEXP algebra) {
	if(OMX_DEBUG) { mxLog("Populating State Space Attributes.  Currently this does very little!"); }

	/* Initialize */
	omxSetExpectationComponent(ox, NULL, "Reset", NULL); //maybe shoulde be on ose?  after next line?
	omxStateSpaceExpectation* ose = (omxStateSpaceExpectation*)(ox->argStruct);

	if( !(ose->returnScores) ){
		if(OMX_DEBUG) { mxLog("Not asking for attributes, this is being skipped!"); }
		return;
	}

	SEXP xpred, ypred, ppred, spred, xupda, pupda, xsmoo, psmoo;
	// yupda, supda

	omxRecompute(ose->A, NULL);
	omxRecompute(ose->B, NULL);
	omxRecompute(ose->C, NULL);
	omxRecompute(ose->D, NULL);
	omxRecompute(ose->Q, NULL);
	omxRecompute(ose->R, NULL);

	// allocate matrices to be returned
	if(OMX_DEBUG_ALGEBRA) { mxLog("Allocating initial population matrices ..."); }
	int nx = ose->C->cols;
	int ny = ose->C->rows;
	if(OMX_DEBUG_ALGEBRA) { mxLog("Find number of rows of data ..."); }
	//int nt = ox->data->dataMat->rows;
	int nt = ox->data->numObs;
	if(OMX_DEBUG_ALGEBRA) {std::cout << "... numObs:\n" << nt << std::endl; }
	if(OMX_DEBUG_ALGEBRA) { mxLog("Done Finding rows of data ..."); }
	int np = ((nx+1)*nx)/2;
	int ns = ((ny+1)*ny)/2;
	Rf_protect(xpred = Rf_allocMatrix(REALSXP, nt+1, nx));
	Rf_protect(ypred = Rf_allocMatrix(REALSXP, nt+1, ny));
	Rf_protect(ppred = Rf_allocMatrix(REALSXP, nt+1, np));
	Rf_protect(spred = Rf_allocMatrix(REALSXP, nt+1, ns));
	Rf_protect(xupda = Rf_allocMatrix(REALSXP, nt+1, nx));
	//Rf_protect(yupda = Rf_allocMatrix(REALSXP, nt+1, ny));
	Rf_protect(pupda = Rf_allocMatrix(REALSXP, nt+1, np));
	//Rf_protect(supda = Rf_allocMatrix(REALSXP, nt+1, ns));
	Rf_protect(xsmoo = Rf_allocMatrix(REALSXP, nt+1, nx));
	Rf_protect(psmoo = Rf_allocMatrix(REALSXP, nt+1, np));

	if(OMX_DEBUG_ALGEBRA) { mxLog("Setting zeroth row ..."); }
	// Set first row of xpred to x0
	int row = 0;
	for(int col = 0; col < nx; col++){
		REAL(xpred)[col * (nt+1) + row] =
			omxMatrixElement(ose->x, col, 0);
		REAL(xupda)[col * (nt+1) + row] =
			omxMatrixElement(ose->x, col, 0);
	}

	// Set first row of ppred to vech(P0)
	int counter = 0;
	for(int i = 0; i < ose->P->cols; i++) {
		for(int j = i; j < ose->P->rows; j++) {
			REAL(ppred)[counter * (nt+1) + row] = omxMatrixElement(ose->P, j, i);
			REAL(pupda)[counter * (nt+1) + row] = omxMatrixElement(ose->P, j, i);
			counter++;
		}
	}

	Eigen::VectorXd oldDefs;
	oldDefs.resize(ox->data->defVars.size());
	oldDefs.setConstant(NA_REAL);

	// Probably should loop through all the data here!!!
	if(OMX_DEBUG_ALGEBRA) { mxLog("Beginning forward loop ..."); }
	for(row=1; row < (nt+1); row++){
		if(OMX_DEBUG_ALGEBRA) { mxLog("Setting first data row ..."); }
		// Set row of data
		for(int i = 0; i < ny; i++) {
			omxSetMatrixElement(ose->y, i, 0, omxDoubleDataElement(ox->data, row-1, i));
		}

		// handle definition variables
		int numVarsFilled = 0;
		numVarsFilled = ox->data->handleDefinitionVarList(ox->currentState, row-1, oldDefs.data());

		/* Run Kalman prediction */
		if(ose->t == NULL){
			omxKalmanPredict(ose);
		} else {
			omxKalmanBucyPredict(ose);
		}

		// Copy latent state
		for(int col = 0; col < nx; col++)
			REAL(xpred)[col * (nt+1) + row] =
				omxMatrixElement(ose->x, col, 0);

		// Copy latent cov
		counter = 0;
		for(int i = 0; i < ose->P->cols; i++) {
			for(int j = i; j < ose->P->rows; j++) {
				REAL(ppred)[counter * (nt+1) + row] = omxMatrixElement(ose->P, j, i);
				counter++;
			}
		}

		// Create Full observed cov prediction
		if(OMX_DEBUG_ALGEBRA) { mxLog("Hand prediction of full observed cov ..."); }
		omxDSYMM(FALSE, 1.0, ose->P, ose->C, 0.0, ose->Y); // Y = C P
		omxCopyMatrix(ose->S, ose->R); // S = R
		omxDGEMM(FALSE, TRUE, 1.0, ose->Y, ose->C, 1.0, ose->S); // S = Y C^T + S THAT IS C P C^T + R

		// Copy observed cov
		counter = 0;
		for(int i = 0; i < ose->S->cols; i++) {
			for(int j = i; j < ose->S->rows; j++) {
				REAL(spred)[counter * (nt+1) + row] = omxMatrixElement(ose->S, j, i);
				counter++;
			}
		}

		/* Run Kalman update */
		omxKalmanUpdate(ose);

		// Copy latent state
		for(int col = 0; col < nx; col++)
			REAL(xupda)[col * (nt+1) + row] =
				omxMatrixElement(ose->x, col, 0);

		// Copy latent cov
		counter = 0;
		for(int i = 0; i < ose->P->cols; i++) {
			for(int j = i; j < ose->P->rows; j++) {
				REAL(pupda)[counter * (nt+1) + row] = omxMatrixElement(ose->P, j, i);
				counter++;
			}
		}

		// Copy observed means prediction
		for(int col = 0; col < ny; col++)
			REAL(ypred)[col * (nt+1) + row] =
				omxMatrixElement(ose->s, col, 0);

		// TODO Add m2ll calculation here.
		// Probably like this
		/*m2ll = y^T S y */ // n.b. y originally is the data row but becomes the data residual!
		//omxDSYMV(1.0, S, y, 0.0, s); // s = S y
		//m2ll = omxDDOT(y, s); // m2ll = y s THAT IS y^T S y
		//m2ll += det; // m2ll = m2ll + det THAT IS m2ll = log(det(S)) + y^T S y
		// Note: this leaves off the S->cols * log(2*pi) THAT IS k*log(2*pi)
	}

	/* TODO Add Backward pass through data for Kalman smoother*/
	// Initialize end of smoothed latents to last updated latents
	// P = last updated P
	// x = last updated x

	// Copy x and P as smoothed estimates for export
	row = nt;
	// Copy latent state
	for(int col = 0; col < nx; col++)
		REAL(xsmoo)[col * (nt+1) + row] =
			omxMatrixElement(ose->x, col, 0);

	// Copy latent cov
	counter = 0;
	for(int i = 0; i < ose->P->cols; i++) {
		for(int j = i; j < ose->P->rows; j++) {
			REAL(psmoo)[counter * (nt+1) + row] = omxMatrixElement(ose->P, j, i);
			counter++;
		}
	}

	// loop backwars through all data
	if(OMX_DEBUG_ALGEBRA) { mxLog("Beginning backward loop ..."); }
	for(row = nt-1; row > -1; row--){
		// handle definition variables
		int numVarsFilled = 0;
		numVarsFilled = ox->data->handleDefinitionVarList(ox->currentState, row, oldDefs.data());

		// Copy Z = updated P from pupda
		counter = 0;
		for(int i = 0; i < nx; i++) {
			for(int j = i; j < nx; j++) {
				double next = REAL(pupda)[counter * (nt+1) + row];
				omxSetMatrixElement(ose->Z, i, j, next);
				omxSetMatrixElement(ose->Z, j, i, next);
				counter++;
			}
		}

		// Copy z = updated x from xupda
		for(int col = 0; col < nx; col++)
			omxSetMatrixElement(ose->z, col, 0, REAL(xupda)[col * (nt+1) + row]);

		// Copy eigenExpA = predicted P from later row of ppred
		counter = 0;
		for(int i = 0; i < nx; i++) {
			for(int j = i; j < nx; j++) {
				double next = REAL(ppred)[counter * (nt+1) + row + 1];
				ose->eigenExpA(i, j) = next;
				ose->eigenExpA(j, i) = next;
				counter++;
			}
		}

		// Copy eigenPreX = predicted x from later row of xpred
		for(int col = 0; col < nx; col++)
			ose->eigenPreX(col, 0) =  REAL(xpred)[col * (nt+1) + row + 1];

		// Run Smoother
		omxRauchTungStriebelSmooth(ose);

		// Copy latent state to xsmoo
		for(int col = 0; col < nx; col++)
			REAL(xsmoo)[col * (nt+1) + row] =
				omxMatrixElement(ose->x, col, 0);

		// Copy latent cov to psmoo
		counter = 0;
		for(int i = 0; i < ose->P->cols; i++) {
			for(int j = i; j < ose->P->rows; j++) {
				REAL(psmoo)[counter * (nt+1) + row] = omxMatrixElement(ose->P, j, i);
				counter++;
			}
		}
	}

	// TODO check on definition variable population
	//  I suspect this does yet work properly for def vars.

	Rf_setAttrib(algebra, Rf_install("xPredicted"), xpred);
	Rf_setAttrib(algebra, Rf_install("yPredicted"), ypred);
	Rf_setAttrib(algebra, Rf_install("PPredicted"), ppred);
	Rf_setAttrib(algebra, Rf_install("SPredicted"), spred);
	Rf_setAttrib(algebra, Rf_install("xUpdated"), xupda);
	Rf_setAttrib(algebra, Rf_install("PUpdated"), pupda);
	Rf_setAttrib(algebra, Rf_install("xSmoothed"), xsmoo);
	Rf_setAttrib(algebra, Rf_install("PSmoothed"), psmoo);

	/*
	omxMatrix *expCovInt, *expMeanInt;
	expCovInt = argStruct->cov;
	expMeanInt = argStruct->means;

	Rf_protect(expCovExt = Rf_allocMatrix(REALSXP, expCovInt->rows, expCovInt->cols));
	for(int row = 0; row < expCovInt->rows; row++)
		for(int col = 0; col < expCovInt->cols; col++)
			REAL(expCovExt)[col * expCovInt->rows + row] =
				omxMatrixElement(expCovInt, row, col);
	if (expMeanInt != NULL && expMeanInt->rows > 0  && expMeanInt->cols > 0) {
		Rf_protect(expMeanExt = Rf_allocMatrix(REALSXP, expMeanInt->rows, expMeanInt->cols));
		for(int row = 0; row < expMeanInt->rows; row++)
			for(int col = 0; col < expMeanInt->cols; col++)
				REAL(expMeanExt)[col * expMeanInt->rows + row] =
					omxMatrixElement(expMeanInt, row, col);
	} else {
		Rf_protect(expMeanExt = Rf_allocMatrix(REALSXP, 0, 0));
	}

	Rf_setAttrib(algebra, Rf_install("expCov"), expCovExt);
	Rf_setAttrib(algebra, Rf_install("expMean"), expMeanExt);

	if(argStruct->populateRowDiagnostics){
		omxMatrix *rowLikelihoodsInt = argStruct->rowLikelihoods;
		Rf_protect(rowLikelihoodsExt = Rf_allocVector(REALSXP, rowLikelihoodsInt->rows));
		for(int row = 0; row < rowLikelihoodsInt->rows; row++)
			REAL(rowLikelihoodsExt)[row] = omxMatrixElement(rowLikelihoodsInt, row, 0);
		Rf_setAttrib(algebra, Rf_install("likelihoods"), rowLikelihoodsExt);
	}
	*/

}

void omxKalmanPredict(omxStateSpaceExpectation* ose) {
	if(OMX_DEBUG) { mxLog("Kalman Predict Called."); }
	/* Creat local copies of State Space Matrices */
	omxMatrix* A = ose->A;
	if(OMX_DEBUG_ALGEBRA) {omxPrintMatrix(A, "....State Space: A"); }
	omxMatrix* B = ose->B;
	if(OMX_DEBUG_ALGEBRA) {omxPrintMatrix(B, "....State Space: B"); }
	omxMatrix* Q = ose->Q;
	if(OMX_DEBUG_ALGEBRA) {omxPrintMatrix(Q, "....State Space: Q"); }
	omxMatrix* u = ose->u;
	if(OMX_DEBUG_ALGEBRA) {omxPrintMatrix(u, "....State Space: u"); }
	omxMatrix* x = ose->x;
	if(OMX_DEBUG_ALGEBRA) {omxPrintMatrix(x, "....State Space: x"); }
	omxMatrix* z = ose->z;
	omxMatrix* P = ose->P;
	if(OMX_DEBUG_ALGEBRA) {omxPrintMatrix(P, "....State Space: P"); }
	//omxMatrix* S = ose->S;
	//omxMatrix* Y = ose->Y;
	omxMatrix* Z = ose->Z;

	/* x = A x + B u */
	omxDGEMV(FALSE, 1.0, A, x, 0.0, z); // x = A x
	if(OMX_DEBUG_ALGEBRA) {omxPrintMatrix(z, "....State Space: z = A x"); }
	if(OMX_DEBUG_ALGEBRA) {omxPrintMatrix(A, "....State Space: A"); }
	omxDGEMV(FALSE, 1.0, B, u, 1.0, z); // x = B u + x THAT IS x = A x + B u
	if(OMX_DEBUG_ALGEBRA) {omxPrintMatrix(z, "....State Space: z = A x + B u"); }
	omxCopyMatrix(x, z); // x = z THAT IS x = A x + B u
	if(OMX_DEBUG_ALGEBRA) {omxPrintMatrix(x, "....State Space: x = A x + B u"); }

	/* P = A P A^T + Q */
	omxDSYMM(FALSE, 1.0, P, A, 0.0, Z); // Z = A P
	if(OMX_DEBUG_ALGEBRA) {omxPrintMatrix(Z, "....State Space: Z = A P"); }
	omxCopyMatrix(P, Q); // P = Q
	if(OMX_DEBUG_ALGEBRA) {omxPrintMatrix(P, "....State Space: P = Q"); }
	omxDGEMM(FALSE, TRUE, 1.0, Z, A, 1.0, P); // P = Z A^T + P THAT IS P = A P A^T + Q
	if(OMX_DEBUG_ALGEBRA) {omxPrintMatrix(P, "....State Space: P = A P A^T + Q"); }
}

void omxKalmanUpdate(omxStateSpaceExpectation* ose) {
	//TODO: Clean up this hack of function.
	if(OMX_DEBUG) { mxLog("Kalman Update Called."); }
	/* Creat local copies of State Space Matrices */
	omxMatrix* C = ose->C;
	if(OMX_DEBUG_ALGEBRA) {omxPrintMatrix(C, "....State Space: C"); }
	omxMatrix* D = ose->D;
	if(OMX_DEBUG_ALGEBRA) {omxPrintMatrix(D, "....State Space: D"); }
	omxMatrix* r = ose->r;
	omxMatrix* s = ose->s;
	omxMatrix* u = ose->u;
	if(OMX_DEBUG_ALGEBRA) {omxPrintMatrix(u, "....State Space: u"); }
	omxMatrix* x = ose->x;
	if(OMX_DEBUG_ALGEBRA) {omxPrintMatrix(x, "....State Space: x"); }
	omxMatrix* y = ose->y;
	if(OMX_DEBUG_ALGEBRA) {omxPrintMatrix(y, "....State Space: y"); }
	omxMatrix* P = ose->P;
	// omxMatrix* Cov = ose->cov; //unused
	omxMatrix* Means = ose->means;
	omxMatrix* Det = ose->det;
	*Det->data = 0.0; // the value pointed to by Det->data is assigned to be zero
	omxMatrix* smallC = ose->smallC;
	omxMatrix* smallD = ose->smallD;
	omxMatrix* smallr = ose->smallr;
	omxMatrix* R = ose->R;
	if(OMX_DEBUG_ALGEBRA) {omxPrintMatrix(R, "....State Space: R on entry"); }
	omxMatrix* smallR = ose->smallR;
	if(OMX_DEBUG_ALGEBRA) {omxPrintMatrix(smallR, "....State Space: small R on entry"); }
	omxMatrix* smallK = ose->smallK;
	omxMatrix* smallS = ose->smallS;
	omxMatrix* smallY = ose->smallY;
	int ny = y->rows;
	int nx = x->rows;
	Eigen::VectorXi toRemoveSS(ny);
	int numRemovesSS = 0;
	Eigen::VectorXi toRemoveNoneLat(nx);
	toRemoveNoneLat.setZero();
	Eigen::VectorXi toRemoveNoneOne(1);
	toRemoveNoneOne.setZero();

	omxMatrix* covInfo = ose->covInfo;
	int info = 0; // Used for computing inverse for Kalman gain

	omxCopyMatrix(smallS, ose->S);

	/* Reset/Resample aliased matrices */
	omxCopyMatrix(smallC, ose->C);
	if(OMX_DEBUG_ALGEBRA) {omxPrintMatrix(smallC, "....State Space: C (Reset)"); }
	omxCopyMatrix(smallD, ose->D);
	omxCopyMatrix(smallR, ose->R);
	if(OMX_DEBUG_ALGEBRA) {omxPrintMatrix(smallR, "....State Space: small R (Reset)"); }
	omxCopyMatrix(smallr, ose->r);
	if(OMX_DEBUG_ALGEBRA) {omxPrintMatrix(smallr, "....State Space: r (Reset)"); }
	omxCopyMatrix(smallK, ose->K);
	omxCopyMatrix(smallS, ose->S);
	omxCopyMatrix(smallY, ose->Y);

	/* r = r - C x - D u */
	/* Alternatively, create just the expected value for the data row, x. */
	omxDGEMV(FALSE, 1.0, smallC, x, 0.0, s); // s = C x
	if(OMX_DEBUG_ALGEBRA) {omxPrintMatrix(s, "....State Space: s = C x"); }
	omxDGEMV(FALSE, 1.0, D, u, 1.0, s); // s = D u + s THAT IS s = C x + D u
	if(OMX_DEBUG_ALGEBRA) {omxPrintMatrix(s, "....State Space: s = C x + D u"); }
	omxCopyMatrix(Means, s); // Means = s THAT IS Means = C x + D u
	//if(OMX_DEBUG_ALGEBRA) {omxPrintMatrix(Means, "....State Space: Means"); }
	omxTransposeMatrix(Means);
	if(OMX_DEBUG_ALGEBRA) {omxPrintMatrix(Means, "....State Space: Means"); }

	//If entire data vector, y, is missing, then set residual, r, to zero.
	//otherwise, compute residual.
	toRemoveSS.setZero();
	for(int j = 0; j < y->rows; j++) {
		double dataValue = omxMatrixElement(y, j, 0);
		int dataValuefpclass = std::fpclassify(dataValue);
		if(dataValuefpclass == FP_NAN || dataValuefpclass == FP_INFINITE) {
			numRemovesSS++;
			toRemoveSS[j] = 1;
			omxSetMatrixElement(r, j, 0, 0.0);
		} else {
			omxSetMatrixElement(r, j, 0, (dataValue -  omxMatrixElement(s, j, 0)));
		}
	}
	if(OMX_DEBUG_ALGEBRA) {omxPrintMatrix(r, "....State Space: Residual (Loop)"); }
	/* Now compute the residual */
	//omxCopyMatrix(r, y); // r = y
	//omxDAXPY(-1.0, s, r); // r = r - s THAT IS r = y - (C x + D u)
	omxCopyMatrix(smallr, ose->r);

	/* Filter S Here */
	// N.B. if y is completely missing or completely present, leave S alone.
	// Otherwise, filter S.
	if(numRemovesSS < ny && numRemovesSS > 0) {
		//if(OMX_DEBUG_ALGEBRA) {omxPrintMatrix(smallS, "....State Space: S"); }
		if(OMX_DEBUG) { mxLog("Filtering S, R, C, r, K, and Y."); }
		omxRemoveRowsAndColumns(smallS, numRemovesSS, numRemovesSS, toRemoveSS.data(), toRemoveSS.data());
		//if(OMX_DEBUG_ALGEBRA) {omxPrintMatrix(smallS, "....State Space: S (Filtered)"); }

		omxRemoveRowsAndColumns(smallR, numRemovesSS, numRemovesSS, toRemoveSS.data(), toRemoveSS.data());

		if(OMX_DEBUG_ALGEBRA) {omxPrintMatrix(smallC, "....State Space: C"); }
		omxRemoveRowsAndColumns(smallC, numRemovesSS, 0, toRemoveSS.data(), toRemoveNoneLat.data());
		if(OMX_DEBUG_ALGEBRA) {omxPrintMatrix(smallC, "....State Space: C (Filtered)"); }

		if(OMX_DEBUG_ALGEBRA) {omxPrintMatrix(r, "....State Space: r"); }
		omxRemoveRowsAndColumns(smallr, numRemovesSS, 0, toRemoveSS.data(), toRemoveNoneOne.data());
		if(OMX_DEBUG_ALGEBRA) {omxPrintMatrix(smallr, "....State Space: r (Filtered)"); }

		//if(OMX_DEBUG_ALGEBRA) {omxPrintMatrix(smallK, "....State Space: K"); }
		omxRemoveRowsAndColumns(smallK, numRemovesSS, 0, toRemoveSS.data(), toRemoveNoneLat.data());
		//if(OMX_DEBUG_ALGEBRA) {omxPrintMatrix(smallK, "....State Space: K (Filtered)"); }

		omxRemoveRowsAndColumns(smallY, numRemovesSS, 0, toRemoveSS.data(), toRemoveNoneLat.data());
	}
	if(numRemovesSS == ny) {
		if(OMX_DEBUG_ALGEBRA) { mxLog("Completely missing row of data found."); }
		if(OMX_DEBUG_ALGEBRA) { mxLog("Skipping much of Kalman Update."); }
		return ;
	}

	/* S = C P C^T + R */
	omxDSYMM(FALSE, 1.0, P, smallC, 0.0, smallY); // Y = C P
	//omxCopyMatrix(S, smallR); // S = R
	memcpy(smallS->data, smallR->data, smallR->rows * smallR->cols * sizeof(double)); // Less safe omxCopyMatrix that keeps smallS aliased to S.
	omxDGEMM(FALSE, TRUE, 1.0, smallY, smallC, 1.0, smallS); // S = Y C^T + S THAT IS C P C^T + R
	if(OMX_DEBUG_ALGEBRA) {omxPrintMatrix(smallS, "....State Space: S = C P C^T + R"); }

	/* Now compute the Kalman Gain and update the error covariance matrix */
	/* S = S^-1 */
	omxDPOTRF(smallS, &info); // S replaced by the lower triangular matrix of the Cholesky factorization
	if(OMX_DEBUG_ALGEBRA) {omxPrintMatrix(smallS, "....State Space: Cholesky of S"); }
	covInfo->data[0] = (double) info;
	for(int i = 0; i < smallS->cols; i++) {
		*Det->data += log(fabs(omxMatrixElement(smallS, i, i)));
	}
	//det *= 2.0; //sum( log( abs( diag( chol(S) ) ) ) )*2
	omxDPOTRI(smallS, &info); // S = S^-1 via Cholesky factorization
	if(OMX_DEBUG_ALGEBRA) {omxPrintMatrix(smallS, "....State Space: Inverse of S"); }
	// If Cholesky of exp cov failed (i.e. non-positive def), Populate 1,1 element of smallS (inverse of exp cov) with NA_REAL
	if(covInfo->data[0] > 0) {
		omxSetMatrixElement(smallS, 0, 0, NA_REAL);
	}

	/* K = P C^T S^-1 */
	/* Computed as K^T = S^-1 C P */
	omxDSYMM(TRUE, 1.0, smallS, smallY, 0.0, smallK); // K = Y^T S THAT IS K = P C^T S^-1
	if(OMX_DEBUG_ALGEBRA) {omxPrintMatrix(smallK, "....State Space: K^T = S^-1 C P"); }

	/* x = x + K r */
	if(OMX_DEBUG_ALGEBRA) {omxPrintMatrix(smallr, "....State Space Check Residual: r"); }
	omxDGEMV(TRUE, 1.0, smallK, smallr, 1.0, x); // x = K r + x
	if(OMX_DEBUG_ALGEBRA) {omxPrintMatrix(x, "....State Space: x = K r + x"); }

	/* P = (I - K C) P */
	/* P = P - K C P */
	omxDGEMM(TRUE, FALSE, -1.0, smallK, smallY, 1.0, P); // P = -K Y + P THAT IS P = P - K C P
	if(OMX_DEBUG_ALGEBRA) {omxPrintMatrix(P, "....State Space: P = P - K C P"); }

	if(OMX_DEBUG_ALGEBRA) {omxPrintMatrix(smallS, "....State Space: Inverse of S"); }

	/*m2ll = y^T S y */ // n.b. y originally is the data row but becomes the data residual!
	//omxDSYMV(1.0, S, y, 0.0, s); // s = S y
	//m2ll = omxDDOT(y, s); // m2ll = y s THAT IS y^T S y
	//m2ll += det; // m2ll = m2ll + det THAT IS m2ll = log(det(S)) + y^T S y
	// Note: this leaves off the S->cols * log(2*pi) THAT IS k*log(2*pi)
}

void omxKalmanBucyPredict(omxStateSpaceExpectation* ose) {
	if(OMX_DEBUG) { mxLog("Kalman Bucy Predict Called."); }
	/* Creat local copies of State Space Matrices */
	omxMatrix* A = ose->A;
	if(OMX_DEBUG_ALGEBRA) {omxPrintMatrix(A, "....State Space: A"); }
	omxMatrix* B = ose->B;
	if(OMX_DEBUG_ALGEBRA) {omxPrintMatrix(B, "....State Space: B"); }
	omxMatrix* Q = ose->Q;
	if(OMX_DEBUG_ALGEBRA) {omxPrintMatrix(Q, "....State Space: Q"); }
	omxMatrix* u = ose->u;
	if(OMX_DEBUG_ALGEBRA) {omxPrintMatrix(u, "....State Space: u"); }
	omxMatrix* x = ose->x;
	if(OMX_DEBUG_ALGEBRA) {omxPrintMatrix(x, "....State Space: x"); }
	omxMatrix* P = ose->P;
	if(OMX_DEBUG_ALGEBRA) {omxPrintMatrix(P, "....State Space: P"); }
	omxMatrix* Z = ose->Z;
	omxMatrix* t = ose->t;
	if(OMX_DEBUG_ALGEBRA) {omxPrintMatrix(t, "....State Space: t"); }
	if(OMX_DEBUG_ALGEBRA) {std::cout << "... State Space oldT on entrance:\n" << ose->oldT << std::endl; }
	ose->deltaT = omxMatrixElement(t, 0, 0) - ose->oldT;
	ose->oldT = omxMatrixElement(t, 0, 0);
	if(OMX_DEBUG_ALGEBRA) {std::cout << "... State Space oldT on exit:\n" << ose->oldT << std::endl; }
	double deltaT = ose->deltaT;
	if(OMX_DEBUG_ALGEBRA) {std::cout << "... State Space deltaT:\n" << deltaT << std::endl; }

	//EigenMatrixAdaptor eigenA(A);
	// intializes eigenA as an instance of EigenMatrixAdaptor class, initialized to omxMatrix A
	// EigenVectorAdaptor
	// eigenB = eigenA.exp(); // matrix exponential
	// for scalar multiplication
	// A * omxElement(B,1,4); A eigen matrix, B omxMatrix
	// Subtract from diagonal 1.0
	// A.diagonal() -= 1

	/* Eigen Matrix reference setting */
	Eigen::MatrixXd &eigenExpA = ose->eigenExpA;
	Eigen::MatrixXd &eigenIA = ose->eigenIA;
	Eigen::MatrixXd &PSI = ose->PSI;
	Eigen::MatrixXd &IP = ose->IP;
	Eigen::MatrixXd &I = ose->I;

	/*R code for the next few lines
		ldim <- nrow(x$A)
		I <- diag(1, nrow=ldim)
		expA <- as.matrix(expm(x$A * x$deltaT))
		intA <- solve(x$A) %*% (expA - I)
		x$x <- expA %*% x$x + intA %*% x$B %*% x$u
	*/

	/*  Z = A
		eigenA = Z
		eigenExpA = eigenA*deltaT  THAT IS eigenExpA = A*deltaT
	*/
	omxCopyMatrix(Z, A);
	EigenMatrixAdaptor eigenA(Z);
	if(OMX_DEBUG_ALGEBRA) {std::cout << "... State Space eigenA:\n" << eigenA << std::endl; }
	eigenExpA = eigenA * deltaT;
	if(OMX_DEBUG_ALGEBRA) {std::cout << "... State Space eigenA:\n" << eigenA << std::endl; }
	if(OMX_DEBUG_ALGEBRA) {std::cout << "... State Space deltaT:\n" << deltaT << std::endl; }
	if(OMX_DEBUG_ALGEBRA) {std::cout << "... State Space eigenExpA:\n" << eigenExpA << std::endl; }

	/* eigenExpA = expm(eigenExpA)  THAT IS eigenExpA = expm(A*deltaT) */
	eigenExpA = eigenExpA.exp();
	if(OMX_DEBUG_ALGEBRA) {std::cout << "... State Space eigenExpA:\n" << eigenExpA << std::endl; }

	/* eigenIA = eigenExpA - I*/
	if(OMX_DEBUG_ALGEBRA) {std::cout << "... State Space I:\n" << I << std::endl; }
	eigenIA = eigenExpA - I;
	if(OMX_DEBUG_ALGEBRA) {std::cout << "... State Space expA - I:\n" << eigenIA << std::endl; }
	eigenIA = eigenA.lu().solve(eigenIA);
	if(OMX_DEBUG_ALGEBRA) {std::cout << "... State Space A^-1 (expA - I):\n" << eigenIA << std::endl; }
	if(OMX_DEBUG_ALGEBRA) {std::cout << "... State Space eigenA:\n" << eigenA << std::endl; }
	/* eigenIA = A^-1 IA  THAT IS eigenIA = A^-1 (expm(A*deltaT) - I) */

	/* x = expm(A*deltaT) * x + IA * B * u */
	EigenMatrixAdaptor eigenx(x); //or vector
	EigenMatrixAdaptor eigenu(u); //or vector
	EigenMatrixAdaptor eigenB(B);
	eigenx.derived() = eigenExpA * eigenx + eigenIA * eigenB * eigenu;
	if(OMX_DEBUG_ALGEBRA) {std::cout << "... State Space XPred:\n" << eigenx << std::endl; }

	/* SUMMARY */
	// EA = expm(A*deltaT)
	// IA = A^-1 (EA - I)
	// x = EA x + IA B u

	/*R code for the next few lines
		ldim <- nrow(x$A)
		psi <- cbind(rbind(x$A, matrix(0, nrow=ldim, ncol=ldim)), rbind(x$Q, -x$A))
		epsi <- as.matrix(expm(psi * x$deltaT)) %*% rbind(x$P, I)
		x$P <- epsi[1:ldim, ] %*% solve(epsi[(ldim+1):(2*ldim), ])
	*/

	/* PSI = block matrix
	-A^T  0
	 Q    A
	*/
	EigenMatrixAdaptor eigenQ(Q);
	PSI << -1.0*eigenA.transpose(), Eigen::MatrixXd::Zero(A->rows, A->rows), eigenQ, eigenA;
	PSI = PSI * deltaT;
	if(OMX_DEBUG_ALGEBRA) {std::cout << "... State Space PSI:\n" << PSI << std::endl; }
	/* PSI = PSI * deltaT */
	/* PSI = expm(PSI) */
	PSI = PSI.exp();
	if(OMX_DEBUG_ALGEBRA) {std::cout << "... State Space expPSI*deltaT:\n" << PSI << std::endl; }

	/* IP = block matrix
	I
	P
	*/
	EigenMatrixAdaptor eigenP(P);
	IP << I, eigenP;
	//IP << eigenP, Eigen::MatrixXd::Identity(rows, rows)
	if(OMX_DEBUG_ALGEBRA) {std::cout << "... State Space IP:\n" << IP << std::endl; }

	/* IP = PSI IP */
	IP = PSI * IP;
	if(OMX_DEBUG_ALGEBRA) {std::cout << "... State Space Blocks:\n" << IP << std::endl; }

	/* eigenIA = block 1 of IP; eigenExpA = block 2 of IP */
	eigenP.derived() = IP.block(0, 0, A->rows, A->cols);
	if(OMX_DEBUG_ALGEBRA) {std::cout << "... State Space Block1:\n" << eigenP << std::endl; }
	eigenExpA = IP.block(A->rows, 0, A->rows, A->cols);
	if(OMX_DEBUG_ALGEBRA) {std::cout << "... State Space Block2:\n" << eigenExpA << std::endl; }
	eigenP.derived() = eigenP.transpose().lu().solve(eigenExpA.transpose());
	if(OMX_DEBUG_ALGEBRA) {std::cout << "... State Space eigenP:\n" << eigenP << std::endl; }
	/* P = B1 B2^-1  THAT IS  solve B2^T P = B1^T for P */

	/* SUMMARY */
	// PSI = 	-A^T  0
	//			 Q    A
	// PSI = expm(PSI * deltaT)
	// IP = 	I
	// 			P
	// IP = PSI IP
	// P = IP[1] IP[2]^-1

	if(OMX_DEBUG_ALGEBRA) {std::cout << "... State Space eigenX" << eigenx << std::endl; }
	if(OMX_DEBUG_ALGEBRA) {omxPrintMatrix(x, "....State Space: x"); }
	if(OMX_DEBUG_ALGEBRA) {std::cout << "... State Space eigenP:\n" << eigenP << std::endl; }
	if(OMX_DEBUG_ALGEBRA) {omxPrintMatrix(P, "....State Space: P"); }
}

void omxRauchTungStriebelSmooth(omxStateSpaceExpectation* ose) {
	if(OMX_DEBUG) { mxLog("Rauch Tung Striebel Smooth Called."); }
	/* Creat local copies of State Space Matrices */
	omxMatrix* A = ose->A;
	if(OMX_DEBUG_ALGEBRA) {omxPrintMatrix(A, "....State Space: A"); }
	omxMatrix* x = ose->x;
	if(OMX_DEBUG_ALGEBRA) {omxPrintMatrix(x, "....State Space: x smoothed"); }
	omxMatrix* z = ose->z;
	if(OMX_DEBUG_ALGEBRA) {omxPrintMatrix(z, "....State Space: x updated"); }
	omxMatrix* P = ose->P;
	if(OMX_DEBUG_ALGEBRA) {omxPrintMatrix(P, "....State Space: P smoothed"); }
	omxMatrix* Z = ose->Z;
	if(OMX_DEBUG_ALGEBRA) {omxPrintMatrix(Z, "....State Space: P updated"); }

	/* Eigen Matrix reference setting */
	Eigen::MatrixXd &eigenPreX = ose->eigenPreX;
	if(OMX_DEBUG_ALGEBRA) {std::cout << "... State Space x predicted:\n" << eigenPreX << std::endl; }
	Eigen::MatrixXd &eigenExpA = ose->eigenExpA;
	if(OMX_DEBUG_ALGEBRA) {std::cout << "... State Space P predicted:\n" << eigenExpA << std::endl; }
	Eigen::MatrixXd &eigenIA = ose->eigenIA; // Storage for Sg (RTS smoother gain matrix)

	/* Eigen Adaptor copies (not really copies, more like wrappers) */
	EigenMatrixAdaptor eigenA(A);
	EigenMatrixAdaptor eigenx(x);
	EigenMatrixAdaptor eigenz(z);
	EigenMatrixAdaptor eigenP(P);
	EigenMatrixAdaptor eigenZ(Z);

	/* Create the RTS Gain matrix*/
	// Sg = Pui * A * Ppi+1 ^-1
	// eigenIA = Z * A * eigenExpA^-1
	// Possible typo above, A should be A^T
	eigenIA = eigenExpA.lu().solve( eigenA * eigenZ ).transpose();
	// try also
	//eigenIA = eigenExpA.ldlt().solve( eigenA.transpose() * eigenZ ).transpose();
	// with #include <Eigen/Cholesky>

	/* Smooth the latent state */
	// xsi = xui + Sg * (xsi+1 - xpi+1)
	// x = z + eigenIA * (x - eigenPreX)
	eigenx.derived() = eigenz + eigenIA * (eigenx - eigenPreX);

	/* Smooth the latent covariance */
	// Psi = Pui + Sg * (Psi+1 - Ppi+1) * Sg^T
	// P = Z + eigenIA * (P - eigenExpA) * eigenIA^T
	eigenP.derived() = eigenZ + eigenIA * (eigenP - eigenExpA) * eigenIA.transpose();

	// TODO add eigenPreX to struct and initialize
}

void omxInitStateSpaceExpectation(omxExpectation* ox) {

	SEXP rObj = ox->rObj;
	if(OMX_DEBUG) { mxLog("Initializing State Space Expectation."); }

	int nx, ny, nu;

	//SEXP slotValue;   //Used by PPML

	/* Create and fill expectation */
	//omxStateSpaceExpectation *SSMexp = (omxStateSpaceExpectation*) R_alloc(1, sizeof(omxStateSpaceExpectation));
	omxStateSpaceExpectation *SSMexp = new omxStateSpaceExpectation;

	omxState* currentState = ox->currentState;

	/* Set Expectation Calls and Structures */
	ox->computeFun = omxCallStateSpaceExpectation;
	ox->destructFun = omxDestroyStateSpaceExpectation;
	ox->componentFun = omxGetStateSpaceExpectationComponent;
	ox->mutateFun = omxSetStateSpaceExpectationComponent;
	ox->populateAttrFun = omxPopulateSSMAttributes;
	ox->argStruct = (void*) SSMexp;
	ox->canDuplicate = false;

	/* Set up expectation structures */
	if(OMX_DEBUG) { mxLog("Initializing State Space Meta Data for expectation."); }

	if(OMX_DEBUG) { mxLog("Processing A."); }
	SSMexp->A = omxNewMatrixFromSlot(rObj, currentState, "A");

	if(OMX_DEBUG) { mxLog("Processing B."); }
	SSMexp->B = omxNewMatrixFromSlot(rObj, currentState, "B");

	if(OMX_DEBUG) { mxLog("Processing C."); }
	SSMexp->C = omxNewMatrixFromSlot(rObj, currentState, "C");

	if(OMX_DEBUG) { mxLog("Processing D."); }
	SSMexp->D = omxNewMatrixFromSlot(rObj, currentState, "D");

	if(OMX_DEBUG) { mxLog("Processing Q."); }
	SSMexp->Q = omxNewMatrixFromSlot(rObj, currentState, "Q");

	if(OMX_DEBUG) { mxLog("Processing R."); }
	SSMexp->R = omxNewMatrixFromSlot(rObj, currentState, "R");

	if(OMX_DEBUG) { mxLog("Processing initial x."); }
	SSMexp->x0 = omxNewMatrixFromSlot(rObj, currentState, "x0");

	if(OMX_DEBUG) { mxLog("Processing initial P."); }
	SSMexp->P0 = omxNewMatrixFromSlot(rObj, currentState, "P0");

	if(OMX_DEBUG) { mxLog("Processing u."); }
	SSMexp->u = omxNewMatrixFromSlot(rObj, currentState, "u");

	if(OMX_DEBUG) { mxLog("Processing t."); }
	SSMexp->t = omxNewMatrixFromSlot(rObj, currentState, "t");

	/* Initialize the place holder matrices used in calculations */
	nx = SSMexp->C->cols;
	ny = SSMexp->C->rows;
	nu = SSMexp->D->cols;

	if(OMX_DEBUG) { mxLog("Processing first data row for y."); }
	SSMexp->y = omxInitMatrix(ny, 1, TRUE, currentState);
	for(int i = 0; i < ny; i++) {
		omxSetMatrixElement(SSMexp->y, i, 0, omxDoubleDataElement(ox->data, 0, i));
	}
	if(OMX_DEBUG_ALGEBRA) {omxPrintMatrix(SSMexp->y, "....State Space: y"); }

	// TODO Make x0 and P0 static (if possible) to save memory
	// TODO Look into omxMatrix.c/h for a possible new matrix from omxMatrix function
	if(OMX_DEBUG) { mxLog("Generating static internals for resetting initial values."); }
	SSMexp->x = 	omxInitMatrix(nx, 1, TRUE, currentState);
	SSMexp->P = 	omxInitMatrix(nx, nx, TRUE, currentState);
	omxCopyMatrix(SSMexp->x, SSMexp->x0);
	omxCopyMatrix(SSMexp->P, SSMexp->P0);

	if(OMX_DEBUG) { mxLog("Generating internals for computation."); }

	SSMexp->covInfo = 	omxInitMatrix(1, 1, TRUE, currentState);
	SSMexp->det = 	omxInitMatrix(1, 1, TRUE, currentState);
	SSMexp->r = 	omxInitMatrix(ny, 1, TRUE, currentState);
	SSMexp->s = 	omxInitMatrix(ny, 1, TRUE, currentState);
	SSMexp->z = 	omxInitMatrix(nx, 1, TRUE, currentState);
	SSMexp->K = 	omxInitMatrix(ny, nx, TRUE, currentState); // Actually the tranpose of the Kalman gain
	SSMexp->S = 	omxInitMatrix(ny, ny, TRUE, currentState);
	SSMexp->Y = 	omxInitMatrix(ny, nx, TRUE, currentState);
	SSMexp->Z = 	omxInitMatrix(nx, nx, TRUE, currentState);

	SSMexp->cov = 		omxInitMatrix(ny, ny, TRUE, currentState);
	SSMexp->means = 	omxInitMatrix(1, ny, TRUE, currentState);

	/* Create alias matrices for missing data filtering */
	SSMexp->smallC = 	omxInitMatrix(ny, nx, TRUE, currentState);
	SSMexp->smallD = 	omxInitMatrix(ny, nu, TRUE, currentState);
	SSMexp->smallR = 	omxInitMatrix(ny, ny, TRUE, currentState);
	SSMexp->smallr = 	omxInitMatrix(ny, 1, TRUE, currentState);
	SSMexp->smallK = 	omxInitMatrix(ny, nx, TRUE, currentState);
	SSMexp->smallS = 	omxInitMatrix(ny, ny, TRUE, currentState);
	SSMexp->smallY = 	omxInitMatrix(ny, nx, TRUE, currentState);

	//SSMexp->deltaT = 	omxInitMatrix(1, 1, TRUE, currentState);
	SSMexp->oldT = 0.0;
	SSMexp->deltaT = 0.0;

	/* Eigen Matrix initialization */
	SSMexp->eigenExpA.resize(nx, nx);
	SSMexp->I.resize(nx, nx);
	SSMexp->I = Eigen::MatrixXd::Identity(nx, nx);
	SSMexp->eigenIA.resize(nx, nx);
	SSMexp->PSI.resize(2*nx, 2*nx);
	SSMexp->IP.resize(2*nx, nx);
	SSMexp->eigenPreX.resize(nx, 1);

	/* Population of Kalman scores*/
	if(OMX_DEBUG) {
		mxLog("Accessing Kalman score population option.");
	}
	SSMexp->returnScores = Rf_asInteger(R_do_slot(rObj, Rf_install("scores")));

	omxCopyMatrix(SSMexp->smallC, SSMexp->C);
	omxCopyMatrix(SSMexp->smallD, SSMexp->D);
	omxCopyMatrix(SSMexp->smallR, SSMexp->R);
	omxCopyMatrix(SSMexp->smallr, SSMexp->r);
	omxCopyMatrix(SSMexp->smallK, SSMexp->K);
	omxCopyMatrix(SSMexp->smallS, SSMexp->S);
	omxCopyMatrix(SSMexp->smallY, SSMexp->Y);

}

omxMatrix* omxGetStateSpaceExpectationComponent(omxExpectation* ox, omxFitFunction* off, const char* component) {
	omxStateSpaceExpectation* ose = (omxStateSpaceExpectation*)(ox->argStruct);
	omxMatrix* retval = NULL;

	if(strEQ("cov", component)) {
		retval = ose->cov;
	} else if(strEQ("means", component)) {
		retval = ose->means;
	} else if(strEQ("pvec", component)) {
		// Once implemented, change compute function and return pvec
	} else if(strEQ("inverse", component)) {
		retval = ose->smallS;
	} else if(strEQ("determinant", component)) {
		retval = ose->det;
	} else if(strEQ("r", component)) {
		retval = ose->r;
	} else if(strEQ("covInfo", component)) {
		retval = ose->covInfo;
	}

	return retval;
}

void omxSetStateSpaceExpectationComponent(omxExpectation* ox, omxFitFunction* off, const char* component, omxMatrix* om) {
	omxStateSpaceExpectation* ose = (omxStateSpaceExpectation*)(ox->argStruct);

	if(!strcmp("y", component)) {
		for(int i = 0; i < ose->y->rows; i++) {
			omxSetMatrixElement(ose->y, i, 0, omxVectorElement(om, i));
		}
		//ose->y = om;
	}
	if(!strcmp("Reset", component)) {
		omxRecompute(ose->x0, NULL);
		omxRecompute(ose->P0, NULL);
		omxCopyMatrix(ose->x, ose->x0);
		omxCopyMatrix(ose->P, ose->P0);
		if(ose->t != NULL){
			ose->oldT = 0.0;
		}
	}
}


// Software License for MTL
//
// Copyright (c) 2007 The Trustees of Indiana University.
//               2008 Dresden University of Technology and the Trustees of Indiana University.
//               2010 SimuNova UG (haftungsbeschränkt), www.simunova.com.
// All rights reserved.
// Authors: Peter Gottschling and Andrew Lumsdaine
//
// This file is part of the Matrix Template Library
//
// See also license.mtl.txt in the distribution.

// Written by Jan Bos
// Edited  by Peter Gottschling

#ifndef ITL_BICGSTAB_ELL_INCLUDE
#define ITL_BICGSTAB_ELL_INCLUDE

#include <boost/numeric/mtl/concept/collection.hpp>
#include <boost/numeric/mtl/vector/dense_vector.hpp>
#include <boost/numeric/mtl/matrix/strict_upper.hpp>
#include <boost/numeric/mtl/matrix/dense2D.hpp>
#include <boost/numeric/mtl/utility/irange.hpp>
#include <boost/numeric/mtl/operation/orth.hpp>
#include <boost/numeric/mtl/utility/exception.hpp>
#include <boost/numeric/linear_algebra/identity.hpp>
#include <boost/numeric/mtl/operation/resource.hpp>
#include <boost/numeric/mtl/operation/size.hpp>
#include <boost/numeric/mtl/interface/vpt.hpp>

#include <boost/numeric/itl/itl_fwd.hpp>
#include <boost/numeric/itl/krylov/base_solver.hpp>

namespace itl {

/// Bi-Conjugate Gradient Stabilized(ell)
template < typename LinearOperator, typename Vector,
	   typename LeftPreconditioner, typename RightPreconditioner,
	   typename Iteration >
int bicgstab_ell(const LinearOperator &A, Vector &x, const Vector &b,
		 const LeftPreconditioner &L, const RightPreconditioner &R,
		 Iteration& iter, size_t l)
{
    mtl::vampir_trace<7006> tracer;
    using mtl::size; using mtl::irange; using mtl::imax; using mtl::mat::strict_upper;
    typedef typename mtl::Collection<Vector>::value_type Scalar;
    typedef typename mtl::Collection<Vector>::size_type  Size;

    if (size(b) == 0) throw mtl::logic_error("empty rhs vector");

    const Scalar                zero= math::zero(Scalar()), one= math::one(Scalar());
    Vector                      x0(resource(x)), y(resource(x));
    mtl::dense_vector<Vector>   r_hat(l+1,Vector(resource(x))), u_hat(l+1,Vector(resource(x)));

    // shift problem
    x0= zero;
    r_hat[0]= b;
    if (two_norm(x) != zero) {
	r_hat[0]-= A * x;
	x0= x;
	x= zero;
    }

    Vector  r0_tilde(r_hat[0]/two_norm(r_hat[0]));
    y= solve(L, r_hat[0]);
    r_hat[0]= y;
    u_hat[0]= zero;

    Scalar                      rho_0(one), rho_1(zero), alpha(zero), Gamma(zero), beta(zero), omega(one);
    mtl::mat::dense2D<Scalar>        tau(l+1, l+1);
    mtl::dense_vector<Scalar>   sigma(l+1), gamma(l+1), gamma_a(l+1), gamma_aa(l+1);

    while (! iter.finished(r_hat[0])) {
	++iter;
	rho_0= -omega * rho_0;

	for (Size j= 0; j < l; ++j) {
	    rho_1= dot(r0_tilde, r_hat[j]);
	    beta= alpha * rho_1/rho_0; rho_0= rho_1;

	    for (Size i= 0; i <= j; ++i)
		u_hat[i]= r_hat[i] - beta * u_hat[i];

	    y= A * Vector(solve(R, u_hat[j]));
	    u_hat[j+1]= solve(L, y);
	    Gamma= dot(r0_tilde, u_hat[j+1]);
	    alpha= rho_0 / Gamma;

	    for (Size i= 0; i <= j; ++i)
		r_hat[i]-= alpha * u_hat[i+1];

	    if (iter.finished(r_hat[j])) {
		x= solve(R, x);
		x+= x0;
		return iter;
	    }

	    r_hat[j+1]= solve(R, r_hat[j]);
	    y= A * r_hat[j+1];
	    r_hat[j+1]= solve(L, y);
	    x+= alpha * u_hat[0];
	}

	// mod GS (MR part)
	irange  i1m(1, imax);
	mtl::dense_vector<Vector>   r_hat_tail(r_hat[i1m]);
	tau[i1m][i1m]= orthogonalize_factors(r_hat_tail);
	for (Size j= 1; j <= l; ++j)
	    gamma_a[j]= dot(r_hat[j], r_hat[0]) / tau[j][j];

	gamma[l]= gamma_a[l]; omega= gamma[l];
	if (omega == zero) return iter.fail(3, "bicg breakdown #2");

	// is this something like a tri-solve?
	for (Size j= l-1; j > 0; --j) {
	    Scalar sum= zero;
	    for (Size i=j+1;i<=l;++i)
		sum += tau[j][i] * gamma[i];
	    gamma[j] = gamma_a[j] - sum;
	}

	gamma_aa[irange(1, l)]= strict_upper(tau[irange(1, l)][irange(1, l)]) * gamma[irange(2, l+1)] + gamma[irange(2, l+1)];

	x+= gamma[1] * r_hat[0];
	r_hat[0]-= gamma_a[l] * r_hat[l];
	u_hat[0]-= gamma[l] * u_hat[l];
	for (Size j=1; j < l; ++j) {
	    u_hat[0] -= gamma[j] * u_hat[j];
	    x+= gamma_aa[j] * r_hat[j];
	    r_hat[0] -= gamma_a[j] * r_hat[j];
	}
    }
    x= solve(R, x); x+= x0; // convert to real solution and undo shift
    return iter;
}

/// Solver class for BiCGStab(ell) method
/** Methods inherited from \ref base_solver. **/
template < typename LinearOperator, typename Preconditioner,
	   typename RightPreconditioner>
class bicgstab_ell_solver
  : public base_solver< bicgstab_ell_solver<LinearOperator, Preconditioner, RightPreconditioner>, LinearOperator >
{
    typedef base_solver< bicgstab_ell_solver<LinearOperator, Preconditioner, RightPreconditioner>, LinearOperator > base;
  public:
    /// Construct solver from a linear operator; generate (left) preconditioner from it
    explicit bicgstab_ell_solver(const LinearOperator& A, size_t l= 8) : base(A), l(l), L(A), R(A) {}

    /// Construct solver from a linear operator and left preconditioner
    bicgstab_ell_solver(const LinearOperator& A, size_t l, const Preconditioner& L) : base(A), l(l), L(L), R(A) {}

    /// Construct solver from a linear operator and left preconditioner
    bicgstab_ell_solver(const LinearOperator& A, size_t l, const Preconditioner& L, const RightPreconditioner& R)
      : base(A), l(l), L(L), R(R) {}

    /// Solve linear system approximately as specified by \p iter
    template < typename HilbertSpaceX, typename HilbertSpaceB, typename Iteration >
    int solve(HilbertSpaceX& x, const HilbertSpaceB& b, Iteration& iter) const
    {
	return bicgstab_ell(this->A, x, b, L, R, iter, l);
    }

  private:
    size_t                l;
    Preconditioner        L;
    RightPreconditioner   R;
};

} // namespace itl

#endif // ITL_BICGSTAB_ELL_INCLUDE


#include "logisticregression.h"
#include "ui_logisticregression.h"

#include <cmath>
#include <iostream>
#include <stdio.h>
#include <vector>
#include <sys/time.h>
#include <chrono>
#include <sys/resource.h>   // check the memory usage
#include <stdio.h>
#include <thread>
#include <fstream>
#include <sstream>

#include <NTL/RR.h>
#include <NTL/xdouble.h>
#include <NTL/ZZ.h>
#include <NTL/ZZX.h>
#include <NTL/BasicThreadPool.h>

#include "HELR/CZZ.h"
#include "HELR/Params.h"
#include "HELR/PubKey.h"
#include "HELR/Scheme.h"
#include "HELR/SchemeAlgo.h"
#include "HELR/SecKey.h"
#include "HELR/TestScheme.h"
#include "HELR/TimeUtils.h"
#include "HELR/Ring2Utils.h"
#include "HELR/StringUtils.h"
#include "HELR/EvaluatorUtils.h"

#include "Database.h"
#include "LRtest.h"
#include "HELR.h"
#include <QFile>
#include <QtDebug>
#include <QFileDialog>
#include "mainwindow.h"
using namespace NTL;
using namespace std;

LogisticRegression::LogisticRegression(QWidget *parent) :
    QWidget(parent),
    ui(new Ui::LogisticRegression)
{
    ui->setupUi(this);
    QPalette bgpal = palette();
    bgpal.setColor (QPalette::Background, QColor (0, 0 , 0, 255));
    bgpal.setColor (QPalette::Foreground, QColor (255,255,255,255)); setPalette (bgpal);
}

LogisticRegression::~LogisticRegression()
{
    delete ui;
}

void LogisticRegression::testHELR()
{
    /*
    if(Argc != 3){
        cout << "-------------------------------------------------------------" << endl;
        cerr << "Enter the File and degree of approximation \t"  << "(e.g. $test edin.txt 3) \n ";
    }
    */
    //qDebug()<< QDir::currentPath();
    //  char* filename  =  Argv[1];
    // int polydeg = atoi(Argv[2]);  // degree of approximation polynomial
//    char* filename = "data/edin.txt";

    /*QFile file(filename);
    if(!file.open(QIODevice::ReadOnly | QIODevice::Text))
    {
        qDebug()<<file.errorString();
    }else{
        cout<<"errorrrr"<<endl;
    }*/
//    int polydeg = 3;
    dMat  zData;
    dMat* zTest = new dMat[5];
    dMat* zTrain = new dMat[5];

    freopen("mid_result.txt","w",stdout);
    int nLine= readData(zData, filename);

    cout << "Sample the learning and test data ..." << endl;
    cvRandomSamplingData(zTrain, zTest, zData, filename);

    //----------------------------------------------------------------
    // Parameters for Logistic regression
    //----------------------------------------------------------------

    long logN= 11;
    long logp= 28;
    long logl= 10;
    long logq, cBit1, cBit2;
    int max_iter;

    struct LRpar LRparams;
    ReadLRparams(LRparams, max_iter, zTrain[0], polydeg, logp);

    SetNumThreads(LRparams.dim1);
    //SetNumThreads(4);

    switch(polydeg){
    case 3:
        cBit1=  (LRparams.logn - LRparams.log2polyscale);          // 1st iteration
        cBit2 =  (3*logp+ LRparams. logn - LRparams.log2polyscale);  // 2nd~ iteration
        logq = cBit1 + (LRparams.max_iter-1)*(cBit2)+ logp + logl;                  // max-bitlength we need
        break;

    case 7:
        cBit1=  (LRparams.logn - LRparams.log2polyscale);          // 1st iteration
        cBit2=  (4*logp+ LRparams.logn - LRparams.log2polyscale);
        logq= cBit1 + (LRparams.max_iter-1)*(cBit2)+ logp + logl;
        break;
    }

    freopen("cipher training.txt","w",stdout);

    cout << "Data dimension with dummy vectors: " << LRparams.dim1 << ", Number of lines: " << nLine << endl;

    cout << "-------------------------------------------------------------" << endl;
    cout << "Key Generation ... (logN,logp,logq, nslots)= (" ;
    cout << logN << "," << logp << "," << logq << "," << LRparams.nslots << ")" <<endl;

    auto start= chrono::steady_clock::now();

    Params params(logN, logq);
    SecKey secretKey(params);
    PubKey publicKey(params, secretKey);
    SchemeAux schemeaux(logN);
    Scheme scheme(params, publicKey, schemeaux);
    SchemeAlgo algo(scheme);

    auto end = std::chrono::steady_clock::now();
    auto diff = end - start;
    cout << "KeyGen time= " << chrono::duration <double, milli> (diff).count()/1000.0 << " s" << endl;

    LogReg LR(scheme, secretKey, LRparams);
    dMat HEtheta_list;

    int count = 5;
    for(int k = 0; k < count; ++k){

        cout << "-------------------------------------------------------------" << endl;
        cout << k << "th Data Encryption ... " << endl;

        struct rusage usage;

        start= chrono::steady_clock::now();

        Cipher* zTrainCipher = new Cipher[LRparams.dim1];

        LR.EncryptData(zTrainCipher, zTrain[k]);

        end = std::chrono::steady_clock::now();
        diff = end - start;
        cout << "Enc time= "  << chrono::duration <double, milli> (diff).count()/1000.0 << "(s), " ;

        int ret = getrusage(RUSAGE_SELF,&usage);
        cout<< "Mem: " << usage.ru_maxrss/(1024)  << "(MB)" << endl;

        cout << "-------------------------------------------------------------" << endl;
        cout << "HE Logistic Regression ... "  << endl;

        Cipher* thetaCipher= new Cipher[LRparams.dim1];

        start= chrono::steady_clock::now();

        LR.HElogreg(thetaCipher, zTrainCipher, zTrain[k]);

        end = std::chrono::steady_clock::now();
        diff = end - start;
        cout << "Eval time= "  << chrono::duration <double, milli> (diff).count()/1000.0 << "(s), " ;

        ret = getrusage(RUSAGE_SELF,&usage);
        cout<< "Mem: " << usage.ru_maxrss/(1024)  << "(MB)" << endl;

        cout << "-------------------------------------------------------------" << endl;
        cout << "Decryption ... "  << endl;

        dVec HEtheta(LRparams.dim1, 0.0);

        CZZ* dtheta = new CZZ[LRparams.dim1];

        for(int i=0; i< LRparams.dim1; ++i){
            dtheta[i] = (scheme.decrypt(secretKey, thetaCipher[i]))[0];

            conv(HEtheta[i], dtheta[i].r);
            HEtheta[i] = scaledown(HEtheta[i], LRparams.logp);
            cout << "[" << HEtheta[i] << "] " ;
        }
        cout << ": enc " << endl;

        getAUC(HEtheta, zTest[k]);
        HEtheta_list.push_back(HEtheta);

        cout << "-------------------------------------------------------------" << endl;
        cout << "Compare with unenc LR " << endl;
        dVec mtheta(LRparams.dim1, 0.0);

        for(int i= 0; i< LRparams.max_iter; i++){
            LR_poly(mtheta, zTrain[k], LRparams);
        }

        for(int i= 0; i< LRparams.dim1; i++)
            cout << "[" << mtheta[i] << "] " ;
        cout << ": unenc " << endl;

        getAUC(mtheta, zTest[k]);

        cout << "MSE (HELR/non-HELR): " << getMSE(HEtheta, mtheta) << endl;

        cout << "-------------------------------------------------------------" << endl;
        cout << "Compare with sigmoid LR " << endl;
        dVec mtheta_sig(LRparams.dim1, 0.0);

        for(int i= 0; i< LRparams.max_iter; i++){
            LR_sigmoid(mtheta_sig, zTrain[k], LRparams);
        }

        for(int i= 0; i< LRparams.dim1; i++)
            cout << "[" << mtheta_sig[i] << "] " ;
        cout << ": unenc " << endl;

        getAUC(mtheta_sig, zTest[k]);

        cout << "MSE (HELR/non-HELR): " << getMSE(HEtheta, mtheta_sig) << endl;

    }
    ofstream fout;
    fout.open("beta_HELR.txt");
    //! write the beta results in the text file
    fout << "-------------------------------------------------------------" << endl;
    fout << "[" << endl;
    //        for(int i = 0; i < LRparams.dim1; ++i){
    //            for(int k = 0; k < count-1; ++k){
    //                fout << HEtheta_list[k][i] << "," ;
    //            }
    //            fout << HEtheta_list[count-1][i] << ";" << endl;
    //        }

    for (int k=0;k<count;k++) {
        for(int i = 0; i < LRparams.dim1-1; ++i){
            fout << HEtheta_list[k][i] << "," ;
        }
        fout << HEtheta_list[k][LRparams.dim1-1]<<";";
        fout<<"\n";
    }
//    for(int i = 0; i < LRparams.dim1; ++i){
//        for(int k = 0; k < count-1; ++k){
//            fout << HEtheta_list[k][i] << "," ;
//        }
//        fout << HEtheta_list[count-1][i] << ";" << endl;
//    }
    //    for(int i = 0; i < HEtheta_list.size(); ++i){
    //        for(int k = 0; k < HEtheta_list[0].size(); ++k){
    //            fout << HEtheta_list[i][k] << "," ;

    //        }
    //        //fout << HEtheta_list[i][i] << ";" << endl;
    //    }
    fout << "];" << endl;
    fout.close();

    delete[] zTest;
    delete[] zTrain;
    ShowTxtToWindowCip();

}

void LogisticRegression::testLR()
{
    /*if(Argc != 2){
        cout << "-------------------------------------------------------------" << endl;
        cerr << "Enter the File and degree of approximation \t"  << "(e.g. $test edin.txt 3) \n ";
    }
    cout<<"Argc:"<<Argc<<endl;*/
//    filename  =  "data/edin.txt";

    /*
    QFile file(filename);
    if(!file.open(QIODevice::ReadOnly | QIODevice::Text))
    {
        qDebug()<<file.errorString();
    }else{
        cout<<"success"<<endl;
    }*/
    dMat  zData;
    dMat* zTest = new dMat[5];
    dMat* zTrain = new dMat[5];

    int nLine= readData(zData, filename);
    /*
    for(int i=0;i<zData.size();++i){
        for(int j=0;j<zData[0].size();++j){
            cout<<zData[i][j]<<" ";
        }
        cout<<endl;
    }*/

    freopen("plaintext training.txt","w",stdout);

    cout << "Sample the learning and test data ..." << endl;
    cvRandomSamplingData(zTrain, zTest, zData, filename);

    //----------------------------------------------------------------
    // Parameters for Logistic regression
    //----------------------------------------------------------------

    long logN= 11;
    long logp= 28;
    long logl= 10;
    long logq, cBit1, cBit2;
    int max_iter;
    long dim ;

    dMat mtheta3_list;
    dMat mtheta7_list;
    dMat mtheta_sig_list;

    ofstream fout;
    fout.open("beta_LR.txt");

    for(int k = 0; k < 5; ++k){

        cout << "-------------------------------------------------------------" << endl;
        cout << "HELR_3 " << endl;

        //int polydeg = 3;  // degree of approximation polynomial

        struct LRpar LRparams;
        ReadLRparams(LRparams, max_iter, zTrain[0], polydeg, logp);

        dVec mtheta3(LRparams.dim1, 0.0);
        for(int i= 0; i< LRparams.max_iter; i++){
            LR_poly(mtheta3, zTrain[k], LRparams);
        }

        for(int i= 0; i< LRparams.dim1; i++)
            cout << "[" << mtheta3[i] << "] " ;
        cout  << endl;

        getAUC(mtheta3, zTest[k]);
        mtheta3_list.push_back(mtheta3);

        cout << "-------------------------------------------------------------" << endl;
        cout << "HELR_7 " << endl;

        polydeg = 7;  // degree of approximation polynomial

        struct LRpar LRparams7;
        ReadLRparams(LRparams7, max_iter, zTrain[0], polydeg, logp);

        dVec mtheta7(LRparams7.dim1, 0.0);
        for(int i= 0; i< LRparams7.max_iter; i++){
            LR_poly(mtheta7, zTrain[k], LRparams7);
        }

        for(int i= 0; i< LRparams7.dim1; i++)
            cout << "[" << mtheta7[i] << "] " ;
        cout  << endl;

        getAUC(mtheta7, zTest[k]);
        mtheta7_list.push_back(mtheta7);

        dim = LRparams7.dim1;

        cout << "-------------------------------------------------------------" << endl;
        cout << "sigmoid LR " << endl;
        dVec mtheta_sig(LRparams.dim1, 0.0);

        for(int i= 0; i< LRparams.max_iter; i++){
            LR_sigmoid(mtheta_sig, zTrain[k], LRparams);
        }

        for(int i= 0; i< LRparams.dim1; i++)
            cout << "[" << mtheta_sig[i] << "] " ;
        cout  << endl;

        getAUC(mtheta_sig, zTest[k]);
        mtheta_sig_list.push_back(mtheta_sig);

        cout << "-------------------------------------------------------------" << endl;
        cout << "MSE (HELR/non-HELR): " << getMSE(mtheta3, mtheta_sig) << endl;
        cout << "MSE (HELR/non-HELR): " << getMSE(mtheta7, mtheta_sig) << endl;

    }

    //! write the beta results in the text file
    fout << "-------------------------------------------------------------" << endl;
    fout << "HELR_3" << endl<<"[\n";

    for (int k=0;k<5;k++) {
        for(int i = 0; i < dim-1; ++i){
            fout << mtheta3_list[k][i] << "," ;
        }
        fout << mtheta3_list[k][dim-1]<<";";
        fout<<"\n";
    }
    fout<<"]"<<endl;

    /*
    for(int i = 0; i < dim; ++i){
        for(int k = 0; k < 4; ++k){
            fout << mtheta3_list[k][i] << "," ;
        }
        fout << mtheta3_list[4][i] << ";" << endl;
    }
    fout<<"]"<<endl;
    */
    fout << "-------------------------------------------------------------" << endl;
    fout << "HELR_7" << endl<<"[";
    for (int k=0;k<5;k++) {
        for(int i = 0; i < dim-1; ++i){
            fout << mtheta7_list[k][i] << "," ;
        }
        fout << mtheta7_list[k][dim-1]<<";";
        fout<<"\n";
    }
    fout<<"]"<<endl;
//    for(int i = 0; i < dim; ++i){
//        for(int k = 0; k < 4; ++k){
//            fout << mtheta7_list[k][i] << "," ;
//        }
//        fout << mtheta7_list[4][i] << ";" << endl;
//    }
//    fout << "]" << endl;

    fout << "-------------------------------------------------------------" << endl;
    fout << "LR" << endl<<"[";
    for (int k=0;k<5;k++) {
        for(int i = 0; i < dim-1; ++i){
            fout << mtheta_sig_list[k][i] << "," ;
        }
        fout << mtheta_sig_list[k][dim-1]<<";";
        fout<<"\n";
    }
    fout << "]" << endl;
//    for(int i = 0; i < dim; ++i){
//        for(int k = 0; k < 4; ++k){
//            fout << mtheta_sig_list[k][i] << "," ;
//        }
//        fout << mtheta_sig_list[4][i] << ";" << endl;
//    }
//    fout<<"]"<<endl;
    fout.close();

    fout << "-------------------------------------------------------------" << endl;

    delete[] zTest;
    delete[] zTrain;

    ShowTxtToWindow();
}

void LogisticRegression::on_pushButton_2_clicked()
{
    testLR();
}

void LogisticRegression::ShowTxtToWindow()//显示文本文件中的内容
{
    QString fileName = "plaintext training.txt";

    if(!fileName.isEmpty())
    {
        QFile *file = new QFile;
        file->setFileName(fileName);
        bool ok = file->open(QIODevice::ReadOnly);
        if(ok)
        {
            QTextStream in(file);
            ui->plainResult->setText(in.readAll());
            file->close();
            delete file;
        }
        else
        {
            QMessageBox::information(this,"错误信息","打开文件:" + file->errorString());
            return;
        }
    }
}

void LogisticRegression::ShowTxtToWindowCip()
{
    QString fileName = "cipher training.txt";

    if(!fileName.isEmpty())
    {
        QFile *file = new QFile;
        file->setFileName(fileName);
        bool ok = file->open(QIODevice::ReadOnly);
        if(ok)
        {
            QTextStream in(file);
            ui->cResult->setText(in.readAll());
            file->close();
            delete file;
        }
        else
        {
            QMessageBox::information(this,"错误信息","打开文件:" + file->errorString());
            return;
        }
    }
}

void LogisticRegression::on_return_2_clicked()
{
    MainWindow *win = new MainWindow;
    this->hide();
    win->show();
}

void LogisticRegression::on_file_clicked()
{
    //定义文件对话框类
    QFileDialog *fileDialog = new QFileDialog(this);
    //定义文件对话框标题
    fileDialog->setWindowTitle(tr("选择数据集"));
    //设置默认文件路径
    fileDialog->setDirectory(".");
    //设置可以选择多个文件,默认为只能选择一个文件QFileDialog::ExistingFiles
    fileDialog->setFileMode(QFileDialog::ExistingFiles);
    //设置视图模式
    fileDialog->setViewMode(QFileDialog::Detail);
    //打印所有选择的文件的路径
    QStringList fileNames;
    if(fileDialog->exec())
    {
        fileNames = fileDialog->selectedFiles();
    }

    QString QString_fileNames = fileNames.join(",");

    QString curPath = QDir::currentPath();
    QString relPath = QString_fileNames.mid(curPath.length()+1);

    QByteArray ba = relPath.toLatin1();
    strcpy(filename,ba.data());
    cout<<filename<<endl;
}

void LogisticRegression::on_comboBox_activated(const QString &arg1)
{
    QMap<QString, int> map_polydeg;
    map_polydeg.insert("3(默认)",3);
    map_polydeg.insert("7",7);
    polydeg = map_polydeg[arg1];
}

void LogisticRegression::on_pushButton_clicked()
{
    testHELR();
}


#include "geodb/klee.hpp"

#include "geodb/interval.hpp"
#include "geodb/utility/range_utils.hpp"

#include <boost/functional/hash.hpp>

#include <memory>

namespace geodb {

namespace {

using interval_t = interval<double>;

size_t log2_ceil(size_t value) {
    geodb_assert(value != 0, "invalid value");

    size_t log2 = 0;
    value--;
    while (value != 0) {
        ++log2;
        value /= 2;
    }
    return log2;
}

class segment_tree {
    struct node {
        /// The interval that this node represents.
        interval_t interval;

        /// The current width of the union of all intervals
        /// active in the subtree rooted at this node.
        double union_width = 0;

        /// Number of active intervals that contain this interval.
        size_t count = 0;

        /// True if this node is a leaf node.
        bool leaf = false;
    };

public:
    /// Construct a new, empty segment tree from the given iterator range.
    ///
    /// The range must contain the Universe of endpoints, i.e. every
    /// possible interval end point.
    /// The points must be sorted in ascending order and must be unique.
    /// The range must have at least two points.
    template<typename RndIter>
    segment_tree(RndIter begin, RndIter end)
    {
        reset(begin, end);
    }

    /// Replaces the structure of this tree using the points
    /// in the given sorted range.
    /// Has the same requirements and effects as the constructor,
    /// but reuses the allocated memory.
    template<typename RndIter>
    void reset(RndIter begin, RndIter end) {
        geodb_assert(std::distance(begin, end) >= 2,
                     "Must have at least two endpoints");

        // Allocate enough space for a complete binary tree.
        size_t leaves = std::distance(begin, end) - 1;
        size_t full_leaves = 1 << log2_ceil(leaves);
        size_t max_size = full_leaves * 2 - 1;

        m_nodes.resize(max_size);
        build(at(0), begin, end - 1);
    }

    /// Sets the given interval to "active", i.e. counting it
    /// in the calculation of the union width.
    /// Only intervals with endpoints in the previously registered universe
    /// can be inserted.
    void insert(const interval_t& interval) {
        insert(at(0), interval);
    }

    /// Removes the given interval, no longer counting it
    /// in the calculation of the union width.
    /// Only intervals that have previously been inserted can be removed.
    void remove(const interval_t& interval) {
        remove(at(0), interval);
    }

    double union_width() const {
        return at(0).union_width;
    }

private:
    /// Recursive build function.
    /// Both begin and end are valid iterators, i.e. its an inclusive range.
    /// All intervals [p1, p2] where p1, p2 are in the original input sequence
    /// will be represented by leaf nodes.
    template<typename Iter>
    void build(node& n, Iter begin, Iter end) {
        geodb_assert(begin != end, "Range is empty");

        n.count = 0;
        n.union_width = 0;
        n.leaf = (begin + 1 == end);
        if (n.leaf) {
            n.interval = interval_t(*begin, *end);
        } else {
            node& l = left(n);
            node& r = right(n);

            Iter mid = begin + (end - begin) / 2;
            build(l, begin, mid);
            build(r, mid, end);
            n.interval = interval_t(l.interval.begin(), r.interval.end());
        }
    }

    void insert(node& n, const interval_t& interval) {
        if (interval.contains(n.interval)) {
            n.count += 1;
            if (n.count == 1) {
                n.union_width = n.interval.end() - n.interval.begin();
            }
        } else {
            geodb_assert(!n.leaf, "must not be a leaf");

            node& l = left(n);
            node& r = right(n);

            if (go_left(l.interval, interval)) {
                insert(l, interval);
            }
            if (go_right(r.interval, interval)) {
                insert(r, interval);
            }
            if (n.count == 0) {
                n.union_width = l.union_width + r.union_width;
            }
        }
    }

    void remove(node& n, const interval_t& interval) {
        if (interval.contains(n.interval)) {
            n.count -= 1;
            if (n.count == 0) {
                n.union_width = n.leaf ? 0
                                       : left(n).union_width + right(n).union_width;
            }
        } else {
            geodb_assert(!n.leaf, "must not be a leaf");

            node& l = left(n);
            node& r = right(n);

            if (go_left(l.interval, interval)) {
                remove(l, interval);
            }
            if (go_right(r.interval, interval)) {
                remove(r, interval);
            }
            if (n.count == 0) {
                n.union_width = l.union_width + r.union_width;
            }
        }
    }

    bool go_left(const interval_t& left, const interval_t& interval) const {
        return interval.begin() < left.end();
    }

    bool go_right(const interval_t& right, const interval_t& interval) const {
        return interval.end() > right.begin();
    }

    node& left(node& n) {
        geodb_assert(!n.leaf, "n has no children");
        return at(index(n) * 2 + 1);
    }

    node& right(node& n) {
        geodb_assert(!n.leaf, "n has no children");
        return at(index(n) * 2 + 2);
    }

    node& at(size_t index) {
        geodb_assert(index < m_nodes.size(), "index out of bounds");
        return m_nodes[index];
    }

    const node& at(size_t index) const {
        geodb_assert(index < m_nodes.size(), "index out of bounds");
        return m_nodes[index];
    }

    size_t index(node& n) {
        return &n - m_nodes.data();
    }

private:
    std::vector<node> m_nodes;
};

enum event_type {
    open, close
};

struct event2d {
    event_type type;
    double x;
    interval_t y;

    event2d(event_type type, double x, const interval_t& y)
        : type(type)
        , x(x)
        , y(y)
    {}

    bool operator<(const event2d& other) const {
        if (x == other.x) {
            return type < other.type;
        }
        return x < other.x;
    }
};

struct event3d {
    event_type type;
    double x;
    rect2d r;

    event3d(event_type type, double x, const rect2d& r)
        : type(type)
        , x(x)
        , r(r)
    {}

    bool operator<(const event3d& other) const {
        if (x == other.x) {
            return type < other.type;
        }
        return x < other.x;
    }
};

} // namespace

/// Constructs the segment tree with the universe of all (unique)
/// rectangle y corner points.
static segment_tree build_segment_tree(const std::vector<rect2d>& rects) {
    std::vector<double> yvalues;
    yvalues.reserve(rects.size() * 2);
    for (const rect2d& rect : rects) {
        if (!rect.empty()) {
            yvalues.push_back(rect.min().y());
            yvalues.push_back(rect.max().y());
        }
    }

    std::sort(yvalues.begin(), yvalues.end());
    yvalues.erase(std::unique(yvalues.begin(), yvalues.end()), yvalues.end());

    return segment_tree(yvalues.begin(), yvalues.end());
}

static std::vector<event2d> rectangle_events(const std::vector<rect2d>& rects) {
    std::vector<event2d> events;
    events.reserve(2 * rects.size());
    for (const rect2d& rect : rects) {
        if (!rect.empty()) {
            interval_t y(rect.min().y(), rect.max().y());
            events.emplace_back(open, rect.min().x(), y);
            events.emplace_back(close, rect.max().x(), y);
        }
    }

    std::sort(events.begin(), events.end());
    return events;
}

static std::vector<event3d> rectangle_events(const std::vector<rect3d>& rects) {
    std::vector<event3d> events;
    events.reserve(2 * rects.size());
    for (const rect3d& rect : rects) {
        if (!rect.empty()) {
            rect2d rect2(vector2d(rect.min().y(), rect.min().z()),
                         vector2d(rect.max().y(), rect.max().z()));
            events.emplace_back(open, rect.min().x(), rect2);
            events.emplace_back(close, rect.max().x(), rect2);
        }
    }

    std::sort(events.begin(), events.end());
    return events;
}

double union_area(const std::vector<rect2d>& rects) {
    if (rects.size() == 0) {
        return 0;
    }

    segment_tree tree = build_segment_tree(rects);
    std::vector<event2d> events = rectangle_events(rects);

    double area = 0.0;
    double lastx = 0;
    for (const event2d& e : events) {
        // Works even for the first iteration because tree.union_width() will be 0.
        area += (e.x - lastx) * tree.union_width();
        if (e.type == open) {
            tree.insert(e.y);
        } else {
            tree.remove(e.y);
        }
        lastx = e.x;
    }
    return area;
}

double union_area(const std::vector<rect3d>& rects) {
    if (rects.size() == 0) {
        return 0;
    }

    std::vector<event3d> events = rectangle_events(rects);
    std::vector<rect2d> active;
    active.reserve(rects.size());

    double area = 0;
    double lastx = 0;
    for (const event3d& e : events) {
        area += (e.x - lastx) * union_area(active);
        if (e.type == open) {
            active.push_back(e.r);
        } else {
            auto pos = std::find(active.begin(), active.end(), e.r);
            geodb_assert(pos != active.end(), "rectangle must be active");
            active.erase(pos);
        }
        lastx = e.x;
    }

    return area;
}

} // namespace geodb