// Copyright (C) 2010 Deltares
//
// This program is free software; you can redistribute it and/or modify
// it under the terms of the GNU General Public License version 2 as
// published by the Free Software Foundation.
//
// This program is distributed in the hope that it will be useful,
// but WITHOUT ANY WARRANTY; without even the implied warranty of
// MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
// GNU General Public License for more details.
//
// You should have received a copy of the GNU General Public License
// along with this program; if not, write to the Free Software
// Foundation, Inc., 51 Franklin St, Fifth Floor, Boston, MA 02110-1301 USA

/**
 * @file
 * @brief xxx
 * @author Dirk Schwanenberg
 * @version 1.0
 * @date 2010
 */

#include "rtcToolsSimulator.h"
#include <utilities/utils.h>
#include <optimizationProblem/objectiveFunction.h>
#include "piDiagInterface.h"

const string RTCTOOLSSIMULATOR_CODE = "RTCTools.runtime.rtcToolsSimulator";

using namespace rtctools;


rtcToolsSimulator::rtcToolsSimulator(
	int iEnsemble,
	timeSeriesMatrixInterface* tsMatrix,
	schematisation *schema,
	objectiveFunction *obj,
	double p,
	rtcRuntimeConfigSettings* runtimeSettings) : iEnsemble(iEnsemble), schema(schema),
	    tsMatrix(tsMatrix), obj(obj), p(p), runtimeSettings(runtimeSettings)
{ }


double rtcToolsSimulator::simulate(int iStart, int iEnd, bool executeObjectiveFunction, int iTerm, double** JInc2DArray)
{
	try {
		// calls simulation stepwise
		for (int i=iStart; i<=iEnd; i++) {
			simulate(i);
		}
	} catch (exception &e) {
		piDiagInterface::addLine(1, "double rtcToolsSimulator::simulate(int iStart, int iEnd) - error during simulation - " + string(e.what()), RTCTOOLSSIMULATOR_CODE);
		throw;
	} catch (...) {
		piDiagInterface::addLine(1, "double rtcToolsSimulator::simulate(int iStart, int iEnd) - unknown error during simulation", RTCTOOLSSIMULATOR_CODE);
		throw;
	}

	// computes objective function value
	double J = 0.0;

	if (executeObjectiveFunction) {
		if (runtimeSettings->parallelInternalSim) {
			double* JArray = new double[obj->getNTerm()];
			if (iTerm<0) {
				std::runtime_error *e = NULL;
				#pragma omp parallel for schedule(dynamic)
				for (int i=0; i<obj->getNTerm(); i++) {
					try {
						JArray[i] = obj->getTerm(i)->evaluate(iStart, iEnd, ((timeSeriesMatrix*)tsMatrix)->getValueMatrix(), JInc2DArray, tsMatrix->getTimes(), p);
					} catch (std::runtime_error &re) {
						e = new std::runtime_error(re);
					}
				}
				if (e) {
					throw *e;
					delete e;
				}
				for (int i=0; i<obj->getNTerm(); i++) J += JArray[i];
			} else {
				J += obj->getTerm(iTerm)->evaluate(iStart, iEnd, ((timeSeriesMatrix*)tsMatrix)->getValueMatrix(), JInc2DArray, tsMatrix->getTimes(), p);
			}
			delete [] JArray;
		} else {
			if (iTerm<0) {
				for (int i=0; i<obj->getNTerm(); i++) {
					J += obj->getTerm(i)->evaluate(iStart, iEnd, ((timeSeriesMatrix*)tsMatrix)->getValueMatrix(), JInc2DArray, tsMatrix->getTimes(), p);
				}
			} else {
				J += obj->getTerm(iTerm)->evaluate(iStart, iEnd, ((timeSeriesMatrix*)tsMatrix)->getValueMatrix(), JInc2DArray, tsMatrix->getTimes(), p);
			}
		}

		// TODO move one level higher
		if (obj->getNTerm()>0 && JInc2DArray!=0) {
			objectiveFunction::OUTPUT iOutput = obj->getOutputIDs();
			double** vMat = ((timeSeriesMatrix*)tsMatrix)->getValueMatrix();

			if (iOutput.JInc>-1 && iOutput.JAcc>-1) {
				// objective function value at initial condition is zero by definition
				vMat[0][iOutput.JInc] = 0.0;
				vMat[0][iOutput.JAcc] = 0.0;
				for (int i=1; i<tsMatrix->getNTimeStep(); i++) {
					vMat[i][iOutput.JInc] = 0.0;
					vMat[i][iOutput.JAcc] = vMat[i-1][iOutput.JAcc];
					for (int j=0; j<obj->getNTerm(); j++) {
						vMat[i][iOutput.JInc] += JInc2DArray[i][5*j+4];
						vMat[i][iOutput.JAcc] += JInc2DArray[i][5*j+4];
					}
				}
			}
		}
	}

	return J;
}


void rtcToolsSimulator::simulate(int iStep)
{
	if ((iStep<1) || (iStep>=this->tsMatrix->getNTimeStep()))
		throw runtime_error("void rtcToolsSimulator::simulate(int iStep) - time step is not available in time series matrix");

	// t, dt
	long long t = tsMatrix->getTime(iStep);
	double dt = ((double)tsMatrix->getDT(iStep))/1000.0;

    // get vectors with old and new states
    double *stateOld = tsMatrix->getState(iStep-1);
    double *stateNew = tsMatrix->getState(iStep);

	// triggers
	int nTrigger = schema->getNTrigger();
    trigger **tr = schema->getTriggers();
	try {
		for (int i=0; i<nTrigger; i++) {
			// this deactivates all referenced rules in the triggers
			// note a rules is not deactivated if NOT referenced in any trigger
			tr[i]->deactivate();
		}
		for (int i=0; i<nTrigger; i++) {
			tr[i]->solve(stateOld, stateNew, t, dt);
		}
	} catch (exception &e) {
		piDiagInterface::addLine(1, "void rtcToolsSimulator::simulate(int iStep) - error during trigger execution - " + string(e.what()), RTCTOOLSSIMULATOR_CODE);
		throw;
	} catch (...) {
		piDiagInterface::addLine(1, "void rtcToolsSimulator::simulate(int iStep) - unknown error during trigger execution", RTCTOOLSSIMULATOR_CODE);
		throw;
	}


	// rules
	int nRule = schema->getNRule();
    rule **ru = schema->getRules();
	try {
		for (int i=0; i<nRule; i++) {
			if (ru[i]->isActive()) {
				ru[i]->solve(stateOld, stateNew, t, dt);
			}
		}
	} catch (exception &e) {
		piDiagInterface::addLine(1, "void rtcToolsSimulator::simulate(int iStep) - error during rule execution - " + string(e.what()), RTCTOOLSSIMULATOR_CODE);
		throw;
	} catch (...) {
		piDiagInterface::addLine(1, "void rtcToolsSimulator::simulate(int iStep) - unknown error during rule execution", RTCTOOLSSIMULATOR_CODE);
		throw;
	}

	// components
	int nComponent = schema->getNComponent();
    vector<component*> co = schema->getComponents();
	try {
		for (int i=0; i<nComponent; i++) {
			set<int> eSet = co[i]->getESet();
			if (eSet.size()==0 || eSet.find(this->iEnsemble)!=eSet.end()) {
				co[i]->solve(stateOld, stateNew, t, dt);
			}
		}
	} catch (exception &e) {
		piDiagInterface::addLine(1, "void rtcToolsSimulator::simulate(int iStep) - error during component execution - " + string(e.what()), RTCTOOLSSIMULATOR_CODE);
		throw;
	} catch (...) {
		piDiagInterface::addLine(1, "void rtcToolsSimulator::simulate(int iStep) - unknown error during component execution", RTCTOOLSSIMULATOR_CODE);
		throw;
	}

	// increments the temporary states in limited memory option except for the last one
	if (runtimeSettings->limitedMemory && (iStep<tsMatrix->getNTimeStep()-1)) {
		tsMatrix->incrementTimeStep();
	}
}


void rtcToolsSimulator::evaluateGradient(int iStart, int iEnd, int iTerm)
{
    int nComponent = schema->getNComponent();
    vector<component*> co = schema->getComponents();

	int nRule = schema->getNRule();
    rule **ru = schema->getRules();

	int nTrigger = schema->getNTrigger();
    trigger **tr = schema->getTriggers();

    // compute derivatives of objective functions
	if (iTerm<0) {
		// all terms
		for (int i=0; i<obj->getNTerm(); i++) {
			obj->getTerm(i)->evaluateDer(iStart, iEnd, ((timeSeriesMatrix*)tsMatrix)->getValueMatrix(), ((timeSeriesMatrix*)tsMatrix)->getObjMatrix(), tsMatrix->getTimes(), p);
		}
	} else {
		// only the term specified
		obj->getTerm(iTerm)->evaluateDer(iStart, iEnd, ((timeSeriesMatrix*)tsMatrix)->getValueMatrix(), ((timeSeriesMatrix*)tsMatrix)->getObjMatrix(), tsMatrix->getTimes(), p);
	}

    // reverse simulation for lambda and gradient
    for (int iStep=iEnd-1; iStep>=iStart-1; iStep--) {

        // t, dt
        long long t = tsMatrix->getTime(iStep);
        double dt = ((double)tsMatrix->getDT(iStep))/1000.0;

        // get vectors with old and new states
        double *stateOld = tsMatrix->getState(iStep);
        double *stateNew = tsMatrix->getState(iStep+1);
		double *dStateOld = ((timeSeriesMatrix*)tsMatrix)->getStateObj(iStep);
        double *dStateNew = ((timeSeriesMatrix*)tsMatrix)->getStateObj(iStep+1);

        // reverse component loop
        for (int i=nComponent-1; i>=0; i--) {
			set<int> eSet = co[i]->getESet();
			if (eSet.size()==0 || eSet.find(this->iEnsemble)!=eSet.end()) {
				co[i]->solveDer(stateOld, stateNew, t, dt, dStateOld, dStateNew);
			}
        }

		// triggers in foward order to reproduce rule activation
		for (int i=0; i<nTrigger; i++) {
			tr[i]->deactivate();
		}
		for (int i=0; i<nTrigger; i++) {
			tr[i]->solve(stateOld, stateNew, t, dt);
		}

		// reverse rule loop
		for (int i=nRule-1; i>=0; i--) {
			if (ru[i]->isActive()) {
				ru[i]->solveDer(stateOld, stateNew, t, dt, dStateOld, dStateNew);
			}
		}
    }
}


void rtcToolsSimulator::eval_g(int n, const double* x, int m, double* g)
{
	int nTimeStep = tsMatrix->getNTimeStep();
	double** tsm = ((timeSeriesMatrix*)tsMatrix)->getValueMatrix();

	// rate of change constraints on variables
	//#pragma omp parallel for schedule(dynamic)
	for (int i=0; i<obj->getNVariableRateOfChange(); i++) {
		objectiveFunction::variableRateOfChange* b = obj->getVariableRateOfChange(i);
		for (int j=0; j<(int)b->rof2DArray[iEnsemble].size(); j++) {
			rateOfChangeConstraint& rof = b->rof2DArray[iEnsemble][j];
			double res = 0.0;
			for (int k=0; k<(int)rof.iVariableTSM.size(); k++) {
				res += rof.factor[k] * tsm[rof.iStep[k]][rof.iVariableTSM[k]];
			}
			if (res!=res) {
				stringstream ss;
				ss << "constraint " << b->id << " includes NaN at time step " << j;
				piDiagInterface::addLine(1, ss.str(), RTCTOOLSSIMULATOR_CODE);
				throw runtime_error(ss.str().c_str());
			}
			g[rof.ncIndx] = res;
		}
	}

	// average constraints on variables
	//#pragma omp parallel for schedule(dynamic)
	for (int i=0; i<obj->getNVariableAverage(); i++) {
		objectiveFunction::variableAverage* b = obj->getVariableAverage(i);
		for (int j=0; j<(int)b->av2DArray[iEnsemble].size(); j++) {
			averageConstraint& av = b->av2DArray[iEnsemble][j];
			double res = 0.0;
			for (int k=0; k<(int)av.iVariableTSM.size(); k++) {
				res += av.factor[k] * tsm[av.iStep[k]][av.iVariableTSM[k]];
			}
			if (res!=res) {
				stringstream ss;
				ss << "constraint " << b->id << " includes NaN at time step " << j;
				piDiagInterface::addLine(1, ss.str(), RTCTOOLSSIMULATOR_CODE);
				throw runtime_error(ss.str().c_str());
			}
			g[av.ncIndx] = res;
		}
	}

	// equality constraints on variables
	//#pragma omp parallel for schedule(dynamic)
	for (int i=0; i<obj->getNVariableEqual(); i++) {
		objectiveFunction::variableEqual* b = obj->getVariableEqual(i);
		for (int j=0; j<(int)b->eq2DArray[iEnsemble].size(); j++) {
			equalConstraint& eq = b->eq2DArray[iEnsemble][j];
			double res = 0.0;
			for (int k=0; k<(int)eq.iVariableTSM.size(); k++) {
				res += eq.factor[k] * tsm[eq.iStep[k]][eq.iVariableTSM[k]];
			}
			if (res!=res) {
				stringstream ss;
				ss << "constraint " << b->id << " includes NaN at time step " << j;
				piDiagInterface::addLine(1, ss.str(), RTCTOOLSSIMULATOR_CODE);
				throw runtime_error(ss.str().c_str());
			}
			g[eq.ncIndx] = res;
		}
	}

	// bound constraints on states
	//#pragma omp parallel for schedule(dynamic)
	for (int i=0; i<obj->getNState(); i++) {

		objectiveFunction::state* b = obj->getState(i);

		for (int j=0; j<(int)b->state2DArray[iEnsemble].size(); j++) {

			stateConstraint& sc = b->state2DArray[iEnsemble][j];

			double v = tsm[sc.iStepResiduum][sc.iResiduum]; 
			if (v!=v) {
				stringstream ss;
				ss << "state constraint " << b->id << " includes NaN at time step " << j;
				piDiagInterface::addLine(1, ss.str(), RTCTOOLSSIMULATOR_CODE);
				throw runtime_error(ss.str().c_str());
			}

			g[sc.ncIndx] = sc.scalingFactor * v;
		}
	}
}


void rtcToolsSimulator::eval_jac_g(int m, int nnz, int* iRow, int* jCol, double* values)
{
	int nTimeStep = tsMatrix->getNTimeStep();
	double dt = ((double)tsMatrix->getDT())/1000.0;
	double** tsm = ((timeSeriesMatrix*)tsMatrix)->getValueMatrix();

	/* build up the structure of the lower triangular matrix for each variable */
	if (iRow && jCol) {

		// loop variable rate of change bounds
		for (int i=0; i<obj->getNVariableRateOfChange(); i++) {
			objectiveFunction::variableRateOfChange* c = obj->getVariableRateOfChange(i);
			for (int j=0; j<(int)c->rof2DArray[iEnsemble].size(); j++) {				
				rateOfChangeConstraint& rof = c->rof2DArray[iEnsemble][j];
				for (int k=0; k<(int)rof.iRow.size(); k++) {
					iRow[rof.nnzIndx[k]] = rof.iRow[k];
					jCol[rof.nnzIndx[k]] = rof.jCol[k];
				}
			}
		}

		// loop variable average bounds
		for (int i=0; i<obj->getNVariableAverage(); i++) {
			objectiveFunction::variableAverage* c = obj->getVariableAverage(i);
			for (int j=0; j<(int)c->av2DArray[iEnsemble].size(); j++) {				
				averageConstraint& av = c->av2DArray[iEnsemble][j];
				for (int k=0; k<(int)av.iRow.size(); k++) {
					iRow[av.nnzIndx[k]] = av.iRow[k];
					jCol[av.nnzIndx[k]] = av.jCol[k];
				}
			}
		}

		// loop variable equalities
		for (int i=0; i<obj->getNVariableEqual(); i++) {
			objectiveFunction::variableEqual* c = obj->getVariableEqual(i);
			for (int j=0; j<(int)c->eq2DArray[iEnsemble].size(); j++) {				
				equalConstraint& eq = c->eq2DArray[iEnsemble][j];
				for (int k=0; k<(int)eq.iRow.size(); k++) {
					iRow[eq.nnzIndx[k]] = eq.iRow[k];
					jCol[eq.nnzIndx[k]] = eq.jCol[k];
				}
			}
		}

		// loop state bounds
		for (int i=0; i<obj->getNState(); i++) {
			objectiveFunction::state* c = obj->getState(i);
			for (int j=0; j<(int)c->state2DArray[iEnsemble].size(); j++) {
				stateConstraint& sc = c->state2DArray[iEnsemble][j];
				for (int k=0; k<(int)sc.iRow.size(); k++) {
					iRow[sc.nnzIndx[k]] = sc.iRow[k];
					jCol[sc.nnzIndx[k]] = sc.jCol[k];
				}
			}
		}
	}

	if (values) {

		// loop variable rate of change bounds
		if (runtimeSettings->parallelInternalCon) {
			#pragma omp parallel for schedule(dynamic)
			for (int i=0; i<(int)obj->getNVariableRateOfChange(); i++) {
				eval_jac_g_variableRateOfChange(m, nnz, iRow, jCol, values, obj->getVariableRateOfChange(i), dt);
			}
		} else {
			for (int i=0; i<(int)obj->getNVariableRateOfChange(); i++) {
				eval_jac_g_variableRateOfChange(m, nnz, iRow, jCol, values, obj->getVariableRateOfChange(i), dt);
			}
		}

		// loop variable average bounds
		for (int i=0; i<obj->getNVariableAverage(); i++) {
			objectiveFunction::variableAverage* c = obj->getVariableAverage(i);
			for (int j=0; j<(int)c->av2DArray[iEnsemble].size(); j++) {
				averageConstraint& av = c->av2DArray[iEnsemble][j];
				for (int k=0; k<(int)av.iRow.size(); k++) {
					values[av.nnzIndx[k]] = av.value[k];
				}
			}
		}

		// loop variable equalities
		for (int i=0; i<obj->getNVariableEqual(); i++) {
			objectiveFunction::variableEqual* c = obj->getVariableEqual(i);
			for (int j=0; j<(int)c->eq2DArray[iEnsemble].size(); j++) {
				equalConstraint& eq = c->eq2DArray[iEnsemble][j];
				for (int k=0; k<(int)eq.iRow.size(); k++) {
					values[eq.nnzIndx[k]] = eq.value[k];
				}
			}
		}

		// loop state bounds
		if (runtimeSettings->parallelInternalCon) {
			#pragma omp parallel for schedule(dynamic)
			for (int i=0; i<(int)obj->getNState(); i++) {
				eval_jac_g_State(m, nnz, iRow, jCol, values, obj->getState(i));
			}
		} else {
			for (int i=0; i<(int)obj->getNState(); i++) {
				eval_jac_g_State(m, nnz, iRow, jCol, values, obj->getState(i));
			}
		}
	}
}

void rtcToolsSimulator::eval_jac_g_variableRateOfChange(int m, int nnz, int* iRow, int* jCol, 
	double* values, objectiveFunction::variableRateOfChange* c, double dt)
{
	objectiveFunction::variable* v = obj->getVariable(c->iVariableOPT);
	double sFac = v->scalingFactor;

	double tCorr = 1.0;
	if (c->dtRef>0) {
		tCorr = (double)c->dtRef/(double)tsMatrix->getDT();
	}

	for (int j=0; j<(int)c->rof2DArray[iEnsemble].size(); j++) {
				
		rateOfChangeConstraint& rof = c->rof2DArray[iEnsemble][j];

		if (rof.nnzIndx.size()==1) {
			// first value is a state and not an optimization variable
			// in this case, there is only a single nonzero entry for the second value of the difference
			values[rof.nnzIndx[0]] = +tCorr*rof.correction/sFac;
		} else {
			values[rof.nnzIndx[0]] = -tCorr*rof.correction/sFac;
			values[rof.nnzIndx[1]] = +tCorr*rof.correction/sFac;
		}
	}
}

void rtcToolsSimulator::eval_jac_g_State(int m, int nnz, int* iRow, int* jCol, 
	double* values, objectiveFunction::state* c)
{
	int nVar = (int)c->iVariableTSM.size();
	int nStep = c->nTimeStep;

	// use pre-allocated adjoint matrix
	double** dState = c->dState[iEnsemble];

	for (int j=0; j<(int)c->state2DArray[iEnsemble].size(); j++) {

		stateConstraint& sc = c->state2DArray[iEnsemble][j];

		// initialize adjoint matrix with zero
		memset(dState[0], 0, sizeof(double)*(nStep+1)*tsMatrix->getNSeries());
		dState[nStep][c->iState] = 1.0;

		// trace back the adjoints backwards in time
		for (int k=sc.iEnd; k>sc.iStart; k--) {
			if (k < 1)
				break;

			double* stateOld = tsMatrix->getState(k-1);
			double* stateNew = tsMatrix->getState(k);
			long long t = tsMatrix->getTime(k);
			double dt = ((double)tsMatrix->getDT(k))/1000.0;
			double* dStateOld = dState[k-sc.iStart-1];
			double* dStateNew = dState[k-sc.iStart];

			// simulation components
			for (int l=(int)c->c.size()-1; l>=0; l--) {
				c->c[l]->solveDer(stateOld, stateNew, t, dt, dStateOld, dStateNew);
			}
		}

		// extract contribution to Jacobian
		for (int k=0; k<(int)sc.iVariableOPT.size(); k++) {
			double sFac = obj->getVariable(sc.iVariableOPT[k])->scalingFactor;
			values[sc.nnzIndx[k]] = sc.scalingFactor * dState[sc.iStepVariable[k]-sc.iStart][sc.iVariableTSM[k]]/sFac;
		}
	}
}