#include <algorithm>
#include "TTree.h"
#include "TCanvas.h"
#include "TH1.h"
#include "TLegend.h"
#include "TFile.h"
#include "TSystem.h"
#include "TMultiLayerPerceptron.h"
#include "TMLPAnalyzer.h"
   
//Declaration of leaves types
typedef struct { 
   Int_t           type;
   Float_t         jt1_et;
   Float_t         jt2_et;
   Float_t         jt3_et;
   Float_t         jt4_et;
   Float_t         ht;
   Float_t         ht_3b;
   Float_t         pt_b1b2;
   Float_t         m12;
   Float_t         m23;
   Float_t         m13;
   Float_t         m14;
   Float_t         m24;
   Float_t         m34;
   Float_t         dr12;
   Float_t         dr23;
   Float_t         dr13;
   Float_t         dr14;
   Float_t         dr24;
   Float_t         dr34;
   Float_t         drmin;
   Float_t         dphi12;
   Float_t         dphi23;
   Float_t         dphi13;
   Float_t         dphi14;
   Float_t         dphi24;
   Float_t         dphi34;
   Float_t         dphimin;
   Float_t         weight;
} NN_t;

void nn(int ntrain=220)
{
    if (!gROOT->GetClass("TMultiLayerPerceptron")) {
        gSystem->Load("libMLP");
    }

    TFile *f = new TFile("nn.root");
    TTree *snb = (TTree*) gDirectory->Get("NN");
//  # of input node = 7 (# of variables)
//  # of first hidden layer = 13
//  # of second hidden layer = 4
//  # of output layer = 1 (variable "type")
//    TMultiLayerPerceptron *mlp = new TMultiLayerPerceptron("dphi12,dphi13,dphi23,ht,m12,m23,pt_b1b2:13:4:@type","weight",snb,"Entry$%2","Entry$/2");
    TMultiLayerPerceptron *mlp = new TMultiLayerPerceptron("dphi12,dphi13,dphi23,ht,m12,m23,pt_b1b2:13:4:type","weight",snb);
// Alternatively, dr12,dr13,dr23,ht_3b can be used.
    mlp->SetLearningMethod(TMultiLayerPerceptron::kBFGS);
    mlp->Train(ntrain,"text,graph,update=10");

    TMLPAnalyzer ana(mlp);
    // Initialisation
    ana.GatherInformations();
    // output to the console
    ana.CheckNetwork();

    TCanvas* c1 = new TCanvas("c1","Network Analysis",1600,1600);
    // shows how each variable influences the network
    ana.DrawDInputs();
    c1->Print("nn_analysis.eps");

    TCanvas* c2 = new TCanvas("c2","Network Structure",2500,1600);
    // shows the network structure
    mlp->Draw();
    c2->Print("nn_structure.eps");

    TCanvas* c3 = new TCanvas("c3","Sig/Bkgd Separation 1",1600,1600);
    // draws the resulting network
    ana.DrawNetwork(0,"type==1","type==0");
    c3->Print("nn_snb1.eps");

    TCanvas* c4 = new TCanvas("c4","Sig/Bkgd Separation 2",1600,1600);
    // Use the NN to plot the results for each sample
    // This will give approx. the same result as DrawNetwork.
    // All entries are used, while DrawNetwork focuses on 
    // the test sample. Also the xaxis range is manually set.
    TH1F *bg = new TH1F("bg", "NN output", 50, -.5, 1.5);
    TH1F *sig = new TH1F("sig", "NN output", 50, -.5, 1.5);

    NN_t NN;
    Double_t params[7];
    snb->SetBranchAddress("NN",&NN.type);
    for ( int i=0; i<snb->GetEntries(); i++) {
        snb->GetEntry(i);
        params[0] = NN.dphi12;
        params[1] = NN.dphi13;
        params[2] = NN.dphi23;
        params[3] = NN.ht;
        params[4] = NN.m12;
        params[5] = NN.m23;
        params[6] = NN.pt_b1b2;
        if ( NN.type==0 )
            bg->Fill(mlp->Evaluate(0, params));
        else if ( NN.type==1 )
            sig->Fill(mlp->Evaluate(0,params));
    }
    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->Scale(1./bg->GetEntries());
    sig->Scale(1./sig->GetEntries());
    bg->Draw();
    sig->Draw("same");
    TLegend *legend = new TLegend(.75, .80, .95, .95);
    legend->AddEntry(bg, "Background (Fakes)");
    legend->AddEntry(sig, "Signal (Higgs)");
    legend->Draw();
    c4->Print("nn_snb2.eps");
}