#include <oxstd.h>
#include <arma.h>

#import "../code/autometrics"
#import "../code/pcfimlex"
#import "../code/simstore"
#import "../code/simdesign"
#import "../code/simutils"

///////////////////////////////////////////////////////////////////////
// PcFimlEx2: adding bias corrected coeffs and dynamic forecasts
class PcFimlEx2 : PcFimlEx
{
	PcFimlEx2();
	DynamicForecast(const cTforc, const vBeta=<>);
	GetParBC2(const dAutoPval);
}
PcFimlEx2::PcFimlEx2()
{
	PcFimlEx();
}
PcFimlEx2::GetParBC2(const dAutoPval)
{
	decl vp = GetPar(), vpbc, cx = sizerc(vp) - m_cU * m_cY;

	if (cx)
	{
		vpbc = GetBiasCorrections(vp, vp ./ GetStdErr(), quant(1 - (dAutoPval / 2), GetcT() - GetcDfLoss()))[1];
		// keep unrestricted coefficients
		vp[ : cx - 1] = vpbc;
	}
	return vp;
}
PcFimlEx2::DynamicForecast(const cTforc, const vBeta)
{
	if (cTforc <= 0)
		return zeros(0, 3);
	decl yforc, y, asel = GetSelInfo(), t2 = GetSelEnd(), asy; 
	decl idx_y = vecindex(asel[1] .== Y_VAR .&& asel[2] .== 0);
	decl vp = vBeta == <> ? GetPar() : vBeta;

	decl sel_w = dropc(asel[0], idx_y);
	decl grp_w = dropc(asel[1], idx_y);
	decl lag_w = dropc(asel[2], idx_y);
	decl order = vecindex(grp_w .== Y_VAR .&& lag_w .!= 0) | vecindex(grp_w .== X_VAR)
		 | vecindex(grp_w .== U_VAR);

	// vBeta is in the model order: lagged Y, X, U
	// make sure that the specification is in the same order
	sel_w = sel_w[order];
	grp_w = grp_w[order];
	lag_w = lag_w[order];
	
	GetGroupLagNames(Y_VAR, 0, 0, &asy);	              /* dependent vars */

	if (sizerc(sel_w))
		yforc = modelforc(<>, GetAll(), GetVarIndex(asy[0]),
    		sel_w, lag_w, vec(vp)', t2 + 1)[ : cTforc - 1];
	else
		yforc = zeros(cTforc, 1);

	y = GetVar(asy[0])[t2 + 1 : t2 + cTforc];
	return yforc ~ y;
}

///////////////////////////////////////////////////////////////////////

///////////////////////////////////////////////////////////////////////
// SimDesignEx: adding transformations of regressors
class SimDesignEx : SimDesign
{
	SimDesignEx(const obj, const iYvar, const iXvar, const iUvar);
	SetTransformX(const iMode);
	virtual TransformX1(const mX);
	GetLabelTransformX1(dCorrX);
	decl m_iTransformX;
}
SimDesignEx::SimDesignEx(const obj, const iYvar, const iXvar, const iUvar)
{
	SimDesign(obj, iYvar, iXvar, iUvar);
	m_iTransformX = 0;
}
SimDesignEx::SetTransformX(const iMode)
{
	m_iTransformX = iMode;
}
SimDesignEx::TransformX1(const mX)
{
	decl mx;
	switch_single (m_iTransformX)
	{
		case 0:  return mX;
		case 1:  return sqrt(12) * (probn(mX) - 0.5);
		case 2:  return sqrt(2) * (sqr(mX) - 1);
		case 3:  return 2.22 * (log(probn(mX)) + 1.26);
		case 4:  mx = cumulate(mX); return mx - meanc(mx);
		case 11: return sqrt(12) * (probn(mX));
		case 12: return sqrt(2) * (sqr(mX));
		case 13: return 2.22 * (log(probn(mX)));
		case 14: return cumulate(mX);
	}
}
SimDesignEx::GetLabelTransformX1(dCorrX)
{
	switch_single (m_iTransformX)
	{
		case  0: return dCorrX ? sprint("N[0,Cx], c_ii=1, c_ij=", dCorrX) : "IIN[0,1]";
		case  1: return dCorrX ? sprint("sqrt(12)U(-0.5,0.5)Px', c_ii=1, c_ij=", dCorrX) : "sqrt(12)IIU(-0.5,0.5)";            
		case  2: return dCorrX ? sprint("sqrt(2)[chi^2(1)-1]Px', c_ii=1, c_ij=", dCorrX) : "sqrt(2)[IIchi^2(1)-1]";            
		case  3: return dCorrX ? sprint("sqrt(2)(log[U(0,1)]+1.26)Px', c_ii=1, c_ij=", dCorrX) : "sqrt(2)(log[IIU(0,1)]+1.26)";
		case  4: return dCorrX ? sprint("cumulate(N[0,Cx]), c_ii=1, c_ij=", dCorrX, " mean 0") : "cumulate(IIN[0,1]) mean 0";
		case 11: return dCorrX ? sprint("sqrt(12)U(0,1)Px', c_ii=1, c_ij=", dCorrX) : "sqrt(12)IIU(0,1)";            
		case 12: return dCorrX ? sprint("sqrt(2)[chi^2(1)]Px', c_ii=1, c_ij=", dCorrX) : "sqrt(2)[IIchi^2(1)]";            
		case 13: return dCorrX ? sprint("sqrt(2)(log[U(0,1)])Px', c_ii=1, c_ij=", dCorrX) : "sqrt(2)(log[IIU(0,1)])";
		case 14: return dCorrX ? sprint("cumulate(N[0,Cx]), c_ii=1, c_ij=", dCorrX) : "cumulate(IIN[0,1])";
	}
}
///////////////////////////////////////////////////////////////////////

///////////////////////////////////////////////////////////////////////
// rest
StoreModelForc(const oStore, const oModel, const iStore, const iRep, const iRepStage, const vForc, const vAct, const vP)
{
	oStore.StoreResults(iStore, iRep, iRepStage, vP, oModel.GetStdErr(), oModel.AutometricsSignificance(),
		oModel.GetParNames(), sqrt(oModel.GetOmega()), oModel.GetLogLik(), oModel.GetcT(), vForc, vAct);
}

ModelAverage(const iMode, const cM, const cXsig, const vTval, const cXinsig,
	const dXcorr, const cStep, const dStepTval, const cTest, const cTforc,
	const dAutoPval, const bBiasCorrect, const iTransformX)
{
	decl iconstant = SimDesign::USE_CONSTANT_U;
	decl csteplen = 4;

	decl design = new SimDesignEx(new PcFimlEx2(), Y_VAR, X_VAR, U_VAR);

	decl ascol, ascol_forc, j, im, k, vp, vp_model, yf;
	decl time = timer();
	
	format(1000);
	
	// add step dummy starting in observation cTest - csteplen
	design.CreateDatabase("", cTest + cTforc, cXsig + cXinsig, cTest - csteplen ~ -1);
	design.CreateReference(0, cXsig + cXinsig, 0, iconstant, FALSE, cStep, 0, cTest);
	vp = vec(vTval)' / sqrt(cTest)						 // relevant
		~ zeros(1, cXinsig)								 // irrelevant
		~ constant(dStepTval / sqrt(csteplen), 1, cStep) // step dummy
		~ zeros(1, iconstant > 0);						 // Unrestriced constant
	design.SetDGP(vp);
	vp_model = ones(1, sizerc(vp));	  // 1 at location of coef that enters initial GUM
//	vp_model[0] = 0;	// misspecified initial GUM
	design.SetModel(vp_model);

	design.SetTransformX(iTransformX);
	design.SetXCorrelation(dXcorr, SimDesign::USE_XCOR_MATCONST);
	design.SetRenewX(TRUE);

	println("DGP X: ", design.GetLabelTransformX1(dXcorr));

	decl store = new SimStore(design.GetRefParDGP(), design.GetRefParNames(), design.GetSigmaInDGP(), cM, cTforc);

	if (iMode == 0)
	{
		println("\n**** Averaging over all 2^k models ***\n");
		ascol = {"2^k(U,eq)", "2^k(U,SC)", "2^k(C,SC)"};
		ascol_forc = {"FA(2^k,U,eq)", "FA(2^k,U,SC)", "MA(2^k,U,eq)", "MA(2^k,U,SC)", "MA(2^k,C,SC)"};
	}
	else if (iMode == 1)
	{
		println("\n**** Model selection using Autometrics, final model only ***\n");
		ascol = {"Auto"};
		ascol_forc = {"Autometrics", "FA(Auto,SC)", "MA(Auto,eq)", "MA(Auto,SC)"};
	}
	else if (iMode == 2)
	{
		println("\n**** Model selection using Autometrics, all final candidates ***\n");
		ascol = {"AuT(all,eq)", "AuT(all,SC)", "AuT(allC,SC)"};
		ascol_forc = {"FA(AuT,U,eq)", "FA(AuT,U,SC)", "MA(AuT,U,eq)", "MA(AuT,U,SC)", "MA(AuT,C,SC)"};
	}														  
	decl refmodel = design.GetRefModelObject(), asdiag = {};
	
	for (im = 0; im < cM; ++im)
	{
		decl model = clone(refmodel);
		decl auto = new Autometrics();

		model.SetPrint(0);
		auto.SetPrint(0);
		auto.SetStrategyPValue(dAutoPval);
	
		design.Generate(model, im);

		design.FormulateDGP(model);
		model.Estimate();
		vp = model.GetPar();
		yf = model.DynamicForecast(cTforc, vp);
		StoreModelForc(store, model, SimStore::STORE_DGP, im, 0, yf[][0], yf[][1], vp);

		design.FormulateModel(model);
		model.Estimate();
		vp = model.GetPar();
		StoreModelForc(store, model, SimStore::STORE_GUM, im, 0, model.DynamicForecast(cTforc, vp)[][0], <>, vp);

		if (iMode == 0)
		{
			decl subset = <>;
			decl asel = model.GetSelInfo(), asel_sub = asel;
			decl idx_xfree = vecindex((asel[1] .== Y_VAR .&& asel[2] .!= 0) .|| asel[1] .== X_VAR);
			decl cxfree = sizerc(idx_xfree);
	
			for (j = 0; ; ++j)
			{
				asel_sub[0] = dropc(asel[0], idx_xfree[subset]);
				asel_sub[1] = dropc(asel[1], idx_xfree[subset]);
				asel_sub[2] = dropc(asel[2], idx_xfree[subset]);
	
				model.ClearModel();
				model.SetSelInfo(asel_sub);
	
				model.Estimate();
				vp = model.GetPar();
				// j==0: store; j>0: add parameters and forecasts to previous
				StoreModelForc(store, model, SimStore::STORE_REP, im, j, model.DynamicForecast(cTforc, vp)[][0], <>, vp);
				StoreModelForc(store, model, SimStore::STORE_RSC, im, j, model.DynamicForecast(cTforc, vp)[][0], <>, vp);
	
				subset = Autometrics::NextSubset(cxfree, subset);
				if (subset == <>)
					break;
			}
		}
		else
		{
			k = auto.Estimate(model, Y_VAR, X_VAR);
			vp = bBiasCorrect ? model.GetParBC2(dAutoPval) : model.GetPar();
			// j==0: store parameters and forecasts
			StoreModelForc(store, model, SimStore::STORE_REP, im, 0, model.DynamicForecast(cTforc, vp)[][0], <>, vp);
			StoreModelForc(store, model, SimStore::STORE_RSC, im, 0, model.DynamicForecast(cTforc, vp)[][0], <>, vp);
	
			if (iMode == 2)
			{
				for (j = 1; ; ++j)
				{
					if (!auto.EstimateFinalModel(j))
						break;
					if (j == k)
						continue;
					vp = bBiasCorrect ? model.GetParBC2(dAutoPval) : model.GetPar();
					// j>0: add parameters and forecasts to previous 
					StoreModelForc(store, model, SimStore::STORE_REP, im, j, model.DynamicForecast(cTforc, vp)[][0], <>, vp);
					StoreModelForc(store, model, SimStore::STORE_RSC, im, j, model.DynamicForecast(cTforc, vp)[][0], <>, vp);
				}
			}
		}
		design.FormulateReferenceModel(model);
		model.Estimate();
		store.StoreForecastsFromAverage(SimStore::STORE_FORC_AVG1, im,	   // unconditional
			model.DynamicForecast(cTforc, store.GetStoredPar(SimStore::STORE_RSC, im, TRUE))[][0]);
		if (iMode == 0 || iMode == 2)
			store.StoreForecastsFromAverage(SimStore::STORE_FORC_AVG2, im,  // conditional
				model.DynamicForecast(cTforc, store.GetStoredPar(SimStore::STORE_RSC, im, FALSE))[][0]);
		store.StoreForecastsFromAverage(SimStore::STORE_FORC_AVG0, im,	 // unconditional
			model.DynamicForecast(cTforc, store.GetStoredPar(SimStore::STORE_REP, im, TRUE))[][0]);

		delete model;
	    delete auto;
	}
	
	if (iMode != 0)
	{
		decl auto = new Autometrics();
		auto.SetStrategyPValue(dAutoPval);
		auto.SetPrint(1);
		auto.ReportSettings();
	    delete auto;
	}
	println("T for estimation", "%13d", cTest, "  #forecasts", "%19d", cTforc);
	println("#significant", "%17d", cXsig,     "  #insignificant", "%15d", cXinsig);
	println("Constant", "%21s", {"none", "fixed", "free"}[iconstant], "  X correlation", "%16.3f", dXcorr);
	println("Step t-value", "%17g", dStepTval, "  bias correction", "%14s", bBiasCorrect ? "yes" : "no");
	println("replications", "%17d", cM, "  time ", "%24s", timespan(time));
	
	decl ret = store.ReportPar(ascol, ascol_forc);

	println("key for modes:");
	println("%-6s", "2^k:", "all subsets");
	println("%-6s", "Auto:", "Autometrics, final model (all the same)");
	println("%-6s", "AuT:", "Autometrics, all terminal candidates");
	println("key for averaging:");
	println("%-6s", "MA:", "forecasts from averaged model");
	println("%-6s", "FA:", "average of forecasts");
	println("%-6s", "eq:", "equal weights");
	println("%-6s", "SC:", "SC weights");
	println("%-6s", "U:", "unconditional average");
	println("%-6s", "C:", "conditional average (only non-zero coefficients)");

	delete design;
	delete store;

	return ret;
}
RunT(const cM, const vTvals, const vStepTvals, const dAutoPval, const bBiasCorrect,
	const bDoAllSubsets, const dXcorr=0.0, const iTransformX=0, const bAltSigns=FALSE)
{
	decl cstep = sizerc(vStepTvals) > 0;
	decl cx_sig = 5 - cstep;
	decl vtval;
	decl cx_insig = cx_sig + cstep;
	decl cx_correlation = dXcorr;
	decl dsteptval = 3;
	decl ctest = 100;
	decl ctforc = 1;
	decl ret = {<>,{}}, ret1, ret2, ret3;
	
	for (decl i = 0; i < sizerc(vTvals); ++i)
	{
		decl vtval = constant(vTvals[i], 1, cx_sig);
		if (bAltSigns)
			vtval[range(1, cx_sig - 1, 2)] .*= -1;
			
		for (decl j = 0; j < max(sizerc(vStepTvals), 1); ++j)
		{
			decl dsteptval = sizerc(vStepTvals) ? vStepTvals[j] : 0;
			
			if (dsteptval)
				println("\n---- t-value: ", vtval[0][0], ", step t-value: ", dsteptval, " ----\n");
			else
				println("\n---- t-value: ", vtval[0][0], " ----\n");
		
			ranseed(-1);
			ret1 = ModelAverage(1, cM, cx_sig, vtval, cx_insig, cx_correlation, cstep, dsteptval, ctest, ctforc, dAutoPval, bBiasCorrect, iTransformX);
		
			ranseed(-1);
			ret2 = ModelAverage(2, cM, cx_sig, vtval, cx_insig, cx_correlation, cstep, dsteptval, ctest, ctforc, dAutoPval, bBiasCorrect, iTransformX);
		
			ranseed(-1);
			if (bDoAllSubsets)
			{
				ret3 = ModelAverage(0, cM, cx_sig, vtval, cx_insig, cx_correlation, cstep, dsteptval, ctest, ctforc, dAutoPval, bBiasCorrect, iTransformX);
				ret3 = {ret3[0][][2:], ret3[1][2:]};
			}
			else
				ret3 = {<>,{}};

			ret[0] |= vTvals[i] ~ (dsteptval ~ (ret1[0] ~ ret2[0][][2:] ~ ret3[0]));
			if (i == 0 && j == 0)
				ret[1] = {"t-val"} ~ {"Step_t-val"} ~ ret1[1] ~ ret2[1][2:] ~ ret3[1];
		}
	}
return ret;
}
Save(const mas, const sFile)
{
	decl db = new Database();
	db.Create(1, 1, 1, sizer(mas[0]), 1);
	db.Append(mas[0] ~ constant(.NaN, sizer(mas[0]), sizeof(mas[1]) - sizec(mas[0])), mas[1]);
	db.SaveIn7(sFile);
}
Run(const cM, const dPvalue, const bBiasCorrect, const bDoAllSubsets, const bSave=TRUE, const dXcorr=0.0, const iTransformX=0)
{
	decl ret = RunT(cM, <0.1,1:5,8>, <>, dPvalue, bBiasCorrect, bDoAllSubsets, dXcorr, iTransformX);
	println("\nSummary MSFE results:", "%c", ret[1], ret[0]);
	if (bSave)
		Save(ret, sprint("23_11_msfe_", (100*dPvalue)) ~ (bBiasCorrect ? "_bc" : "") );
}
RunStep(const cM, const dPvalue, const bBiasCorrect, const bDoAllSubsets, const bSave=TRUE, const dXcorr=0.0, const iTransformX=0)
{
	decl ret = RunT(cM, <3>, <0.1,1:5,8>, dPvalue, bBiasCorrect, bDoAllSubsets, dXcorr, iTransformX);
	println("\nSummary MSFE results:", "%c", ret[1], ret[0]);
	if (bSave)
		Save(ret, sprint("23_11_msfe_step_", (100*dPvalue)) ~ (bBiasCorrect ? "_bc" : "") );
}
main()
{
	decl cm = 1000;


	Run(cm, 0.01, FALSE, FALSE);
//	Run(cm, 0.05, FALSE, TRUE );
	Run(cm, 0.05, FALSE, FALSE);
	Run(cm, 0.10, FALSE, FALSE);
	Run(cm, 0.25, FALSE, FALSE);
	Run(cm, 0.50, FALSE, FALSE);

//  bias correction
//	Run(cm, 0.05, TRUE, FALSE);

//	Run(cm, 0.5, TRUE, FALSE, FALSE, 0.0, 0);
}	