#include <oxstd.h>
#import <database>
#import <maximize>

enum
{   Y_VAR, X_VAR
};

/*-------------------------- Probit : Database -----------------------------*/
class Probit : Database
{
    decl m_iResult;                           // result from MaxBFGS
    decl m_cT;                             // number of observations
    decl m_cX;                      // number of regressors (if any)

    decl m_mY;                         // dependent variable [cT][1]
    decl m_mX;                   // regressor data vector [cT][m_cX]
    decl m_vPar;                               // parameters: [p][1]
    decl m_dLogLik;                                // log-likelihood
    decl m_mCovar;              // stores variance-covariance matrix

    Probit();                                         // constructor
    fProbit(const vP, const adFunc, const avScore, const amHess);
    InitData();                        // extract data from database
    InitPar();                              // initialize parameters
    DoEstimation(vPar);                         // do the maximation
    Covar();                   // compute variance-covariance matrix
    Output();                                   // print the results
    Estimate();                         // full estimation procedure
};
/*------------------------ END Probit : Database ---------------------------*/

Probit::Probit()
{
    Database();                    // initialize base class
    m_iResult = -1;
    m_vPar = m_mCovar = <>;
    println("Probit example 4, object created on ", date());
}

Probit::fProbit(const vP, const adFunc, const avScore, const amHess)
{
    decl prob = probn(m_mX * vP);
    decl tail = 1 - prob;

    adFunc[0] = double(
        meanc(m_mY .* log(prob) + (1-m_mY) .* log(tail)));

    if (avScore)
    {
        decl weight = (m_mY - prob) .* densn(m_mX * vP)
            ./ (prob .* tail);
        avScore[0] = meanc(weight .* m_mX)';
    }

return 1;                            // 1 indicates success
}

Probit::InitData()
{
    m_mY = GetGroupLag(Y_VAR, 0, 0);          // get Y data
    m_mX = GetGroup(X_VAR); // and get X data, SetSelSample
    m_cT = m_iT2sel - m_iT1sel + 1;       //  sets m_iT?sel
    m_vPar = m_mCovar = <>;        // throw old values away
    return TRUE;
}
Probit::InitPar()
{
    m_vPar = zeros(columns(m_mX), 1);    // start from zero
    return TRUE;
}
Probit::DoEstimation(vStart)
{
    m_iResult = MaxBFGS(fProbit, &vStart, &m_dLogLik, 0, FALSE);
       // m_iResult has return code, vStart will have new params
    m_dLogLik *= m_cT;             // change to unscaled log-lik
    return {vStart, "BFGS", FALSE};       // three return values
}
Probit::Estimate()
{
    if (!InitData() || !InitPar())             // initialize all
        return FALSE;

    [m_vPar] = DoEstimation(m_vPar);                 // estimate
    
    if (m_iResult == MAX_CONV || m_iResult == MAX_WEAK_CONV)
    {
        Output();                // if successful: print results
        return TRUE;
    }
    return FALSE;                                      // failed
}

Probit::Covar()
{
    decl mhess;
    if (Num2Derivative(fProbit, m_vPar, &mhess))
    {
        m_mCovar = -invert(mhess) / m_cT;
    }
}   
Probit::Output()
{
    print("\n", MaxConvergenceMsg(m_iResult),
        " using analytical derivatives\n",
        "Log-likelihood = ", m_dLogLik, "; parameters:", m_vPar);

    Covar();
    print("standard errors:", sqrt(diagonal(m_mCovar)'));
}

main()
{
    decl probitobj = new Probit();        // create object
                             // load data file into object
    probitobj.LoadIn7("data/finney.in7");
    probitobj.Deterministic(FALSE);     // create constant
                                    // Formulate the model
    probitobj.Select(Y_VAR, { "vaso",0,0 } );
    probitobj.Select(X_VAR, { "Constant",0,0,
        "Lrate",0,0, "Lvolume",0,0 } );
    probitobj.SetSelSample(-1, 1, -1, 1);   // full sample

    MaxControl(-1, 1);             // print each iteration
    probitobj.Estimate();                      // estimate

    delete probitobj;
}
