/*--------------------------------------------------------------------------
 * ranmc.ox - ranmc code: convenient distributions for Monte Carlo experiments
 *
 *       (C) Jurgen Doornik 1998
 *
 *--------------------------------------------------------------------------*/

#include <oxstd.h>
#include <oxprob.h>
#include <ranmc.h>

RanMC::Choleski(const mSig)
{			 /* replace mSig=PP' taking zero diagonal elements into account */
	decl c = columns(mSig), sel, p, psel;
	sel = vecindex(diagonal(mSig) .!= 0);/* only part with non-zero variance*/
	
	psel = choleski(mSig[sel][sel]);
	p = zeros(c, c);
    if (columns(psel) == 0)
        return 0;

	p[sel][sel] = psel;				 /* put back at corresponding positions */
    return p;
}

RanMC::CheckDist(const sFunc, iDist, mPar1, mPar2)
{
    if (columns(mPar1) == 1)		  /* if vector: both must be row vector */
        mPar1 = mPar1';
    if (columns(mPar2) == 1)
        mPar2 = mPar2';
    if (columns(mPar1) != columns(mPar2))
        eprint(sFunc, " argument error\n");

    if (iDist == MVNORMAL_CORR)		 /* mPar2 has sd on diag, corr off-diag */
	{
		decl sd = diagonal(mPar2);
		mPar2 = sd' .* setdiagonal(mPar2, 1) .* sd;
		iDist = MVNORMAL;
	}
		
    if (iDist == MVNORMAL && columns(mPar2) > 0)
    {	                            /* replace Sigma=PP' by P' for MVNORMAL */
		mPar2 = setupper(mPar2, mPar2');		/* copy lower diag to upper */
		decl mp = RanMC::Choleski(mPar2);

		if (columns(mp) == 0)
            eprint(sFunc, " decomposition error\n");
        mPar2 = mp';				                         /* will use P' */
    }
    else if (iDist == NORMAL)
	{
        mPar2 = sqrt(mPar2);
	}
	return { iDist, mPar1, mPar2 };
}

RanMC::RanDist(const iDist, const cT, const cY, const mDf1, const mDf2)
{
    decl i, m;							   /* NB: does not handle MVNHETERO */

    if (iDist == MVNORMAL)
        return rann(cT, cY) * mDf2 + mDf1;
	else if (iDist == NORMAL)
        return rann(cT, cY) .* mDf2 + mDf1;
    else if (iDist == NO_DIST)
        return zeros(cT, cY);
    else if (iDist == MVNARCH)
    {
        m = new matrix[cT][cY];
        m[0][] = rann(1, cY) * choleski(mDf2);
        for (i = 1; i < cT; ++i)
            m[i][] = rann(1, cY) * choleski(mDf2 + mDf1*(m[i-1][]'m[i-1][])*mDf1');
        return m;
    }

    m = new matrix[1][0];
    if (iDist == LOGNORMAL)
        for (i = 0; i < cY; ++i)
            m ~= exp(rann(cT, 1) * sqrt(mDf2[0][i]) + mDf1[0][i]);
    else if (iDist == T_DIST)
        for (i = 0; i < cY; ++i)
            m ~= rant(cT, 1, mDf1[0][i]);
    else if (iDist == F_DIST)
        for (i = 0; i < cY; ++i)
            m ~= ranf(cT, 1, mDf1[0][i], mDf2[0][i]);
    else if (iDist == EXPONENTIAL)
        for (i = 0; i < cY; ++i)
            m ~= ranexp(cT, 1, mDf1[0][i]);

return m;
}
RanMC::RanDist1(const iDist, const cY, const mDf1, const mDf2, const mUlag, const mYlag)
{
    decl i, m;

    if (iDist == MVNORMAL)
        return rann(1, cY) * mDf2 + mDf1;
	else if (iDist == NORMAL)
        return rann(1, cY) .* mDf2 + mDf1;
    else if (iDist == NO_DIST)
        return zeros(1, cY);
    else if (iDist == MVNARCH)
        return rann(1, cY) * RanMC::Choleski(mDf2 + mDf1*(mUlag'mUlag)*mDf1');
    else if (iDist == MVNHETERO)
        return rann(1, cY) * RanMC::Choleski(mDf2 + mDf1*(mYlag'mYlag)*mDf1');

    m = new matrix[1][cY];
    if (iDist == LOGNORMAL)
        for (i = 0; i < cY; ++i)
            m[0][i] = exp(rann(1, 1) * sqrt(mDf2[0][i]) + mDf1[0][i]);
    else if (iDist == T_DIST)
        for (i = 0; i < cY; ++i)
            m[0][i] = rant(1, 1, mDf1[0][i]);
    else if (iDist == F_DIST)
        for (i = 0; i < cY; ++i)
            m[0][i] = ranf(1, 1, mDf1[0][i], mDf2[0][i]);
    else if (iDist == EXPONENTIAL)
        for (i = 0; i < cY; ++i)
            m[0][i] = ranexp(1, 1, mDf1[0][i]);
			
return m;
}

RanMC::WriteDist(const sPar, const iDist, const mDf1, const mDf2)
{
    decl i, cm = columns(mDf1);

    if (iDist == MVNORMAL)
    {
        decl sigma = mDf2'mDf2, fmean, fvar;
        fmean = (mDf1 == 0) ? FALSE : TRUE;
        fvar = (sigma - unit(rows(sigma)) == 0) ? FALSE : TRUE;

		if (cm == 1)
			print(sPar, " ~ N(", fmean ? "mu," : "0,", fvar ? "sigma" : "1", ")\n");
		else
			print(sPar, " ~ MVN(", fmean ? "mu," : "0,", fvar ? "sigma" : "I", ")\n");
        if (fmean)
            print("mu\'=", mDf1);
        if (rows(sigma) < 10 && fvar)
            print("sigma=", sigma);
        return;
    }
    else if (iDist == NO_DIST)
    {   print(sPar, " = 0\n");
        return;
    }
    else if (iDist == MVNARCH)
    {
        print(sPar, " ~ MVN(0,S + M u[t-1]u[t-1]\'M\')\n");
        print("S=", mDf2, "M=", mDf1);
        return;
    }
    else if (iDist == MVNHETERO)
    {
        print(sPar, " ~ MVN(0,S + M y[t-1]y[t-1]\'M\')\n");
        print("S=", mDf2, "M=", mDf1);
        return;
    }

	decl df2 = mDf2;
	if (iDist == NORMAL)
		df2 = sqr(df2);

    for (i = 0; i < cm; ++i)
    {
		if (cm > 1)
			print(sPar, i + 1, " ~ ");
		else
			print(sPar, " ~ ");
        if (iDist == NORMAL)
            print("N");
        else if (iDist == LOGNORMAL)
            print("LN");
        else if (iDist == T_DIST)
            print("t");
        else if (iDist == F_DIST)
            print("F");
        else if (iDist == EXPONENTIAL)
            print("Exp");
        print("(", mDf1[0][i], ",", df2[0][i], ")\n");
    }
}

