#ifndef MARS_MJTrainEnergy
#define MARS_MJTrainEnergy

#ifndef MARS_MJTrainRanForest
#include "MJTrainRanForest.h"
#endif

class MDataSet;

class MJTrainEnergy : public MJTrainRanForest
{
private:
    TString fTrainParameter;
    TString fResultFunction;

    // To be moved to a base class
    TString fNameOutput;

public:
    MJTrainEnergy() : fNameOutput("MEnergyEst") { SetTrainLin(); }

    void SetTrainLog() { SetTrainFunc("log10(MMcEvt.fEnergy)", "pow(10, x)"); }
    void SetTrainLin() { SetTrainFunc("MMcEvt.fEnergy", "x"); }

    void SetTrainFunc(const char *par, const char *res)
    {
        fTrainParameter = par;
        fResultFunction = res;
    }

    Bool_t Train(const char *out, const MDataSet &set, Int_t num);

    ClassDef(MJTrainEnergy, 0)//Class to train Random Forest energy estimator
};


#endif
