// 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;

int main(int argv, char ** argc){
    
    //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
    
    
    
    
    
    
    return 0;
}