//Based on original teststat code
//Ian Connelly
//21/May/12

// Take in argument two numbers
// Use this to allocate two TFiles
// Functions
// - One to initialise pdfs (argument TFiles)
// - One to generate single event - ie 1 Poisson and xi
// - One to evaluate these values
// -> Loop these two for number of events

#include <stdlib.h>
#include <iostream>
#include "TFile.h"
#include "RooRealVar.h"
#include "RooWorkspace.h"
#include "RooArgSet.h"
#include "RooPoisson.h"
#include "RooDataSet.h"
#include "RooKeysPdf.h"
#include "RooPlot.h"
#include "TMath.h"
#include "TH1.h"
#include "TString.h"
#include "RooHistPdf.h"

using namespace RooFit;
using std::cout;
using std::endl;
using std::cerr;
using std::map;

void Initialise(TFile* f1, TFile* f2, TString tstr1, TString tstr2);
void GenerateEvent();
void EvaluateEvent(int poisP1, int poisP2, int poisP3, RooDataSet* xiP1, RooDataSet* xiP2, RooDataSet* xiP3);

//Global Variables
//Pdfs
RooPoisson* pdfnHadGenP1;
RooPoisson* pdfnHadGenP2;
RooPoisson* pdfnHadGenP3;

RooKeysPdf* pdfXiGenP1;
RooKeysPdf* pdfXiGenP2;
RooKeysPdf* pdfXiGenP3;
RooKeysPdf* pdfXiComP1;
RooKeysPdf* pdfXiComP2;
RooKeysPdf* pdfXiComP3;

//RooRealVars
RooRealVar* xi;
RooRealVar* hpm;
RooRealVar* nHad;
RooRealVar* poisMeanGenP1;
RooRealVar* poisMeanGenP2;
RooRealVar* poisMeanGenP3;


//Histograms

//Others
int argv;
char ** argc;
double meanGenP1, meanGenP2, meanGenP3, meanComP1, meanComP2, meanComP3;

int main(int argv, char ** argc){
    
   // RooMsgService::instance().getStream(1).removeTopic(NumIntegration);
    
    int num1 = atoi(argc[1]);
    int num2 = atoi(argc[2]);
    //Initialise mappings
    map<int,TString> fileRef;
    fileRef[0]= "sig_pyth";
    fileRef[1]= "sig_herw";
    fileRef[2]= "bkg_zjj0";
    fileRef[3]= "bkg_zjj1";
    fileRef[4]= "bkg_tt";
    TString name1 = fileRef.find(num1)->second;
    TString name2 = fileRef.find(num2)->second;
    TString filePath = "/scratch3/connelly/testarea/analysis/kernalestimation/output/v2/";
    TString fileName1 = "workspace"+name1+".root";
    TString fileName2 = "workspace"+name2+".root";
    TFile* file1 = new TFile(filePath+fileName1,"READ");
    TFile* file2 = new TFile(filePath+fileName2,"READ");
    cout << "===> Files initalised." << endl;
    
    Initialise(file1,file2,name1,name2);
    int numEvents = 2;
    for (int i = 0;  i < numEvents; i++){
        GenerateEvent();
    }
}

void Initialise(TFile* f1, TFile* f2, TString tstr1, TString tstr2){
    //Do it so will generate from f1 and compare with f2
    //Hence when generating under the other hypothesis, switch in function call
    RooWorkspace* w1 = (RooWorkspace*)f1->Get("w");
    RooWorkspace* w2 = (RooWorkspace*)f2->Get("w");
    nHad = (RooRealVar*)w1->var("nHadrons");
    xi   = (RooRealVar*)w1->var("xi");
    hpm  = (RooRealVar*)w1->var("hadronPlaneMomentum");

    meanGenP1 = w1->data("data_"+tstr1+"_varExcP1")->mean(*nHad);
    meanGenP2 = w1->data("data_"+tstr1+"_varExcP2")->mean(*nHad);
    meanGenP3 = w1->data("data_"+tstr1+"_varExcP3")->mean(*nHad);
    
    cout << "===> Generating means : " << meanGenP1 << " " << meanGenP2 << " " << meanGenP3 << endl;
    
    meanComP1 = w2->data("data_"+tstr2+"_varExcP1")->mean(*nHad);
    meanComP2 = w2->data("data_"+tstr2+"_varExcP2")->mean(*nHad);
    meanComP3 = w2->data("data_"+tstr2+"_varExcP3")->mean(*nHad);
    
    cout << "===> Comparison means : " << meanComP1 << " " << meanComP2 << " " << meanComP3 << endl;

    //Need pointers here to persist to next function
    poisMeanGenP1 = new RooRealVar("poisMeanGenP1","Mean Gen P1",meanGenP1,"");
    poisMeanGenP2 = new RooRealVar("poisMeanGenP2","Mean Gen P2",meanGenP2,"");
    poisMeanGenP3 = new RooRealVar("poisMeanGenP3","Mean Gen P3",meanGenP3,"");
    
    pdfnHadGenP1 = new RooPoisson("pdfnHadGenP1","Poisson Gen P1",*nHad,*poisMeanGenP1);
    pdfnHadGenP2 = new RooPoisson("pdfnHadGenP2","Poisson Gen P2",*nHad,*poisMeanGenP2);
    pdfnHadGenP3 = new RooPoisson("pdfnHadGenP3","Poisson Gen P3",*nHad,*poisMeanGenP3);

    pdfXiGenP1 = (RooKeysPdf*)w1->pdf("pdfkeys_"+tstr1+"_xi_P1");
    pdfXiGenP2 = (RooKeysPdf*)w1->pdf("pdfkeys_"+tstr1+"_xi_P2");
    pdfXiGenP3 = (RooKeysPdf*)w1->pdf("pdfkeys_"+tstr1+"_xi_P3");
    
    pdfXiComP1 = (RooKeysPdf*)w2->pdf("pdfkeys_"+tstr2+"_xi_P1");
    pdfXiComP2 = (RooKeysPdf*)w2->pdf("pdfkeys_"+tstr2+"_xi_P2");
    pdfXiComP3 = (RooKeysPdf*)w2->pdf("pdfkeys_"+tstr2+"_xi_P3");
    
    cout << "===> Generation and comparison pdfs are initialised." << endl;
    
}

void GenerateEvent(){
    //Only generate a single event here
    //Try to avoid pointers - NB RooDataSet HAS to be pointer else seg fault on repeated generation
    //Could pass by reference to EvaluateEvent within this function
    RooDataSet* gennHadP1 = new RooDataSet("gennHadP1","Gen nHad P1",RooArgSet(*nHad));
    RooDataSet* gennHadP2 = new RooDataSet("gennHadP2","Gen nHad P2",RooArgSet(*nHad));
    RooDataSet* gennHadP3 = new RooDataSet("gennHadP3","Gen nHad P3",RooArgSet(*nHad));
    
    //Generate single event - number of hadrons
    gennHadP1 = (RooDataSet*)pdfnHadGenP1->generate(*nHad,1);
    gennHadP2 = (RooDataSet*)pdfnHadGenP2->generate(*nHad,1);
    gennHadP3 = (RooDataSet*)pdfnHadGenP3->generate(*nHad,1);
    
    //Need to get the get then number of hadrons generated
    RooArgSet setHadP1 = *gennHadP1->get();
    nHad = (RooRealVar*)setHadP1.find(nHad->GetName());
    int genHadP1 = static_cast<int>(nHad->getVal());
    
    RooArgSet setHadP2 = *gennHadP2->get();
    nHad = (RooRealVar*)setHadP2.find(nHad->GetName());
    int genHadP2 = static_cast<int>(nHad->getVal());

    RooArgSet setHadP3 = *gennHadP3->get();
    nHad = (RooRealVar*)setHadP3.find(nHad->GetName());
    int genHadP3 = static_cast<int>(nHad->getVal());

    cout << genHadP1 << "  " << genHadP2 << "  " << genHadP3 << endl; 
    
    RooDataSet* genXiP1 = new RooDataSet("genXiP1","Gen Xi P1",RooArgSet(*xi));
    RooDataSet* genXiP2 = new RooDataSet("genXiP2","Gen Xi P2",RooArgSet(*xi));
    RooDataSet* genXiP3 = new RooDataSet("genXiP3","Gen Xi P3",RooArgSet(*xi));

    genXiP1 = (RooDataSet*)pdfXiGenP1->generate(*xi,genHadP1);
    genXiP2 = (RooDataSet*)pdfXiGenP2->generate(*xi,genHadP2);
    genXiP3 = (RooDataSet*)pdfXiGenP3->generate(*xi,genHadP3);
    
    cout << "===>Event generated." << endl;
    
    EvaluateEvent(genHadP1,genHadP2,genHadP3,genXiP1,genXiP2,genXiP3);


}

void EvaluateEvent(int poisP1, int poisP2, int poisP3, RooDataSet* xiP1, RooDataSet* xiP2, RooDataSet* xiP3){
    //Now evaluate these values using the generation and comparison pdfs
    //Maybe change from void to double and return the llratio and then change the generate event to do the same?
    // or change generate event to fill histo, or evaluate to fill histo?
    // or a filling function called within evaluate at the end?
    
    //Easy part - Poisson probabilities
    double probnHadGenP1 = TMath::Poisson(poisP1, meanGenP1);
    double probnHadGenP2 = TMath::Poisson(poisP2, meanGenP2);
    double probnHadGenP3 = TMath::Poisson(poisP3, meanGenP3);
    
    double probnHadComP1 = TMath::Poisson(poisP1, meanComP1);
    double probnHadComP2 = TMath::Poisson(poisP2, meanComP2);
    double probnHadComP3 = TMath::Poisson(poisP3, meanComP3);

    //Harder part - Product of probs
    double probXiGenP1;
    
    //Plane 1
    RooArgSet* setPdfXiGenP1 = pdfXiGenP1->getObservables(xiP1);
    RooArgSet* setPdfXiComP1 = pdfXiComP1->getObservables(xiP1);
    RooArgSet* setXiP1;
    cout << "Int" << pdfXiGenP1->createIntegral(*xi)->getVal() << endl;
    for (int j = 0; j < poisP1; j++){
        //setXiP1 = (RooArgSet*)xiP1->get(j);
        //xi = (RooRealVar*)setXiP1->find(xi->GetName());
        
        *setPdfXiGenP1 = *xiP1->get(j);
        probXiGenP1 = pdfXiGenP1->getVal(setPdfXiGenP1); 
        
        cout << xi->getVal() << " " << probXiGenP1 << endl;
        
    }
    
    for (int j = 0; j < poisP2; j++){
        
    }
    
    for (int j = 0; j < poisP3; j++){
        
    }
}