#include "TCanvas.h"
#include "TNtuple.h"
#include "TRandom3.h"
#include "RooRealVar.h"
#include "RooGaussian.h"
#include "RooAddPdf.h"
#include "RooDataSet.h"
#include "RooChebychev.h"
#include "RooPlot.h"
#include <iostream>
using namespace RooFit;
TRandom3 psn;
TRandom3 rnd;
RooRealVar mu_pBrest("mu_pBrest","mu_pBrest",2.1,3.5);

//Read PDF
TFile *f1 = new TFile("conbbsigPDFmom_side.root");
RooWorkspace* w1 = (RooWorkspace*) f1->Get("w");
RooAbsPdf* sig1 = w1->pdf("Momsig");
RooAbsPdf* bb1 = w1->pdf("Mombb");
RooAbsPdf* ulnu1 = w1->pdf("Momulnu");
RooAbsPdf* con1 = w1->pdf("Momcon");
TFile *f2 = new TFile("conbbsigPDFmom.root");
RooWorkspace* w2 = (RooWorkspace*) f2->Get("w");
RooAbsPdf* sig2 = w2->pdf("Momsig");
RooAbsPdf* bb2 = w2->pdf("Mombb");
RooAbsPdf* ulnu2 = w2->pdf("Momulnu");
RooAbsPdf* con2 = w2->pdf("Momcon");

//Initial values of parameters of interest
const double sigScale_35 = 3.4994675552871413;
const double bbScale_35 = 72.03643845412496;
const double conScale_35 = 109.71811992503847;
const double ulnuScale_35 = 199.66866843592646;
const double sigScale1_35 = 1.5690308845937644;
const double bbScale1_35 = 45.265897732556304;
const double conScale1_35 = 80.17297578008359;
const double u0lnuScale1_35 = 62.64556114928302;
const double uplnuScale1_35 = 64.69958513979667;
const double ulnuScale1_35 = 127.34514628907968;
const double sigScale2_35 = 1.9304366706918128;
const double bbScale2_35 = 26.770540721568807;
const double conScale2_35 = 29.545144144954552;
const double u0lnuScale2_35 = 34.7662783010236;
const double uplnuScale2_35 = 37.55724384586087;
const double ulnuScale2_35 = 72.32352214688447;
const double sigScale1_35r = sigScale1_35/sigScale_35;
const double bbScale1_35r = bbScale1_35/bbScale_35;
const double conScale1_35r = conScale1_35/conScale_35;
const double ulnuScale1_35r = ulnuScale1_35/ulnuScale_35;
const double sigScale2_35r = sigScale2_35/sigScale_35;
const double bbScale2_35r = bbScale2_35/bbScale_35;
const double conScale2_35r = conScale2_35/conScale_35;
const double ulnuScale2_35r = ulnuScale2_35/ulnuScale_35;
void do_fit(float *res)
{

    std::cout<<"START"<<std::endl;

    //--------------Here is the simultaneous PDF model------------------
    RooCategory cat("cat","cat");
    cat.defineType("cat1",1);
    cat.defineType("cat2",2);
    RooRealVar signum("signum","signum",sigScale_35*1000,-300.,1000000.);
    RooRealVar bbnum("bbnum","bbnum",bbScale_35*1000,-300.,10000000.);
    RooRealVar ulnunum("ulnunum","ulnunum",ulnuScale_35*1000,-300.,20000000.);
    RooRealVar connum("connum","connum",conScale_35*1000,-300.,10000000.);
    RooRealVar signum1r("signum1r","signum1r",sigScale1_35r);
    RooRealVar bbnum1r("bbnum1r","bbnum1r",bbScale1_35r);
    RooRealVar ulnunum1r("ulnunum1r","ulnunum1r",ulnuScale1_35r);
    RooRealVar connum1r("connum1r","connum1r",conScale1_35r);
    RooRealVar signum2r("signum2r","signum2r",sigScale2_35r);
    RooRealVar bbnum2r("bbnum2r","bbnum2r",bbScale2_35r);
    RooRealVar ulnunum2r("ulnunum2r","ulnunum2r",ulnuScale2_35r);
    RooRealVar connum2r("connum2r","connum2r",conScale2_35r);
    RooProduct signum1("signum1","signum1",RooArgSet(signum,signum1r));
    RooProduct bbnum1("bbnum1","bbnum1",RooArgSet(bbnum,bbnum1r));
    RooProduct ulnunum1("ulnunum1","ulnunum1",RooArgSet(ulnunum,ulnunum1r));
    RooProduct connum1("connum1","connum1",RooArgSet(connum,connum1r));
    RooProduct signum2("signum2","signum2",RooArgSet(signum,signum2r));
    RooProduct bbnum2("bbnum2","bbnum2",RooArgSet(bbnum,bbnum2r));
    RooProduct ulnunum2("ulnunum2","ulnunum2",RooArgSet(ulnunum,ulnunum2r));
    RooProduct connum2("connum2","connum2",RooArgSet(connum,connum2r));
    RooAddPdf Model1("Model1","Model1",RooArgList(*sig1,*bb1,*ulnu1,*con1),RooArgList(signum1,bbnum1,ulnunum1,connum1));
    RooAddPdf Model2("Model2","Model2",RooArgList(*sig2,*bb2,*ulnu2,*con2),RooArgList(signum2,bbnum2,ulnunum2,connum2));
    RooSimultaneous Model("Model","Model",cat);
    Model.addPdf(Model1,"cat1");
    Model.addPdf(Model2,"cat2");

    //Generate toy MC
    RooDataSet *toyData1 = Model1.generate(RooArgSet(mu_pBrest),rnd.Poisson((sigScale1_35+bbScale1_35+conScale1_35+ulnuScale1_35)*1000));
    RooDataSet *toyData2 = Model2.generate(RooArgSet(mu_pBrest),rnd.Poisson((sigScale2_35+bbScale2_35+conScale2_35+ulnuScale2_35)*1000));
    RooDataSet toyData("toyData","joint toyData",mu_pBrest,Index(cat),Import("cat1",*toyData1),Import("cat2",*toyData2));

    std::cout<<"END"<<std::endl;
    //Fit to toy MC
    Model.fitTo(toyData,Extended(kTRUE),Minos(kTRUE));

    res[0] = signum.getVal();
    res[1] = signum.getError();
    res[2] = (res[0]-sigScale_35*1000)/res[1];
    delete toyData1;
    delete toyData2;
}
void toy_35()
{
    float res[3];
    TNtuple* n_toy = new TNtuple ( "n_toy", "toy MC studies", "yield:error:pull");
    float res1[3];

    //Suppress all the output from RooFit
    RooMsgService::instance().setGlobalKillBelow(RooFit::FATAL); 
    RooMsgService::instance().setSilentMode(true);

    //Generate 100 experiments
    for(int i=0;i<100;i++) {
        do_fit(res);
        n_toy->Fill(res);
	if(!(res[1]>0))
	break;
	std::cout<<"This is the "<<i<<"th fit \n";
    }
    RooMsgService::instance().setSilentMode(false);
    
    //Fit the pull with a gaussian
    RooRealVar pull("pull","pull",-5.,5.);
    
    RooRealVar mean("mean","mean", 0.,-1.,1.);
    RooRealVar width("width","width",1.,0.2,1.2);
    RooGaussian gaussian("gaussian","gaussian",pull,mean,width);
    
    RooDataSet *toy = new RooDataSet("toy","toy",n_toy,RooArgSet(pull));
    gaussian.fitTo(*toy,Minos(kTRUE));
    
    //Plot the yield, error and the fit result of the pull
    TCanvas* c1 = new TCanvas("c1","c1",1200,400);
    c1->Divide(3,1);
    n_toy->SetFillColor(50);
    
    c1->cd(1);
    n_toy->Draw("yield");
    
    c1->cd(2);
    n_toy->Draw("error");
    
    c1->cd(3);
    RooPlot* frame = pull.frame();
    toy->plotOn(frame,Binning(50));
    gaussian.plotOn(frame,LineWidth(3));
    frame->Draw();
    c1->SaveAs("toy_result_35000.eps");
}