#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;

    TList fPreTasksOn;
    TList fPreTasksOff;
    TList fPostTasksOn;
    TList fPostTasksOff;

    Bool_t fAutoTrain;
    Bool_t fUseRegression;
    Bool_t fEnableWeightsOn;
    Bool_t fEnableWeightsOff;

    Float_t fFlux;

    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(MDataSet &set, UInt_t &on, UInt_t &off);

public:
    MJTrainSeparation() :
        fNumTrainOn((UInt_t)-1), fNumTrainOff((UInt_t)-1),
        fNumTestOn((UInt_t)-1), fNumTestOff((UInt_t)-1),
        fAutoTrain(kFALSE), fUseRegression(kFALSE),
        fEnableWeightsOn(kFALSE), fEnableWeightsOff(kFALSE), fFlux(2e-7)
    { }

    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 AddPreTaskOn(MTask *t)                     { Add(fPreTasksOn,  t); }
    void AddPreTaskOn(const char *rule,
                      const char *name="MWeight")   { AddPar(fPreTasksOn, rule, name); }

    void AddPreTaskOff(MTask *t)                    { Add(fPreTasksOff, t); }
    void AddPreTaskOff(const char *rule,
                       const char *name="MWeight")  { AddPar(fPreTasksOff, rule, name); }

    void AddPostTaskOn(MTask *t)                    { Add(fPostTasksOn,  t); }
    void AddPostTaskOn(const char *rule,
                       const char *name="MWeight")  { AddPar(fPostTasksOn, rule, name); }

    void AddPostTaskOff(MTask *t)                   { Add(fPostTasksOff, t); }
    void AddPostTaskOff(const char *rule,
                        const char *name="MWeight") { AddPar(fPostTasksOff, rule, name); }

    void SetWeightsOn(const char *rule)  { if (fEnableWeightsOn) return;  fEnableWeightsOn=kTRUE;  AddPostTaskOn(rule); }
    void SetWeightsOn(MTask *t)          { if (fEnableWeightsOn) return;  fEnableWeightsOn=kTRUE;  AddPostTaskOn(t);    }

    void SetWeightsOff(const char *rule) { if (fEnableWeightsOff) return; fEnableWeightsOff=kTRUE; AddPostTaskOff(rule); }
    void SetWeightsOff(MTask *t)         { if (fEnableWeightsOff) return; fEnableWeightsOff=kTRUE; AddPostTaskOff(t);    }

    void SetFlux(Float_t f) { fFlux = f; }

    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