/*
**  EstMLDurTr
**
**  Purpose:
**    Estimate the duration model by Maximum Likelihood,
**    using a transformation
**
**  Date:
**    22/8/2012
**
**  Author:
**    Charles Bos
*/
#include <oxstd.h>
#include <oxfloat.h>
//#include <packages/gnudraw/gnudraw.h>
//#include <packages/oxutils/oxutils.h>
#import <maximize>

static decl s_vY, s_mX;

/*
**  InitData(const avY, const amX, const sData)
**
**  Purpose:
**    Read data from the Duration model
**
**  Inputs:
**    sData     string, name of data file
**
**  Outputs:
**    avY       iN x 1 vector of endogenous variables
**    amX       iN x iK matrix of explanatory variables
**
**  Return value:
**    none
*/
InitData(const avY, const amX, const sData)
{
  decl mYX;

  mYX= loadmat(sData);
  amX[0]= mYX[][1:];
  avY[0]= mYX[][0];
}

/*
**  InitP(const avP, const vY, const mX)
**
**  Purpose:
**    Initialise the vector of parameters
**
**  Inputs:
**    ...
**
**  Return value:
**    ir        integer, result of ols
*/
InitP(const avP, const vY, const mX)
{
  decl ir, vBeta;
  
  // Rough initialisation: Init beta at OLS estimates
  ir= olsc(vY, mX, &vBeta);
  
  // Add in alpha=1, as for the exponential density
  avP[0]= 1|vBeta;
  
  return ir;
}

/*
**  TransPar(const avPTr, const vP)
**
**  Purpose:
**    Transform the parameters alpha, beta to log(alpha), beta
**
**  Inputs:
**    vP        iK+1 vector with alpha and beta
**
**  Outputs:
**    avPTr     iK+1 vector, with log(alpha) and beta
**
**  Return value:
**    none
*/
TransPar(const avPTr, const vP)
{
  avPTr[0]= vP;                 // Set vPTr to a vector of the correct size, including beta
  avPTr[0][0]= log(vP[0]);      // alphaTr= log(alpha)
}

/*
**  TransBackPar(const avP, const vPTr)
**  TransBackPar(const avP, const vPTr, const amG)
**
**  Purpose:
**    Transform the parameters log(alpha), beta to alpha, beta 
**
**  Inputs:
**    vPTr      iK+1 vector with log alpha and beta
**
**  Outputs:
**    avP       iK+1 vector, with alpha and beta
**    amG       (optional) iK+1 x iK+1 matrix, jacobian of transformation
**
**  Return value:
**    ir        TRUE if all went well
*/
TransBackPar(const avP, const vPTr, ...)
{
  decl ava, mG, iK;
  
  avP[0]= vPTr;                 // Set vP to a vector of the correct size, including beta
  avP[0][0]= exp(vPTr[0]);      // alpha= exp(alphaTr)
  
  ava= va_arglist();
  if (sizeof(ava))
    {
      // Calculate analytical Jacobian
      iK= sizerc(vPTr)-1;
      mG= unit(iK+1);
      mG[0][0]= avP[0][0];    
      ava[0][0]= mG;
    }
  
  return !ismissing(avP[0]);
}

/*
**  AvgLnLDur(const vPTr, const adLnPdf, const avScore, const amHess)
**
**  Purpose:
**    Calculate the average loglikelihood of the Duression model
**
**  Inputs:
**    vPTr      iK+1 x 1 vector with log alpha and beta
**
**  Outputs:
**    adLnPdf   double, average loglikelihood value
**
**  Return value:
**    ir        TRUE if all went okay
*/
AvgLnLDur(const vPTr, const adLnPdf, const avScore, const amHess)
{
  decl vP, dAlpha, vBeta, vLnLambda, vLambda, vLY, vLL, vDa, mDb, mG;
  
  adLnPdf[0]= M_NAN;  
  
  TransBackPar(&vP, vPTr, &mG);
  [dAlpha, vBeta]= {vP[0], vP[1:]};

  vLnLambda= s_mX*vBeta;  
  vLambda= exp(vLnLambda);  
  vLY= vLambda .* s_vY;
  
  vLL= log(dAlpha) + dAlpha * vLnLambda 
    + (dAlpha-1)*log(s_vY) - vLY.^dAlpha;
  adLnPdf[0]= double(meanc(vLL));
  print ("+");
  
  if (avScore)
    {
      vDa= 1/dAlpha + vLnLambda + log(s_vY) 
        - log(vLY) .* vLY.^dAlpha;
      mDb= (dAlpha - dAlpha * vLY .^ dAlpha) .* s_mX;
      
      // Don't forget to apply the jacobian of the transformation of the parameters
      avScore[0]= ((meanc(vDa)~meanc(mDb)) * mG)';
    }

  return !ismissing(adLnPdf[0]);
}

/*
**  Estimate(const avP, const adLnPdf, const vY, const mX)
**
**  Purpose:
**    Estimate the model
**
**  Inputs:
**    avP       iP x 1 vector of initial parameters
**    vY        iN x 1 vector of data
**    mX        iN x iK matrix of explanatory variables
**
**  Outputs:
**    avP       iP x 1 vector of estimated parameters
**    adLnPdf   double, optimal value of average LL
**
**  Return value:
**    ir        integer, MaxBFGS indication of convergence
*/
Estimate(const avP, const adLnPdf, const vY, const mX)
{
  decl ir, vS1, vS2, vPTr, vP1;
  
  s_vY= vY;
  s_mX= mX;
  
  TransPar(&vPTr, avP[0]);

  println ("Checking transformation and scores...");
  TransBackPar(&vP1, vPTr);
  ir= AvgLnLDur(vPTr, adLnPdf, &vS1, 0);
  ir= Num1Derivative(AvgLnLDur, vPTr, &vS2);           
  println ("\nAvgLnLDur returns ", ir, ", with AvgLL= ", adLnPdf[0], 
           " and ", "%c", {"p", "pTr", "p-back", "sc-anal", "sc-num", "s1-s2"},
           avP[0]~vPTr~vP1~vS1~vS2~(vS1-vS2));
           
  ir= MaxBFGS(AvgLnLDur, &vPTr, adLnPdf, 0, TRUE);
  TransBackPar(avP, vPTr);
  
  return ir;
}



main()
{
  decl sData, vY, mX, iN, vP0, vP, ir, dLnPdf;
  
  // Magic numbers
  sData= "data/genrdur.fmt";
  
  // Initialisation
  InitData(&vY, &mX, sData);
  InitP(&vP0, vY, mX);  
  
//  vP0= <10; -10; 20>;
  vP0[0]= .5;
  vP= vP0;
  ir= Estimate(&vP, &dLnPdf, vY, mX);
  print ("\nMaxBFGS returns ", MaxConvergenceMsg(ir), 
         " at pars ", 
         "%c", {"p0", "est"},
         "%r", {"alpha", "b0", "b1"}, 
         vP0~vP, 
         " with average LL= ", dLnPdf);
  
}