#include <oxstd.h>
#include <oxdraw.h>
#include <oxfloat.h>
#include <oxprob.h>

#include "bprobitex.h"

/*------------------------ BprobitEx : Bprobit ----------------------*/
BprobitEx::BprobitEx()
{
	// initialize base class
	Modelbase();

	// set default for method, model class, and model counter
	m_iMethod = M_PROBIT;
	m_iModelClass = MC_BINARY;
	m_iModel = 0;
	// ... remaining defaults

	// print startup message
	println("---- ", GetPackageName(), " ", GetPackageVersion(),
        " session started at ", time(), " on ", date(), " ----");
	
}
BprobitEx::GetPackageName()
{	// return package name
	return "BprobitEx";
}
BprobitEx::GetPackageVersion()
{	// return package version
	return "1.0";
}
BprobitEx::GetParNames()
{	// return array of strings with parameter names
	return m_asX;
}
BprobitEx::SetSelSample(const iYear1, const iPeriod1, const iYear2,
	const iPeriod2)
{	// cross-section: always select full sample
	m_iT1sel = 0;
    m_iT2sel = rows(m_mData) - 1;
}

BprobitEx::InitData()
{
	// get and inspect the data selected for estimation
	m_iT1est = m_iT1sel;  m_iT2est = m_iT2sel;
	m_iT1sel = 0;  m_iT2sel = rows(m_mData) - 1;

	// get data: Y,X,W,Sel; not using modelbase version because
	// that defaults to time-series sample selection
    m_mY = GetGroup(Y_VAR);
    m_cY = columns(m_mY);
    if (!m_cY)
    {   println("\n*** Error: need explanatory variable");
        return FALSE;
    }
    m_mX = GetGroup(X_VAR);
    m_cX = columns(m_mX);
    if (!m_cX)
    {   println("\n*** Error: need some regressors");
        return FALSE;
    }
    m_mW = GetGroup(W_VAR);
    m_cW = columns(m_mW);
    if (columns(m_mW) > 1)
    {   println("\n*** Error: only one weight variable allowed");
        return FALSE;
    }
    m_mSel = GetGroup(SEL_VAR);
    m_cSel = columns(m_mSel);
    if (columns(m_mSel) > 1)
    {   println("\n*** Error: only one selection variable allowed");
        return FALSE;
    }
	// check for missing values
	decl vdrop = sumr(isdotmissing(m_mX));
	if (m_cW) vdrop = vdrop .|| isdotmissing(m_mW);
	if (m_cSel) vdrop = vdrop .|| isdotmissing(m_mSel) .|| m_mSel .== 0; 
	vdrop = vdrop .|| sumr(isdotmissing(m_mY));
	m_mIdx = range(0, rows(m_mY) - 1)' + m_iT1sel;
	// now drop those observations with missing values
	if (sumc(vdrop))
	{
		m_mY = deleteifr(m_mY, vdrop);
		m_mX = deleteifr(m_mX, vdrop);
		m_mIdx = deleteifr(m_mIdx, vdrop);
		if (m_cW) m_mW = deleteifr(m_mW, vdrop);
		// don't drop from selection variable
	}
	// set sample size
	m_cT = rows(m_mY);
    if (m_cT <= 2)
    {   println("\n*** Error: only ", m_cT, " observations.");
        return FALSE;
    }
	// get names of variables
	GetGroupNames(Y_VAR, &m_asY);
	GetGroupNames(X_VAR, &m_asX);
	GetGroupNames(W_VAR, &m_asW);
	GetGroupNames(SEL_VAR, &m_asSel);
	// update status: data processed successsfully
	m_iModelStatus = MS_DATA;
    return TRUE;
}

BprobitEx::DoEstimation(vP)
{
	// do the maximization using BFGS
	if (m_iMethod == M_PROBIT)
	    m_iResult = MaxBFGS(fProbit, &vP, &m_dLogLik, 0, FALSE);
	else if (m_iMethod == M_LOGIT)
		println("\n*** Error: not implemented");
	else if (m_iMethod == M_POISSON)
		println("\n*** Error: not implemented");
	else if (m_iMethod == M_NEGBIN)
		println("\n*** Error: not implemented");

	// maximization is done in terms of loglik/T
    m_dLogLik *= m_cT;

	// may need to do more processing if estimation succeeded
	if (m_iResult <= MAX_WEAK_CONV)
	{
		++m_iModel;
	}
    return {vP, "BFGS", FALSE};       // three return values
}

BprobitEx::Output()
{
	Buffering(11);	  // set text marker and start buffering

	// print the header even if estimation failed
	OutputHeader(GetMethodLabel());

	if (m_iModelStatus != MS_ESTIMATED)
		return;

	// output the rest
	OutputPar();
    OutputLogLik();

	Buffering(0);	  // stop buffering
}
///////////////////////////////////////////////////////////////////////
// OxPack specific
BprobitEx::Buffering(const iBufferOn)
{
	if (iBufferOn == 11)
	{
		"OxPackSetMarker"(TRUE);
		"OxPackBufferOn"();
	}
	else if (iBufferOn == 1)
		"OxPackBufferOn"();
	else
		"OxPackBufferOff"();
}
BprobitEx::IsCrossSection()
{
	return 1;
}
BprobitEx::GetModelLabel()
{
	return sprint("CS(", "%2d", m_iModel, ")");
}
BprobitEx::GetMethodLabel()
{
	decl asmethods = {"Logit", "Probit", "Poisson"};
	return asmethods[m_iMethod];
}
BprobitEx::SendSpecials()
{
    return {"Constant"};
}   
BprobitEx::SendVarStatus()
{
	if (m_iModelClass == MC_BINARY)
		return
        {{ "&Y endogenous",'Y',STATUS_GROUP + STATUS_ENDOGENOUS, Y_VAR},
         { "&X variable", 'X', STATUS_GROUP, X_VAR},
         { "&Weight",     'W', STATUS_GROUP + STATUS_ONEONLY, W_VAR},
         { "&Select By",  'S', STATUS_GROUP + STATUS_ONEONLY, SEL_VAR}
		};
	else
		return
        {{ "&Y endogenous",'Y',STATUS_GROUP + STATUS_ENDOGENOUS, Y_VAR},
         { "&X variable", 'X', STATUS_GROUP, X_VAR}
		};
}
BprobitEx::SendMenu(const sMenu)
{
	if (sMenu == "ModelClass")
	{
		return 
			{{ "&1: Binary Discrete Choice", "modelclass0",
					m_iModelClass == MC_BINARY},
			 { "&2: Count Data",			 "modelclass1",
			 		m_iModelClass == MC_COUNT}
			 };
	}
	else if (sMenu == "Test")
	{	return 
			{{ "&Graphic Analysis...", m_iModelClass == MC_COUNT
					? "" : "OP_TEST_GRAPHICS"},
			 { "&Predictions...",      m_iModelClass == MC_COUNT
			 		? "" : "OP_TEST_FORECAST"},
			 0,
			 { "&Further Output...", "Further Output"},
			 { "&Outliers...", "Outliers"},
			 0,
			 { "&Exclusion Restrictions...", "OP_TEST_SUBSET"},
			 { "&Linear Restrictions...",    "OP_TEST_LINRES"},
			 0,
			 { "Store in D&atabase...", "OP_TEST_STORE"}
			 };
	}
}

BprobitEx::ReceiveModel()
{
    Select(SEL_VAR, "OxPackGetData"("SelGroup", SEL_VAR));

	// m_iMethod determined in ReceiveDialog, prevent
	// Modelbase::ReceiveModel() from changing it.
	decl imethod = m_iMethod;
	Modelbase::ReceiveModel();
	m_iMethod = imethod;
}

BprobitEx::DoOption(const sOpt, const val)
{
	if (sOpt == "covunw")
		m_fCovarUnweighted = val;
}
BprobitEx::DoSettingsDlg()
{
	decl adlg, asoptions, avalues;
	
	// settings: choose model specific settings
	// depends on the model class
	if (m_iModelClass == MC_BINARY)
		adlg =
		{	{ "Choose a model:" },
			{ "Logit", CTL_RADIO,
					m_iMethod == M_PROBIT ? 1 : 0, "method" },
			{ "Probit", CTL_RADIO }
		};
	else if (m_iModelClass == MC_COUNT)
		adlg =
		{	{ "Choose a model:" },
			{ "Poisson", CTL_RADIO,
					m_iMethod == M_NEGBIN ? 1 : 0, "method" },
			{ "Negative binomial", CTL_RADIO }
		};

	if (!"OxPackDialog"("Settings", adlg, &asoptions, &avalues))
		return FALSE;

	// process user actions for settings dialog
	if (m_iModelClass == MC_BINARY)
		m_iMethod = avalues[0] == 1 ? M_PROBIT : M_LOGIT;
	else if (m_iModelClass == MC_COUNT)
		m_iMethod = avalues[0] == 1 ? M_NEGBIN : M_POISSON;
	return TRUE;
}
BprobitEx::DoEstimateDlg(const iFirstMethod, const cMethods, const sMethods,
		const bForcAllowed, const bRecAllowed, const bMaxDlgAllowed)
{
	if (Modelbase::DoEstimateDlg(0, 1, "BFGS method", FALSE, FALSE, FALSE))
	{
		return TRUE;
	}
	return FALSE;
}
BprobitEx::DoGraphicsDlg()
{
	decl adlg, asoptions, avalues;
	
	adlg =
	{	{ "Graphic Analysis" },
		{ "Histograms of probabilities for each state",
				CTL_CHECK, 1, "hist" },
		{ "Histograms of probabilities of observed state",
				CTL_CHECK, 1, "histobs" },
		{ "Number of bars:", CTL_INT, 10, "bars" },
		{ "Cumulative correct predictions for each state",
				CTL_CHECK, 0, "ccp" },
		{ "Unsorted", CTL_RADIO, 1, "sort" },
		{ "Sorted by probability", CTL_RADIO },
		{ "Sorted by log-likelihood contribution", CTL_RADIO },
		{ "Cumulative response for each state",
				CTL_CHECK, 0, "crf" }
	};

	if (!"OxPackDialog"("Graphic Analysis", adlg, &asoptions, &avalues))
		return FALSE;

	// process user actions for graphics dialog

	decl iplot = 0;
	
	// call functions if requested
	if (iplot)
		ShowDrawWindow();
	return TRUE;
}

BprobitEx::ReceiveMenuChoice(const sDialog)
{
	if (sDialog == "modelclass0")
	{
		// set model class, and activate model formulation dialog
		m_iModelClass = MC_BINARY;
		return 1;
	}
	else if (sDialog == "modelclass1")
	{
		// set model class, and activate model formulation dialog
		m_iModelClass = MC_COUNT;
		return 1;
	}
	else if (sDialog == "OP_OPTIONS")
    {
		// append BprobitEx options to the Modelbase ones
		decl adlg;
		adlg ~=
			{	{ "Further options", CTL_GROUP, 1 },
				{ "Use unweighted covariance matrix",
						CTL_CHECK, m_fCovarUnweighted, "covunw" }
			};
		return DoOptionsDlg(adlg);
	}
	else if (sDialog == "OP_TEST_GRAPHICS")
	{
		return DoGraphicsDlg();
	}
	else if (sDialog == "OP_TEST_FORECAST")
	{
	}
	else if (sDialog == "Outliers")
	{
	}
	else if (sDialog == "Further Output")
	{
	}
	// allow base class to process unhandled cases
	return Modelbase::ReceiveMenuChoice(sDialog);
}

BprobitEx::LoadOptions()
{
	// load persistent settings
	Modelbase::LoadOptions();
	
	m_iModelClass = "OxPackReadProfileInt"(0,"class", 0);
	m_fCovarUnweighted = "OxPackReadProfileInt"(0,"covunw", 0);
}
BprobitEx::SaveOptions()
{
	// save persistent settings
	"OxPackWriteProfileInt"(0,"class", m_iModelClass);
	"OxPackWriteProfileInt"(0,"covunw", m_fCovarUnweighted);

	Modelbase::SaveOptions();
}
BprobitEx::GetModelSettings()
{
	// process model settings from model recall
	// could be m_iModelClass specific
	return
	{	{	m_iMethod,		"method"	}
	};
}
BprobitEx::SetModelSettings(const aValues)
{
	// allows for model settings storage in model history
	if (!isarray(aValues))
		return;
	for (decl i = 0; i < sizeof(aValues); ++i)
	{
		if (aValues[i][1] == "method")
			m_iMethod = aValues[i][0];
	}
}
BprobitEx::BatchVarStatus(const sTypes, const vcTypes)
{
	// inspect variable types to decide on appropriate model class
	// only switch if W found
	// model class may also need changing in BatchMethod
	for (decl i = sizeof(sTypes) - 1; i >= 0; --i)
		if (sTypes[i] == 'W')
		{	m_iModelClass = MC_BINARY;
		}

	return m_iModelClass;
}
