// ke_plane.cxx
// kernal estimation to define plane in TH2 histograms
// Ian Connelly - Friday 13th April 2012!
//

#include "RooDataHist.h"
#include "RooDataSet.h"
#include "RooArgList.h"
#include "RooAbsArg.h"
#include "RooRealVar.h"
#include "RooArgSet.h"
#include "RooKeysPdf.h"
#include "RooNDKeysPdf.h"
#include "Roo2DKeysPdf.h"
#include "RooPlot.h"
#include <TFile.h>
#include <TTree.h>
#include <TH1.h>
#include <TH2.h>
#include <TString.h>
#include <TCanvas.h>
#include <vector>
#include <sstream>
#include <stdlib.h>
#include <map>
#include "RooWorkspace.h"

using namespace RooFit;
using std::cout;
using std::endl;
using std::cin;
using std::vector;

//Global Vars
int argv;
char ** argc;

void originalProgram();
void newProgram();
void makeWorkspace(RooWorkspace* w, TFile* f, TString tstr);

int main(int argv, char ** argc){
    
    //originalProgram();
    newProgram();
    return 0;
}

void originalProgram(){
    
    //Initialisation
    
    //Pass a number to the program depending on which file is being analysed and then ensure the mapping works correctly
    //ascii to int
    int fileReference = atoi(argc[1]);
    
    TString filename;
    vector<TFile*> files;
    vector<bool> signalFlag;
    vector<RooDataSet> data_incP1;
    vector<RooDataSet> data_excP1;
    vector<RooDataSet> data_incP2;
    vector<RooDataSet> data_excP2;
    vector<RooDataSet> data_incP3;
    vector<RooDataSet> data_excP3;
    
    vector<vector<vector<RooDataSet> > > data;
    vector<vector<RooDataSet> > data_file;
    vector<RooDataSet> data_inc;
    vector<RooDataSet> data_exc;

    RooWorkspace* w = new RooWorkspace("w","workspace");
    
    map<int,TString> label;
    label[0]="sig_pyth";
    label[1]="sig_herw";
    label[2]="bkg_zjj0";
    label[3]="bkg_zjj1";
    label[4]="bkg_tt";
    
    cout << "Using : " << fileReference << " : " << label.find(fileReference)->second << endl;
    
    map<int,TString> plane;
    plane[0]="P1";
    plane[1]="P2";
    plane[2]="P3";

    
    //Use the prog.exe < input syntax to use a list of TFiles which I wish to use.
    //Should get around having to read a directory or type a filename in each time, whilst ensuring continued flexibility of program.
    
    //Import TFiles

    cout << "Ideally run this program using an input file of the following structure : " << endl;
    cout << "> filepath.root" << endl;
    cout << "> y/n (depending on if the root file above it is a signal or background file)"<< endl;
    cout << "Continue this syntax until all files are listed in input file and exectute as : ./kernalestimator.exe < input" << endl;
    cout << "Alternatively type the above syntax directly into terminal and ctrl+d when finished entering." << endl;
    cout << "===================================================" << endl;

    
    while(cin >> filename){
        bool signal = 0;
        if(filename == "y"){
            signal = 1;
            signalFlag.push_back(signal);
            continue; // The two continues ensures that the while will loop to the next input if it is a flag and not filename entered.
        }
        if(filename == "n"){
            signalFlag.push_back(signal);
            continue;
        }
        TFile* tempfile = new TFile(filename, "READ");
        files.push_back(tempfile);
    }
    
    cout << "Number of files loaded is : " << files.size() << endl;
    cout << "Files loaded are : " << endl;
    for(int i=0; i < files.size(); i++){
        files.at(i)->Print();
        cout << signalFlag.at(i) << endl;
        //files.at(i)->ls();
    }
    cout << "===================================================" << endl;

    //Get the TTree histograms
    //Define RooDataSet
    //Variables which are used - Have to provide min/max values
    RooRealVar xi("xi", "xi", -50, 50,"");
    RooRealVar hadronPlaneMomentum("hadronPlaneMomentum","hadronPlaneMomentum",0,10000000,"MeV"); //Seems to be one very high projected momenta particle??
    RooRealVar nHadrons("nHadrons","nHadrons",0,100,"");


    for(int i=0; i < files.size(); i++){
        TTree* temptree1 = (TTree*)files.at(i)->Get("t_p1VarInc");
        TTree* temptree2 = (TTree*)files.at(i)->Get("t_p2VarInc");
        TTree* temptree3 = (TTree*)files.at(i)->Get("t_p3VarInc");
        TTree* temptree4 = (TTree*)files.at(i)->Get("t_p1VarExc");
        TTree* temptree5 = (TTree*)files.at(i)->Get("t_p2VarExc");
        TTree* temptree6 = (TTree*)files.at(i)->Get("t_p3VarExc");
        
        //TString name = label.find(i)->second;
        TString name = label.find(fileReference)->second;
        TString name_var1 = "data_"+name+"_varIncP1";
        TString name_var2 = "data_"+name+"_varIncP2";
        TString name_var3 = "data_"+name+"_varIncP3";
        TString name_var4 = "data_"+name+"_varExcP1";
        TString name_var5 = "data_"+name+"_varExcP2";
        TString name_var6 = "data_"+name+"_varExcP3";

        
        
        RooDataSet varIncP1 (name_var1,"Inclusive variables from Plane 1", temptree1, RooArgSet(xi,hadronPlaneMomentum));
        RooDataSet varExcP1 (name_var4,"Exclusive variables from Plane 1", temptree4, RooArgSet(nHadrons));
        RooDataSet varIncP2 (name_var2,"Inclusive variables from Plane 2", temptree2, RooArgSet(xi,hadronPlaneMomentum));
        RooDataSet varExcP2 (name_var5,"Exclusive variables from Plane 2", temptree5, RooArgSet(nHadrons));
        RooDataSet varIncP3 (name_var3,"Inclusive variables from Plane 3", temptree3, RooArgSet(xi,hadronPlaneMomentum));
        RooDataSet varExcP3 (name_var6,"Exclusive variables from Plane 3", temptree6, RooArgSet(nHadrons));
        
        data_incP1.push_back(varIncP1);
        data_excP1.push_back(varExcP1);
        data_incP2.push_back(varIncP2);
        data_excP2.push_back(varExcP2);
        data_incP3.push_back(varIncP3);
        data_excP3.push_back(varExcP3);
        
        data_inc.push_back(varIncP1);
        data_inc.push_back(varIncP2);
        data_inc.push_back(varIncP3);
        data_exc.push_back(varExcP1);
        data_exc.push_back(varExcP2);
        data_exc.push_back(varExcP3);

        data_file.push_back(data_inc);
        data_file.push_back(data_exc);

        data.push_back(data_file);
    }
    
    //I know RooNDKeysPdf is the primary one, but would maybe the 2D one would be quicker?
    
    //RooPlot* plot = new RooPlot(xi,-20,20,100);
    
    for(int i=0; i < files.size(); i++){
        // RooNDKeysPdf* keyspdf_incP1 = new RooNDKeysPdf("keyspdf_incP1","PDF derived from kernal estimation - Inclusive Plane 1",RooArgList(xi, hadronPlaneMomentum), data_incP1.at(i),"av");
        // Roo2DKeysPdf keyspdf_incP1 ("keyspdf_incP1","PDF",xi, hadronPlaneMomentum, data_incP1.at(i), "av");
        //w->import(keyspdf_incP1);
        
        //Plane loop
        for(int j=0;j <3; j++){
            //TString name = label.find(i)->second;
            TString name = label.find(fileReference)->second;
            TString planename = plane.find(j)->second;
            TString pdfname1 = "pdfkeys_"+name+"_xi_"+planename;
            TString pdfname2 = "pdfkeys_"+name+"_hpm_"+planename;
            TString pdfname3 = "pdfkeys_"+name+"_nHad_"+planename;
            TString dataname = "data_inc"+planename;
            cout << "Generating " << pdfname1 << endl;
            RooKeysPdf keys1 (pdfname1,"keypdf xi",xi,data.at(i).at(0).at(j));
            cout << "Generating " << pdfname2 << endl;
            RooKeysPdf keys2 (pdfname2,"keypdf hadmtm", hadronPlaneMomentum, data.at(i).at(0).at(j));
            cout << "Generating " << pdfname3 << endl;
            RooKeysPdf keys3(pdfname3,"keypdf nHadrons", nHadrons, data.at(i).at(1).at(j));
            
            cout << "Importing pdfs and data to workspace." << endl;
            w->import(keys1);
            w->import(keys2);
            w->import(keys3);
        }
        //NOTE TO ACCESS THE DATASETS FIRST at IS THE FILE NUMBER SECOND IS THE INCLUSIVE OR EXCLUSIVE AND THIRD IS THE PLANE
        
        
        //Cuold just use the data vector here to access the data elements
        w->import(data_incP1.at(i));
        w->import(data_incP2.at(i));
        w->import(data_incP3.at(i));
        w->import(data_excP1.at(i));
        w->import(data_excP2.at(i));
        w->import(data_excP3.at(i));
        //data_incP1.at(i).plotOn(plot);
    }
    
    w->writeToFile("workspace.root");
    
    TFile* f = new TFile("f.root","RECREATE");
    //TCanvas* c = new TCanvas("c","c", 900,900);
    //plot->Draw();
    //c->Print("plot.ps");
    //c->Write();
    f->Write();
    f->Close();
    
    //Apply the Kernal Estimation function
    
}

void newProgram(){
    
    //This iteration of the program will use the newer analysis with slightly different ttree names and options
    //int fileReference = atoi(argc[1]);
    TString filename;

    vector<TFile*> files;
        
    map<int,TString> label;
    label[0]="sig_pyth";
    label[1]="sig_herw";
    label[2]="bkg_zjj0";
    label[3]="bkg_zjj1";
    label[4]="bkg_tt";
        
    while(cin >> filename){
        TFile* tempfile = new TFile(filename, "READ");
        files.push_back(tempfile);
    }
    
    cout << "Number of files loaded is : " << files.size() << endl;
    cout << "Files loaded are : " << endl;
    for(int i=0; i < files.size(); i++){
        files.at(i)->Print();
    }
    cout << "===================================================" << endl;

    for(int i = 0; i < files.size(); i++){
        TString file_name = label.find(i)->second;
        RooWorkspace* _w = new RooWorkspace("w","Workspace "+file_name); //Define Workspace
        TFile* _f = files.at(i); // Retrieve TFile
        makeWorkspace(_w,_f,file_name);
        delete _w;
    }
    
    
    
}

//Function to make the workspace for each file
void makeWorkspace(RooWorkspace* w, TFile* f, TString tstr){
    
    RooRealVar xi("xi", "xi", 0.1, 0.9,"");
    RooRealVar hadronPlaneMomentum("hadronPlaneMomentum","hadronPlaneMomentum",0,10000000,"MeV"); //Seems to be one very high projected momenta particle??
    RooRealVar nHadrons("nHadrons","nHadrons",0,100,"");
    
    //Check the correct trees are accessed in ttbar especially as they have "revision numbers" due to hadd
    
    TTree* temptree1 = (TTree*)f->Get("t_p1VarIncNew"); // These trees only record the xi,hpm for  xi in 0.1->0.9 like the nHadrons counts
    TTree* temptree2 = (TTree*)f->Get("t_p2VarIncNew");
    TTree* temptree3 = (TTree*)f->Get("t_p3VarIncNew");
    TTree* temptree4 = (TTree*)f->Get("t_p1VarExc");
    TTree* temptree5 = (TTree*)f->Get("t_p2VarExc");
    TTree* temptree6 = (TTree*)f->Get("t_p3VarExc");
    
    TTree* temptree7 = (TTree*)f->Get("t_p2VarIncNewIn");  // New trees
    TTree* temptree8 = (TTree*)f->Get("t_p2VarIncNewOut"); 
    
    TString name_var1 = "data_"+tstr+"_varIncP1";
    TString name_var2 = "data_"+tstr+"_varIncP2";
    TString name_var3 = "data_"+tstr+"_varIncP3";
    TString name_var4 = "data_"+tstr+"_varExcP1";
    TString name_var5 = "data_"+tstr+"_varExcP2";
    TString name_var6 = "data_"+tstr+"_varExcP3";
    TString name_var7 = "data_"+tstr+"_varIncP2In";
    TString name_var8 = "data_"+tstr+"_varIncP2Out";

    
    RooDataSet varIncP1 (name_var1,"Inclusive variables from Plane 1", temptree1, RooArgSet(xi,hadronPlaneMomentum));
    RooDataSet varExcP1 (name_var4,"Exclusive variables from Plane 1", temptree4, RooArgSet(nHadrons));
    RooDataSet varIncP2 (name_var2,"Inclusive variables from Plane 2", temptree2, RooArgSet(xi,hadronPlaneMomentum));
    RooDataSet varExcP2 (name_var5,"Exclusive variables from Plane 2", temptree5, RooArgSet(nHadrons));
    RooDataSet varIncP3 (name_var3,"Inclusive variables from Plane 3", temptree3, RooArgSet(xi,hadronPlaneMomentum));
    RooDataSet varExcP3 (name_var6,"Exclusive variables from Plane 3", temptree6, RooArgSet(nHadrons));
    RooDataSet varIncP2In  (name_var7,"Inclusive variables inside R=0.4 cone of Plane 2", temptree7, RooArgSet(xi, hadronPlaneMomentum));
    RooDataSet varIncP2Out (name_var8,"Inclusive variables outside R=0.4 cone of Plane 2", temptree8, RooArgSet(xi, hadronPlaneMomentum));
    
    //Need to save the datasets to file
    w->import(varExcP1);
    w->import(varExcP2);
    w->import(varExcP3);
    w->import(varIncP1);
    w->import(varIncP2);
    w->import(varIncP3);
    w->import(varIncP2In);
    w->import(varIncP2Out);
    
    vector<RooDataSet> _dataP1;
    vector<RooDataSet> _dataP2;
    vector<RooDataSet> _dataP3;
    vector<vector<RooDataSet> > _data;
    
    _dataP1.push_back(varIncP1); // at 0 - inclusive, at 1 - exclusive
    _dataP1.push_back(varExcP1);
    
    _dataP2.push_back(varIncP2); // at 0 - inclusive, at 1 - exclusive, 2 - in, 3 - out
    _dataP2.push_back(varExcP2);
    _dataP2.push_back(varIncP2In);
    _dataP2.push_back(varIncP2Out);
    
    _dataP3.push_back(varIncP3); // at 0 - inclusive, at 1 - exclusive
    _dataP3.push_back(varExcP3);
    
    _data.push_back(_dataP1); _data.push_back(_dataP2); _data.push_back(_dataP3);

    
    map<int,TString> plane;
    plane[0]="P1";
    plane[1]="P2";
    plane[2]="P3";
    
    //Plane loop
    for(int j=0;j <3; j++){
        TString planename = plane.find(j)->second;
        
        TString pdfname1 = "pdfkeys_"+tstr+"_xi_"+planename;
        TString pdfname2 = "pdfkeys_"+tstr+"_hpm_"+planename;
        TString pdfname3 = "pdfkeys_"+tstr+"_nHad_"+planename;
        
        TString pdfname4 = "pdfkeys_"+tstr+"_xi_"+planename+"In";
        TString pdfname5 = "pdfkeys_"+tstr+"_hpm_"+planename+"In";
        TString pdfname6 = "pdfkeys_"+tstr+"_xi_"+planename+"Out";
        TString pdfname7 = "pdfkeys_"+tstr+"_hpm_"+planename+"Out";

        TString dataname = "data_inc"+planename;
        
        cout << "Generating " << pdfname1 << endl;
        RooKeysPdf keys1 (pdfname1,"keypdf xi",xi,_data.at(j).at(0));
        cout << "Generating " << pdfname2 << endl;
        RooKeysPdf keys2 (pdfname2,"keypdf hadmtm", hadronPlaneMomentum, _data.at(j).at(0));
        cout << "Generating " << pdfname3 << endl;
        RooKeysPdf keys3(pdfname3,"keypdf nHadrons", nHadrons, _data.at(j).at(1));
        
        cout << "Importing pdfs and data to workspace." << endl;
        w->import(keys1);
        w->import(keys2);
        w->import(keys3);
        
        if(j == 1){
            //In plane 2 - generate pdf for in/out
            RooKeysPdf keys4 (pdfname4, "keypdf xi IN", xi, _data.at(1).at(2));
            RooKeysPdf keys5 (pdfname5, "keypdf hpm IN", hadronPlaneMomentum, _data.at(1).at(2));
            RooKeysPdf keys6 (pdfname6, "keypdf xi OUT", xi, _data.at(1).at(3));
            RooKeysPdf keys7 (pdfname7, "keypdf hpm OUT", hadronPlaneMomentum, _data.at(1).at(3));
            
            w->import(keys4);
            w->import(keys5);
            w->import(keys6);
            w->import(keys7);
        }
    }  
    
    
    
    w->writeToFile("output/v2/workspace"+tstr+".root");
    
}