#ifndef MARS_MJTrainSeparation
#define MARS_MJTrainSeparation

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

#ifndef MARS_MDataSet
#include "MDataSet.h"
#endif

class MH3;

class MJTrainSeparation : public MJTrainRanForest
{
private:
    MDataSet fDataSetTest;
    MDataSet fDataSetTrain;

    UInt_t fNumTrainOn;
    UInt_t fNumTrainOff;

    UInt_t fNumTestOn;
    UInt_t fNumTestOff;

    Bool_t fAutoTrain;
    Bool_t fUseRegression;

    void DisplayResult(MH3 &h31, MH3 &h32);

    Bool_t   GetEventsProduced(MDataSet &set, Double_t &num, Double_t &min, Double_t &max) const;
    Double_t GetDataRate(MDataSet &set, Double_t &num) const;
    Double_t GetNumMC(MDataSet &set) const;
    Bool_t   AutoTrain();

public:
    MJTrainSeparation() :
        fNumTrainOn((UInt_t)-1), fNumTrainOff((UInt_t)-1),
        fNumTestOn((UInt_t)-1), fNumTestOff((UInt_t)-1),
        fAutoTrain(kFALSE), fUseRegression(kTRUE)
    { }

    void SetDataSetTrain(const MDataSet &ds, UInt_t non=(UInt_t)-1, UInt_t noff=(UInt_t)-1)
    {
        ds.Copy(fDataSetTrain);

        fDataSetTrain.SetNumAnalysis(1);

        fNumTrainOn = non;
        fNumTrainOff = noff;
    }
    void SetDataSetTest(const MDataSet &ds, UInt_t non=(UInt_t)-1, UInt_t noff=(UInt_t)-1)
    {
        ds.Copy(fDataSetTest);

        fDataSetTest.SetNumAnalysis(1);

        fNumTestOn = non;
        fNumTestOff = noff;
    }

    void EnableAutoTrain(Bool_t b=kTRUE)      { fAutoTrain     =  b; }
    void EnableRegression(Bool_t b=kTRUE)     { fUseRegression =  b; }
    void EnableClassification(Bool_t b=kTRUE) { fUseRegression = !b; }

    Bool_t Train(const char *out);

    ClassDef(MJTrainSeparation, 0)//Class to train Random Forest gamma-/background-separation
};

#endif