//
// taken from Wouter Verkerke's advanced statistics course
//
#include <iostream> 

#include "TFile.h"
#include "TTree.h"
#include "TH2F.h"
#include "TPad.h"
#include "TCut.h"
#include "TF1.h"

#if not defined(__CINT__) || defined(__MAKECINT__)
#include "TMVA/Factory.h"
#include "TMVA/Tools.h"
#endif

#ifndef __CINT__
#include "RooGlobalFunc.h"
#endif
#include "RooWorkspace.h"
#include "RooDataSet.h"
#include "RooAbsPdf.h"
#include "RooTreeDataStore.h"

using namespace RooFit;
using namespace std;

void classification()
{
  // plot things
  TFile* fin = TFile::Open( "sample.root" );
  TTree* tree_sig = ( TTree* ) fin->Get( "tree_sig" );
  TTree* tree_bkg = ( TTree* ) fin->Get( "tree_bkg" );

  TH2F* h = new TH2F( "h", "Classification", 10, -10, 10, 10, -10, 10 );
  h->SetStats( kFALSE );
  h->Draw();

  tree_sig->SetMarkerColor( kRed );
  tree_sig->Draw( "y:x", "", "same" );

  tree_bkg->SetMarkerColor( kBlue );
  tree_bkg->Draw( "y:x", "", "same" );

  TF1* f = new TF1( "f", "-x + 7.46268656716417844e-02", -10, 10 );
  f->Draw( "same" );

  gPad->SaveAs( "classification.pdf" );
  //  return;

  // load TMVA library
  TFile* fout = TFile::Open( "TMVA.root", "recreate" );

  TMVA::Tools::Instance();
  TMVA::Factory* factory = 
    new TMVA::Factory( "classification", 
		       fout,
		       "!V:"
		       "!Silent:"
		       "Color:"
		       "DrawProgressBar:"
		       "Transformations=I;D;P;G,D" );

  factory->AddVariable( "x", 'F' );
  factory->AddVariable( "y", 'F' );

  factory->AddSignalTree    ( tree_sig, 1.0 );
  factory->AddBackgroundTree( tree_bkg, 1.0 );

  TCut mycut_sig = "";
  TCut mycut_bkg = "";
  factory->PrepareTrainingAndTestTree( mycut_sig, 
				       mycut_bkg,
				       "nTrain_Signal=0:"
				       "nTrain_Background=0:"
				       "SplitMode=Random:"
				       "NormMode=NumEvents:!V" );

  // cut optimization
  // factory->BookMethod( TMVA::Types::kCuts, 
  // 		       "Cuts", 
  // 		       "!H:"
  // 		       "!V:"
  // 		       "FitMethod=MC:"
  // 		       "EffSel:"
  // 		       "SampleSize=200000:"
  // 		       "VarProp=FSmart" );

  // Fisher discriminant
  factory->BookMethod( TMVA::Types::kFisher, 
		       "Fisher", 
		       "H:"
		       "!V:"
		       "Fisher:"
		       "CreateMVAPdfs:"
		       "PDFInterpolMVAPdf=Spline2:"
		       "NbinsMVAPdf=60:"
		       "NsmoothMVAPdf=10" );

  // ANN: multi-layer perceptron
  factory->BookMethod( TMVA::Types::kMLP, 
		       "MLP", 
		       "H:"
		       "!V:"
		       "NeuronType=sigmoid:"
		       "VarTransform=N:"
		       "NCycles=600:"
		       "HiddenLayers=N+5:"
		       "TestRate=5" );

   factory->TrainAllMethodsForClassification();
   factory->TestAllMethods();
   factory->EvaluateAllMethods();    

   fout->Close();
}

void makesample()
{
  // Factory for PDFs
  RooWorkspace w( "w", kTRUE );
  w.factory("Gaussian::sig(expr('sqrt(((x+2)*(x+2))+((y+2)*(y+2)))',x[-10,10],y[-10,10]),5,1)") ;
  w.factory("Gaussian::bkg(expr('sqrt(((x-2)*(x-2))+((y-2)*(y-2)))',x,y),5,1)") ;
  w.defineSet("obs","x,y") ;

  // Generate events
  RooDataSet* data_sig = w.pdf("sig")->generate( *w.set("obs"), 10000 ) ;
  RooDataSet* data_bkg = w.pdf("bkg")->generate( *w.set("obs"), 10000 ) ;   

  TFile fout( "sample.root", "recreate" );
  RooTreeDataStore* sigstore = (RooTreeDataStore*) data_sig->store();
  RooTreeDataStore* bkgstore = (RooTreeDataStore*) data_bkg->store();
  sigstore->tree().Clone()->Write("tree_sig") ;
  bkgstore->tree().Clone()->Write("tree_bkg") ;
  fout.Close() ;
  cout << "finished" << endl ; 
}