// author: C. Delaere
// Exercise on Neural Networks.
// Note also the examples in root tutorials: root/tutorials/mlp

void NNsimple() 
{
  // we will look for 2 variables x,y
  // x,y are in [0,1]
  // the signal region is defined as (x>1/3 && y>2/3)||(x>2/3 && y>1/3)
  // the background region is defined as y < 4/3-x
  
  // Fill signal tree
  TTree* tree = new TTree("MonteCarlo", "Filtered Monte Carlo Events");
  Float_t x = -1;
  Float_t y = -1;
  Int_t signal = 1;
  tree->Branch("x",&x,"x/F");
  tree->Branch("y",&y,"y/F");
  tree->Branch("signal",&signal,"signal/I");
  // TODO: you way want to change the number of signal and background event that you generate
  for(int i=0;i<10000;++i) {
    // one signal event ...
    x=-1; y=-1; signal=1;
    while(!((x>1/3. && y>2/3.)||(x>2/3. && y>1/3.))) {
      x = gRandom->Uniform();
      y = gRandom->Uniform();
    }
    tree->Fill();
  }
  for(int i=0;i<10000;++i) {
    // one background event ...
    x=1; y=1; signal = 0;
    while(!(y < 4/3.-x)) {
      x = gRandom->Uniform();
      y = gRandom->Uniform();
    }
    tree->Fill();
  }

  // Build and train the network
  // TODO: here, you have to set the proper structure
  TMultiLayerPerceptron *mlp = new TMultiLayerPerceptron("x,y:5:3:signal",tree);
  // TODO: change the number of iterations in the training (here 1000)
  mlp->Train(1000, "text,graph,update=10");

  // Result canvas
  TCanvas* mlpa_canvas = new TCanvas("mlpa_canvas","Network analysis");
  mlpa_canvas->Divide(1,2);
  TMLPAnalyzer ana(mlp);
  ana.GatherInformations();
  ana.CheckNetwork();
  mlpa_canvas->cd(1);
  mlp->Draw();
  mlpa_canvas->cd(2);
  ana.DrawNetwork(0,"signal==1","signal==0"); // that draws the result
  // retrieve the histograms for later use
  THStack* stack = mlpa_canvas->FindObject("__NNout_TMLPA");
  TH1F* b = (TH1F*)(stack->GetHists()->At(0));
  TH1F* s = (TH1F*)(stack->GetHists()->At(1));

  // this is an example to show how to use the NN later on: we plot the NN in 2D
  Double_t vx[121];
  Double_t vy[121];
  Double_t  f[121];
  Double_t  v[2];
  // we simply run on a 2D grid and store x, y, NN(x,y).
  for (Int_t ix=0; ix<=10; ix++) {
     v[0]=ix/10.;
     for (Int_t iy=0; iy<=10; iy++) {
        v[1]=iy/10.;
        Int_t idx=ix*11+iy;
        vx[idx]=v[0];
        vy[idx]=v[1];
        f[idx]=mlp->Evaluate(0, v); // this is how we evaluate the NN at point v={x,y}...
     }
  }
  // now we can create a TGraph2D an draw it !
  TCanvas* ana_canvas = new TCanvas("analysis_canvas","Analysis optimization");
  ana_canvas->Divide(2,2);
  ana_canvas->cd(1);
  TGraph2D* NN = new TGraph2D("NetworkOutput","Network Output",121, vx, vy, f);
  NN->Draw("cont4z");

  // purity and efficiency for a cut x>...
  ana_canvas->cd(2);
  int nbins = b->GetNbinsX();
  unsigned int count_s = s->GetEntries();
  unsigned int count_b = b->GetEntries();

  float* efficiency = new float[nbins];
  float* purity = new float[nbins];
  float* bkgfraction = new float[nbins];
  float* merit1 = new float[nbins];
  float* merit2 = new float[nbins];

  for(int i=1;i<=nbins;++i) {
    count_s -= s->GetBinContent(i);
    count_b -= b->GetBinContent(i);
    efficiency[i-1] = float(count_s)/float(s->GetEntries());
    bkgfraction[i-1] = float(count_b)/float(b->GetEntries());
    purity[i-1] = count_b>0 ? float(count_s)/float(count_s+count_b) : 1.;
    merit1[i-1] = count_b>0 ? count_s/sqrt(count_s+count_b) : 1.;
    merit2[i-1] = count_b>0 ? count_s/sqrt(count_b) : 1.;
  }
  TGraph* interest = new TGraph(nbins,bkgfraction,efficiency);
  interest->Draw("al");
  interest->GetXaxis()->SetTitle("Background rejection");
  interest->GetYaxis()->SetTitle("Signal efficiency");
  interest->SetTitle("Interest");

  // figure of merit
  float* cut = new float[nbins];
  for(int i=1;i<=nbins;++i) {
    cut[i-1] = s->GetBinLowEdge(i)+s->GetBinWidth(i);
  }
  TGraph* fmerit1 = new TGraph(nbins,cut,merit1);
  fmerit1->SetTitle("Figure of merit for measurement");
  fmerit1->SetLineColor(kBlue);
  TGraph* fmerit2 = new TGraph(nbins,cut,merit2);
  fmerit2->SetTitle("Figure of merit for discovery");
  fmerit2->SetLineColor(kRed);
  ana_canvas->cd(3);
  fmerit1->Draw("al");
  fmerit1->GetXaxis()->SetTitle("NN cut");
  ana_canvas->cd(4);
  fmerit2->Draw("al");
  fmerit2->GetXaxis()->SetTitle("NN cut");
  
  // add markers on the interest plot
  float bestCut1 = fmerit1->GetX()[TMath::LocMax(fmerit1->GetN(),fmerit1->GetY())];
  float bestCut2 = fmerit2->GetX()[TMath::LocMax(fmerit2->GetN(),fmerit2->GetY())];
  s->FindBin(bestCut1);
  TMarker* star1 = new TMarker(interest->GetX()[s->FindBin(bestCut1)-1],
                               interest->GetY()[s->FindBin(bestCut1)-1],29);
  star1->SetMarkerColor(kBlue);
  star1->SetMarkerSize(2);
  TMarker* star2 = new TMarker(interest->GetX()[s->FindBin(bestCut2)-1],
                               interest->GetY()[s->FindBin(bestCut2)-1],29);
  star2->SetMarkerColor(kRed);
  star2->SetMarkerSize(2);
  ana_canvas->cd(2);
  star1->Draw("same");
  star2->Draw("same");

}