#include "TMVA/Factory.h"
#include "TMVA/Tools.h"
#include <iostream>
using namespace TMVA;

int nbin = 100;
bool iter = 0;
int colors_val[4] = {kRed, kGreen, kMagenta, kYellow};
float elayer_param[10];

bool isValidPid(int i)
{
  return (i == 1) || (i == 2);
}

void process(TTree *tin, Reader* reader, const char* method, const char* weights, TPad* canvas, TLegend* legend)
{
  int pid;
  double esum;
  float elayer[10];
  tin->GetBranch("pid")->SetAddress(&pid);
  tin->GetBranch("esum")->SetAddress(&esum);
  tin->GetBranch("elayer")->SetAddress(&elayer);

  reader->BookMVA(method, weights);
  double start = -3.0;
  double end = 3.0;
  TH1D *bg = new TH1D("bg", (string(method) + std::string(" pid==1")).c_str(),nbin, start,end);
  TH1D *sig = new TH1D("sig",(string(method) + std::string(" pid==2")).c_str(),nbin, start,end);

  for( int i=0; i<tin->GetEntries(); i++ )
  {
    tin->GetEvent(i);
    if(!isValidPid(pid)) continue;
    for( int j=0; j<10; j++ ) {elayer_param[j]=elayer[j];}
    double out = reader->EvaluateMVA(method);
    if( pid==1 ) {bg->Fill(out);}
    else if( pid==2 ) {sig->Fill(out);}
  }


  bg->SetLineColor(kBlue);
  bg->SetFillStyle(3008);
  bg->SetFillColor(kBlue);

  sig->SetLineColor(kRed);
  sig->SetFillStyle(3003);
  sig->SetFillColor(kRed);

  bg->SetStats(0);
  sig->SetStats(0);

  bg->Draw();
  sig->Draw("same");

  canvas->cd(3);
  int nsig = sig->GetEntries();
  int nbg = bg->GetEntries();
  double *a = new double[nbin];
  double *b = new double[nbin];
  int csig = 0, cbg = 0;
  for (int i = 0; i < nbin; ++i)
  {
    csig += sig->GetBinContent(i);
    cbg  += bg->GetBinContent(i);
    a[i] = double(nsig - csig) / nsig;
    b[i] = double(cbg) / nbg;
  }
  TGraph *graph = new TGraph(100, b, a);
  graph->SetLineWidth(2);
  graph->SetLineColor(colors_val[iter]);
  graph->Draw(iter == 0 ? "ACP" : "CP");
  ++iter;
  legend->AddEntry(graph, method, "l");
}

void task5read()
{
  TPad* canvas = new TCanvas("canvas","",1000,700);
  canvas->Divide(3,1);

  gSystem->Load("libTMVA");

  Reader* reader = new Reader("Color");

  TFile *fin = new TFile("tree.root");
  TTree *tin = NULL;
  fin->GetObject("tree", tin);

  for( int i=0; i<10; i++ )
    reader->AddVariable(Form("elayer[%d]",i),  elayer_param + i);

  auto legend = new TLegend(0.6, 0.6, 0.9, 0.9);
  canvas->cd(1);
  process(tin, reader, "Likelihood", "weights/myTMVA_Likelihood.weights.xml", canvas, legend);
  canvas->cd(2);
  process(tin, reader, "BDT","weights/myTMVA_BDT.weights.xml", canvas, legend);
  legend->Draw();
}

#include "TMVA/Factory.h"
#include "TMVA/Tools.h"

void task5() {

  gSystem->Load("libTMVA");

  TMVA::Tools::Instance();

  TFile* fout = TFile::Open("tmva_out.root", "RECREATE");
  TChain *tin = new TChain("tree");
  tin->AddFile("tree.root");
  TMVA::Factory *factory = new TMVA::Factory("myTMVA", fout, "!V:!Silent:Color:DrawProgressBar:Transformations=I;D;P;G,D:AnalysisType=Classification");
  factory->SetInputTrees(tin,"pid==1","pid==2");

  factory->AddVariable("el0 := elayer[0]",'D');
  factory->AddVariable("el1 := elayer[1]",'D');
  factory->AddVariable("el2 := elayer[2]",'D');
  factory->AddVariable("el3 := elayer[3]",'D');
  factory->AddVariable("el4 := elayer[4]",'D');
  factory->AddVariable("el5 := elayer[5]",'D');
  factory->AddVariable("el6 := elayer[6]",'D');
  factory->AddVariable("el7 := elayer[7]",'D');
  factory->AddVariable("el8 := elayer[8]",'D');
  factory->AddVariable("el9 := elayer[9]",'D');


  factory->PrepareTrainingAndTestTree("","");

  factory->BookMethod(TMVA::Types::kLikelihood, "Likelihood", "H:!V:TransformOutput");
  factory->BookMethod(TMVA::Types::kBDT, "BDT", "H:!V:NTrees=200:nEventsMin=100:BoostType=AdaBoost");

  factory->TrainAllMethods();

  factory->TestAllMethods();

  factory->EvaluateAllMethods();
}