/* $Revision: 1.3 $ $Date: 2003/03/31 21:38:05 $ */
/*===============================================================================================================
 *   KALCVF The Kalman filter
 *
 *   State space model is defined as follows:
 *     z(t+1) = a+F*z(t)+eta(t)     (state or transition equation)
 *       y(t) = b+H*z(t)+eps(t)     (observation or measurement equation)
 *
 *   [logl, <pred, vpred, <filt, vfilt>>] = kalcvf(data, lead, a, F, b, H, var, <z0, vz0>)
 *   computes the one-step prediction and the filtered estimate, as well as their covariance matrices.
 *   The function uses forward recursions, and you can also use it to obtain k-step estimates.
 *
 *   The inputs to the KALCVF function are as follows:
 *     data is a Ny×T matrix containing data (y(1), ... , y(T)).
 *     lead is the number of steps to forecast after the end of the data.
 *        a is an Nz×1 vector for a time-invariant input vector in the transition equation,
 *             or a Nz×(T+lead) vector containing input vectors in the transition equation.
 *        F is an Nz×Nz matrix for a time-invariant transition matrix in the transition equation,
 *             or a Nz×Nz×(T+lead) matrix containing transition matrices in the transition equation.
 *        b is an Ny×1 vector for a time-invariant input vector in the measurement equation,
 *             or a Ny×(T+lead) vector containing input vectors in the measurement equation.
 *        H is an Ny×Nz matrix for a time-invariant measurement matrix in the measurement equation,
 *             or a Ny×Nz×(T+lead) matrix containing measurement matrices in the measurement equation.
 *      var is an (Ny+Nz)×(Ny+Nz) matrix for a time-invariant variance matrix for
 *             the error in the transition equation and the error in the measurement equation,
 *             or a (Ny+Nz)×(Ny+Nz)×(T+lead) matrix containing variance matrices for
 *             the error in the transition equation and the error in the measurement equation,
 *             that is, [eta(t)', eps(t)']'.
 *       z0 is an optional Nz×1 initial state vector.
 *      vz0 is an optional Nz×Nz covariance matrix of an initial state vector.
 *
 *   The KALCVF function returns the following output:
 *     logl is a value of the average log likelihood function of the SSM
 *             under assumption that observation noise eps(t) is normally distributed
 *     pred is an optional Nz×(T+lead) matrix containing one-step predicted state vectors.
 *    vpred is an optional Nz×Nz×(T+lead) matrix containing mean square errors of predicted state vectors.
 *     filt is an optional Nz×T matrix containing filtered state vectors.
 *    vfilt is an optional Nz×Nz×T matrix containing mean square errors of filtered state vectors.
 *
 *   The initial state vector and its covariance matrix of the time invariant Kalman filters
 *   are computed under the stationarity condition:
 *          z0 = (I-F)\a
 *         vz0 = (I-kron(F,F))\V(:)
 *   where F and V are the time invariant transition matrix and the covariance matrix of transition equation noise,
 *   and vec(V) is an Nz^2×1 column vector that is constructed by the stacking Nz columns of matrix V.
 *   Note that all eigenvalues of the matrix F are inside the unit circle when the SSM is stationary.
 *   When the preceding formula cannot be applied, the initial state vector estimate is set to a
 *   and its covariance matrix is given by 1E6I. Optionally, you can specify initial values.
 *
 *   This is a MEX-file for MATLAB.
 *   Copyright 2002-2003 Federal Reserve Bank of Atlanta
 *   Iskander Karibzhanov 5-28-02.
 *   Master of Science in Computational Finance
 *   Georgia Institute of Technology
 *===============================================================================================================
 * Revision history:
 *
 *  03/19/2003  -  algorithm and interface were adapted from SAS/IML KALCVF subroutine for use in MATLAB MEX file
 *
 *===============================================================================================================*/

#include "matlib.h"

/* all warning messages were supressed because in case of failure in dlyap
   there exist default solution in the main algorithm */
#define _DLYAP_NOMSG_

#ifdef KALCVF_MEX
#undef KALCVF_MEX
void mexFunction(int nlhs, mxArray *plhs[], int nrhs, const mxArray *prhs[]) {

	double *data, *a, *F, *b, *H, *var, *z0, *vz0, *logl, *pred, *vpred, *filt, *vfilt;
	int T, Nz, Ny, lead, rows, cols, ndims, inca=0, incF=0, incb=0, incH=0, incvar=0;
	const int *dims;

	/*===== check number of input and output arguments =====*/

	if ( (nrhs != 7) && (nrhs != 9) )
		mexErrMsgTxt("Seven or nine input arguments required.");
	if ( (nlhs > 5) )
		mexErrMsgTxt("More than five output arguments specified.");

	/*===== check number of forecasting steps =====*/
	lead = (int)mxGetScalar(prhs[1]);
	if ( (mxGetM(prhs[1])*mxGetN(prhs[1]) != 1) || (lead < 0) )
		mexErrMsgTxt("lead must be positive integer scalar.");

	/*===== check input matrix dimensions =====*/

	/* data must be Ny×T matrix */
	Ny = mxGetM(prhs[0]);
	T  = mxGetN(prhs[0]);

	/* a must be Nz×1 or Nz×(T+lead) vector */
	Nz = mxGetM(prhs[2]);
	if ( (cols=mxGetN(prhs[2])) == 1 )
		inca = 0;
	else if ( cols == (T+lead) )
		inca = Nz;
	else {
		sprintf(msg, "data has Ny=%d rows and T=%d columns, lead is equal to %d, "
			"a has Nz=%d rows, but the number of columns in a (%d) is neither one "
			"nor T+lead=%d.", Ny, T, lead, Nz, cols, T+lead);
		mexErrMsgTxt(msg);
	}

	/* F must be Nz×Nz or Nz×Nz×(T+lead) matrix */
	dims = mxGetDimensions(prhs[3]);
	if ( (dims[0] != Nz) || (dims[1] != Nz) ) {
		sprintf(msg, "a and F must have the same number of rows. F must be square in the first two dimensions. "
			"a has Nz=%d rows, but F has %d rows and %d columns.", Nz, dims[0], dims[1]);
		mexErrMsgTxt(msg);
	}
	ndims = mxGetNumberOfDimensions(prhs[3]);
	if ( ndims == 2 )
		incF = 0;
	else if ( ndims == 3 ) {
		if ( dims[2] == (T+lead) )
			incF = Nz*Nz;
		else {
			sprintf(msg, "F must have (T+lead)=%d elements in third dimension.", T+lead);
			mexErrMsgTxt(msg);
		}
	} else
		mexErrMsgTxt("F must be two- or three-dimensional matrix.");

	/* b must be Ny×1 or Ny×(T+lead) vector */
	if ( (rows=mxGetM(prhs[4])) != Ny ) {
		sprintf(msg, "data and b must have the same number of rows. "
			"data has Ny=%d rows, but the number of rows in b is %d.", Ny, rows);
		mexErrMsgTxt(msg);
	}
	if ( (cols=mxGetN(prhs[4])) == 1 )
		incb = 0;
	else if ( cols == (T+lead) )
		incb = Ny;
	else {
		sprintf(msg, "data has Ny=%d rows and T=%d columns, lead is equal to %d, "
			"b has Ny=%d rows, but the number of columns in b (%d) is neither one "
			"nor T+lead=%d.", Ny, T, lead, Ny, cols, T+lead);
		mexErrMsgTxt(msg);
	}

	/* H must be Ny×Nz or Ny×Nz×(T+lead) matrix */
	dims = mxGetDimensions(prhs[5]);
	if ( dims[0] != Ny ) {
		sprintf(msg, "data and H must have the same number of rows. "
			"data has Ny=%d rows, but the number of rows in H is %d.", Ny, dims[0]);
		mexErrMsgTxt(msg);
	}
	if ( dims[1] != Nz ) {
		sprintf(msg, "H must have the same number of columns as rows in matrix a. "
			"a has Nz=%d rows, but the number of columns in H is %d.", Nz, dims[1]);
		mexErrMsgTxt(msg);
	}
	ndims = mxGetNumberOfDimensions(prhs[5]);
	if ( ndims == 2 )
		incH = 0;
	else if ( ndims == 3 ) {
		if ( dims[2] == (T+lead) )
			incH = Ny*Nz;
		else {
			sprintf(msg, "H must have (T+lead)=%d elements in third dimension.", T+lead);
			mexErrMsgTxt(msg);
		}
	} else
		mexErrMsgTxt("H must be two- or three-dimensional matrix.");

	/* var must be (Ny+Nz)×(Ny+Nz) or (Ny+Nz)×(Ny+Nz)×(T+lead) matrix */
	dims = mxGetDimensions(prhs[6]);
	if ( (dims[0] != Ny+Nz) || (dims[1] != Ny+Nz) ) {
		sprintf(msg, "var must contain variance matrix for the errors in transition and measurement equations. "
			"var must be square p.d.f. %d×%d or %d×%d×%d matrix, but your var has %d rows and %d columns.",
			Ny+Nz, Ny+Nz, Ny+Nz, Ny+Nz, T+lead, dims[0], dims[1]);
		mexErrMsgTxt(msg);
	}
	ndims = mxGetNumberOfDimensions(prhs[6]);
	if ( ndims == 2 )
		incvar = 0;
	else if ( ndims == 3 ) {
		if ( dims[2] == (T+lead) )
			incvar = (Ny+Nz)*(Ny+Nz);
		else {
			sprintf(msg, "var must have (T+lead)=%d elements in third dimension.", T+lead);
			mexErrMsgTxt(msg);
		}
	} else
		mexErrMsgTxt("var must be two- or three-dimensional matrix.");

	/* if specified, z0 must be Nz×1 vector and vz0 must be Nz×Nz matrix */
	if ( nrhs==9 ) {
		rows = mxGetM(prhs[7]);
		cols = mxGetN(prhs[7]);
		if ( ( rows != Nz ) || ( cols != 1 ) ) {
			sprintf(msg, "a and z0 must have the same number of rows. z0 must be Nz×1 vector. "
				"a has Nz=%d rows, but z0 has %d rows and %d columns.", Nz, rows, cols);
			mexErrMsgTxt(msg);
		}
		rows=mxGetM(prhs[8]);
		cols=mxGetN(prhs[8]);
		if ( ( rows != Nz ) || ( cols != Nz ) ) {
			sprintf(msg, "a and vz0 must have the same number of rows. vz0 must be square Nz×Nz matrix. "
				"a has Nz=%d rows, but vz0 has %d rows and %d columns.", Nz, rows, cols);
			mexErrMsgTxt(msg);
		}
	}

	/*===== get pointers to input arguments =====*/

	data = mxGetPr(prhs[0]);
	a    = mxGetPr(prhs[2]);
	F    = mxGetPr(prhs[3]);
	b    = mxGetPr(prhs[4]);
	H    = mxGetPr(prhs[5]);
	var  = mxGetPr(prhs[6]);
	if ( nrhs == 9 ) {
		z0 = mxGetPr(prhs[7]);
		vz0 = mxGetPr(prhs[8]);
	} else {
		z0  = NULL;
		vz0 = NULL;
	}

	/*===== create output scalar value of the average log likelihood function =====*/

	logl = mxGetPr(plhs[0] = mxCreateDoubleMatrix(1, 1, mxREAL));

	/*===== get pointers to optional output arguments =====*/

	if ( nlhs < 5 ) vfilt = NULL;
	else {
		int dim[3] = { Nz, Nz, T };
		vfilt = mxGetPr(plhs[4]=mxCreateNumericArray(3, dim, mxDOUBLE_CLASS, mxREAL));
	}
	if ( nlhs < 4 ) filt = NULL;
	else filt  = mxGetPr(plhs[3]=mxCreateDoubleMatrix(Nz, T, mxREAL));
	if ( nlhs < 3 ) vpred = NULL;
	else {
		int dim[3] = { Nz, Nz, T+lead };
		vpred = mxGetPr(plhs[2]=mxCreateNumericArray(3, dim, mxDOUBLE_CLASS, mxREAL));
	}
	if ( nlhs < 2 ) pred = NULL;
	else pred  = mxGetPr(plhs[1]=mxCreateDoubleMatrix(Nz, T+lead, mxREAL));

	/*===== compute Kalman Filter =====*/

	*logl = kalcvf(data, T, lead, Ny, Nz, a, inca, F, incF, b, incb, H, incH, var, incvar, z0, vz0, pred, vpred, filt, vfilt);

}
#endif

/* The Kalman filter */
#define PI 3.141592653589793238
double kalcvf(double *data, int T, int lead, int Ny, int Nz,
              double *a_all, int inca, double *F_all, int incF,
              double *b_all, int incb, double *H_all, int incH,
              double *var_all, int incvar, double *z0, double *vz0,
              double *pred, double *vpred, double *filt, double *vfilt) {

	double *a, *F, *b, *H, *var, *V, *G, *R, *P;
	double logl=0.0, prod;
	int Nyz, NyzNyz, NyNz, NzNz, Nz1, incpred=0, incvpred=0, info=0, t, i;

	Nyz = Ny+Nz;
	NyzNyz = Nyz*Nyz;
	NyNz = Ny*Nz;
	NzNz = Nz*Nz;
	Nz1 = Nz+1;

	/* If specified, pointers pred and vpred will be incremented by incpred and incvpred
	   to save data from each iteration. If not specified, allocate memory for vector pred
	   and matrix vpred, and leave zero increments */
	if (pred) incpred = Nz;
	else pred = (double *)mxCalloc(Nz, sizeof(double));
	if (vpred) incvpred = NzNz;
	else vpred = (double *)mxCalloc(NzNz, sizeof(double));

	/* initialize pred and vpred */
	if (z0)
		cblas_dcopy(Nz, z0, one, pred, one);
	else {
		cblas_dcopy(Nz, a_all, one, pred, one);
		dlins(Nz, one, F_all, pred, &info);
	}
	if (vz0)
		cblas_dcopy(NzNz, vz0, one, vpred, one);
	else {
		dlyap(Nz, F_all, var_all, Nyz, vpred, &info);
		if ( info > 0 ) {
			double million=1e6;
			cblas_dcopy(NzNz, &zero_d, zero, vpred, one);
			cblas_daxpy(Nz, one_d, &million, zero, vpred, Nz1);
		}
	}

	/* allocate memory for temporary variables */
	a   = (double *) mxCalloc(Nz,     sizeof(double));
	F   = (double *) mxCalloc(NzNz,   sizeof(double));
	b   = (double *) mxCalloc(Ny,     sizeof(double));
	H   = (double *) mxCalloc(NyNz,   sizeof(double));
	var = (double *) mxCalloc(NyzNyz, sizeof(double));
	P   = (double *) mxCalloc(NzNz,   sizeof(double));
	/* var = [V G; G' R] */
	V = var;          /* V(t) = Var(eta(t)) */
	G = var+Nyz*Nz;   /* G(t) = Cov(eta(t),eps(t)) */
	R = G+Nz;         /* R(t) = Var(eps(t)) */

	for (t=0; t<T; t++) {
		/* H = H_all(:,:,t) */
		cblas_dcopy(NyNz, &H_all[t*incH], one, H, one);
		/* b = b_all(:,t) */
		cblas_dcopy(Ny, &b_all[t*incb], one, b, one);
		/* b = -H*pred-b */
		cblas_dgemv(CblasColMajor, CblasNoTrans, Ny, Nz, mone_d, H, Ny, pred, one, mone_d, b, one);
		/* b = data(:,t)+b */
		cblas_daxpy(Ny, one_d, &data[t*Ny], one, b, one);
      /* P = vpred */
		cblas_dcopy(NzNz, vpred, one, P, one);
		/* P = chol(P)' */
      dpotrf(chl, &Nz, P, &Nz, &info);
		/* H = H*P */
		cblas_dtrmm(CblasColMajor, CblasRight, CblasLower, CblasNoTrans, CblasNonUnit, Ny, Nz, one_d, P, Nz, H, Ny);
		/* var = var_all(:,:,t) */
		cblas_dcopy(NyzNyz, &var_all[t*incvar], one, var, one);
		/* R = H*H'+R */
		cblas_dsyrk(CblasColMajor, CblasLower, CblasNoTrans, Ny, Nz, one_d, H, Ny, one_d, R, Nyz);
		/* R = chol(R)' */
		dpotrf(chl, &Ny, R, &Nyz, &info);
		/* b = R\b */
		cblas_dtrsv(CblasColMajor, CblasLower, CblasNoTrans, CblasNonUnit, Ny, R, Nyz, b, one);
		/* logl += 2*log(prod(diag(R)))+b'*b */
		prod = R[0];
		for (i=1; i<Ny; i++)
			prod *= R[i*(Nyz+1)];
		logl += 2*log(prod)+cblas_ddot(Ny, b, one, b, one);
		if (t<T-1 || lead>0) {
			/* F = F_all(:,:,t) */
			cblas_dcopy(NzNz, &F_all[t*incF], one, F, one);
			/* a = a_all(:,t) */
			cblas_dcopy(Nz, &a_all[t*inca], one, a, one);
			/* a = F*pred+a */
			cblas_dgemv(CblasColMajor, CblasNoTrans, Nz, Nz, one_d, F, Nz, pred, one, one_d, a, one);
			/* F = F*P */
			cblas_dtrmm(CblasColMajor, CblasRight, CblasLower, CblasNoTrans, CblasNonUnit, Nz, Nz, one_d, P, Nz, F, Nz);
			/* G = F*H'+G */
			cblas_dgemm(CblasColMajor, CblasNoTrans, CblasTrans, Nz, Ny, Nz, one_d, F, Nz, H, Ny, one_d, G, Nyz);
			/* G = G/R' */
			cblas_dtrsm(CblasColMajor, CblasRight, CblasLower, CblasTrans, CblasNonUnit, Nz, Ny, one_d, R, Nyz, G, Nyz);
			/* a = G*b+a */
			cblas_dgemv(CblasColMajor, CblasNoTrans, Nz, Ny, one_d, G, Nyz, b, one, one_d, a, one);
			/* V = F*F'+V */
			cblas_dsyrk(CblasColMajor, CblasLower, CblasNoTrans, Nz, Nz, one_d, F, Nz, one_d, V, Nyz);
			/* V = -G*G'+V */
			cblas_dsyrk(CblasColMajor, CblasLower, CblasNoTrans, Nz, Ny, mone_d, G, Nyz, one_d, V, Nyz);
		}
		if (filt || vfilt) {
			/* H = H*P' */
			cblas_dtrmm(CblasColMajor, CblasRight, CblasLower, CblasTrans, CblasNonUnit, Ny, Nz, one_d, P, Nz, H, Ny);
			/* H = R\H */
			cblas_dtrsm(CblasColMajor, CblasLeft, CblasLower, CblasNoTrans, CblasNonUnit, Ny, Nz, one_d, R, Nyz, H, Ny);
			if (filt) {
				/* filt = pred */
				cblas_dcopy(Nz, pred, one, filt, one);
				/* filt = H'*b+filt */
				cblas_dgemv(CblasColMajor, CblasTrans, Ny, Nz, one_d, H, Ny, b, one, one_d, filt, one);
				/* increment filt */
				filt += Nz;
			}
			if (vfilt) {
				/* vfilt = vpread */
				cblas_dcopy(NzNz, vpred, one, vfilt, one);
				/* vfilt = -H'*H+vfilt */
				cblas_dsyrk(CblasColMajor, CblasLower, CblasTrans, Nz, Ny, mone_d, H, Ny, one_d, vfilt, Nz);
				/* vfilt = tril(vfilt)+tril(vfilt,-1)' */
				dtril(Nz, vfilt, Nz);
				/* increment vfilt */
				vfilt += NzNz;
			}
		}
		if (t<T-1 || lead>0) {
			/* pred(:,t+1) = a */
			if (incpred) pred += incpred;
			cblas_dcopy(Nz, a, one, pred, one);
			/* vpred(:,:,t+1) = tril(V)+tril(V,-1)' */
			if (incvpred) vpred += incvpred;
			for (i=0; i<Nz; i++)
				cblas_dcopy(Nz, &V[i], Nyz, &vpred[i], Nz);
			dtril(Nz, vpred, Nz);
		}
	}
	if (lead>1 && (incpred || incvpred)) {
		for (t=T+1; t<T+lead; t++) {
			/* F = F_all(:,:,t) */
			cblas_dcopy(NzNz, &F_all[t*incF], one, F, one);
			/* a = a_all(:,t) */
			cblas_dcopy(Nz, &a_all[t*inca], one, a, one);
			/* a = F*pred+a */
			cblas_dgemv(CblasColMajor, CblasNoTrans, Nz, Nz, one_d, F, Nz, pred, one, one_d, a, one);
			/* pred(:,t) = a */
			if (incpred) pred += incpred;
			cblas_dcopy(Nz, a, one, pred, one);
			/* P = vpred */
			cblas_dcopy(NzNz, vpred, one, P, one);
			/* P = chol(P)' */
			dpotrf(chl, &Nz, P, &Nz, &info);
			/* F = F*P */
			cblas_dtrmm(CblasColMajor, CblasRight, CblasLower, CblasNoTrans, CblasNonUnit, Nz, Nz, one_d, P, Nz, F, Nz);
			/* var = var_all(:,:,t) */
			cblas_dcopy(NyzNyz, &var_all[t*incvar], one, var, one);
			/* vpred = V */
			if (incvpred) vpred += incvpred;
			for (i=0; i<Nz; i++)
				cblas_dcopy(Nz, &V[i], Nyz, &vpred[i], Nz);
			/* vpred = F*F'+vpred */
			cblas_dsyrk(CblasColMajor, CblasLower, CblasNoTrans, Nz, Nz, one_d, F, Nz, one_d, vpred, Nz);
			dtril(Nz, vpred, Nz);
		}
	}

	logl = -(Ny*log(2*PI)+logl/T)/2;

	if (!incpred) mxFree(pred);
	if (!incvpred) mxFree(vpred);
	mxFree(a);
	mxFree(F);
	mxFree(b);
	mxFree(H);
	mxFree(var);
	mxFree(P);
	return (logl);
}
#undef PI