#include <iostream>
#include <string>
#include <sstream>
#include <vector>

#include <TLatex.h>
#include <TCanvas.h>
#include <TApplication.h>
#include <TFile.h>
#include <TH1.h>
#include <TStyle.h>
#include <TPad.h>

#include <RooMsgService.h>
#include <RooRandom.h>
#include <RooRealVar.h>
#include <RooGenericPdf.h>
#include <RooBernstein.h>
#include <RooArgList.h>
#include <RooPlot.h>
#include <RooDataHist.h>
#include <RooHistPdf.h>
#include <RooAddPdf.h>
#include <RooChi2Var.h>
#include <RooNLLVar.h>
#include <RooMinimizer.h>
#include <RooGaussian.h>

#include "generateSimulation.h"

using namespace std;
using namespace RooFit;

TApplication Runner("gui",0,NULL);

RooRealVar *x;

int num_simulations=200;
double BKG=250000;
TH1D *H,*Z;

// Parameters

RooRealVar *a0_gen;
RooRealVar *a1_gen;
RooRealVar *a2_gen;
RooRealVar *a3_gen;
RooRealVar *a4_gen;

RooRealVar *a0_fit;
RooRealVar *a1_fit;
RooRealVar *a2_fit;
RooRealVar *a3_fit;
RooRealVar *a4_fit;

RooRealVar *n_bkg;
RooRealVar *n_H;
RooRealVar *n_Z;

void inizializeParameters();

int main(int narg,char* argv[]){
    RooMsgService::instance().setGlobalKillBelow(RooFit::FATAL);
    RooRandom::randomGenerator()->SetSeed(0);
    
    for(int i=1;i<narg;i++){
        string tmp(argv[i]);
        
        if (tmp.find("--nsim=")==string::npos){cout<<"No valid commands specified"<<endl<<endl;
            exit(0);}
    }
    
    for(int i=1;i<narg;i++){
        string tmp(argv[i]);
        
        if (tmp.find("--nsim=")!=string::npos){
            istringstream is(tmp.substr(7,tmp.length()-7));
            is>>num_simulations;
        }
        
    }

    x=new RooRealVar("x","m_{bb}    [GeV]",60,300);
    
    
    TFile *f_in=new TFile("histograms.root");
    H=(TH1D*)f_in->Get("Hmass");
    Z=(TH1D*)f_in->Get("Zmass");
    
    RooDataHist H_H("H_H","H_H",*x,H);
    RooDataHist H_Z("H_Z","H_Z",*x,Z);
    
    RooHistPdf pdf_H("pdf_H","pdf_H",*x,H_H,2);
    RooHistPdf pdf_Z("pdf_Z","pdf_Z",*x,H_Z,2);
    
    
    
    
    a0_gen=new RooRealVar("a0_gen","a0_gen",27.7044,0.,500.);
    a1_gen=new RooRealVar("a1_gen","a1_gen",10.9279,0.,500.);
    a2_gen=new RooRealVar("a2_gen","a2_gen",7.2648,0.,500.);
    a3_gen=new RooRealVar("a3_gen","a3_gen",2.09955,0.,500.);
    a4_gen=new RooRealVar("a4_gen","a4_gen",1.42476,0.,500.);
    
    RooBernstein Bern_gen("Bern_gen","Bern_gen",*x,RooArgList(*a0_gen,*a1_gen,*a2_gen,*a3_gen,*a4_gen));
    
    
    
    
    
    a0_fit=new RooRealVar("a0_fit","a0_fit",27.7044,0.,500.);
    a1_fit=new RooRealVar("a1_fit","a1_fit",10.9279,0.,500.);
    a2_fit=new RooRealVar("a2_fit","a2_fit",7.2648,0.,500.);
    a3_fit=new RooRealVar("a3_fit","a3_fit",2.09955,0.,500.);
    a4_fit=new RooRealVar("a4_fit","a4_fit",1.42476,0.,500.);
    
    RooBernstein Bern_fit("Bern_fit","Bern_fit",*x,RooArgList(*a0_fit,*a1_fit,*a2_fit,*a3_fit,*a4_fit));
    
    
    
    n_bkg=new RooRealVar("n_bkg","n_bkg",BKG,0.,3*BKG);
    n_H=new RooRealVar("n_H","n_H",H->Integral(),-100*H->Integral(),100*H->Integral());
    n_Z=new RooRealVar("n_Z","n_Z",Z->Integral(),-100*Z->Integral(),100*Z->Integral());
    
    
    RooAddPdf *pdf_fit=new RooAddPdf("pdf_fit","pdf_fit",RooArgList(Bern_fit,pdf_H,pdf_Z),RooArgList(*n_bkg,*n_H,*n_Z));
    
    
    
    vector<double> chi2_results_H,nll_results_H;
    vector<double> chi2_results_Z,nll_results_Z;
    
    cout<<endl;
    
    for(int i=0;i<num_simulations;i++){
        
        cout <<"\x1b[A";
        printf("%c[2K", 27);
        cout<<"\rSimulation number: "<<i+1<<"/"<<num_simulations<<endl;
        
        TH1D *simulation=(TH1D*)generateSimulation(x,&Bern_gen,&pdf_H,&pdf_Z,BKG,H->Integral(),Z->Integral());
        
        RooDataHist H_simulation("H_simulation","H_simulation",*x,simulation);
        
        
        
        inizializeParameters();
        
        RooChi2Var CHI2("CHI2","CHI2",*pdf_fit,H_simulation,Extended(true));
        RooMinimizer *minuit_chi2=new RooMinimizer(CHI2);
        minuit_chi2->setPrintLevel(-1);
        int status_migrad=minuit_chi2->migrad();
        int status_hesse=minuit_chi2->hesse();
        int status_minos=minuit_chi2->minos();
        
        if(status_migrad!=0 || status_hesse!=0 || status_minos!=0)cout<<"FIT FAILED"<<endl<<endl;
        
        chi2_results_H.push_back((n_H->getValV()/H->Integral()-1)/n_H->getError()*H->Integral());
        chi2_results_Z.push_back((n_Z->getValV()/Z->Integral()-1)/n_Z->getError()*Z->Integral());
        
        
        
        
        inizializeParameters();
        
        RooNLLVar NLL("NLL","NLL",*pdf_fit,H_simulation,Extended(true));
        RooMinimizer *minuit_nll=new RooMinimizer(NLL);
        minuit_nll->setPrintLevel(-1);
        status_migrad=minuit_nll->migrad();
        status_hesse=minuit_nll->hesse();
        status_minos=minuit_nll->minos();
        
        if(status_migrad!=0 || status_hesse!=0 || status_minos!=0)cout<<"FIT FAILED"<<endl<<endl;
        
        nll_results_H.push_back((n_H->getValV()/H->Integral()-1)/n_H->getError()*H->Integral());
        nll_results_Z.push_back((n_Z->getValV()/Z->Integral()-1)/n_Z->getError()*Z->Integral());
        
    }
    
    
    TH1D *plot_H_chi2=new TH1D("plot_H","plot_H",60,-6.,6.);
    TH1D *plot_Z_chi2=(TH1D*)plot_H_chi2->Clone();
    TH1D *plot_H_nll=(TH1D*)plot_H_chi2->Clone();
    TH1D *plot_Z_nll=(TH1D*)plot_H_chi2->Clone();
    
    for(int i=0;i<chi2_results_H.size();i++){
        plot_H_chi2->Fill(chi2_results_H[i]);
        plot_Z_chi2->Fill(chi2_results_Z[i]);
        
        plot_H_nll->Fill(nll_results_H[i]);
        plot_Z_nll->Fill(nll_results_Z[i]);
    }
    
    RooRealVar y("y","y",-6.,6.);
    
    RooDataHist H_chi2_H("H_chi2_H","H_chi2_H",y,plot_H_chi2);
    RooDataHist H_chi2_Z("H_chi2_Z","H_chi2_Z",y,plot_Z_chi2);
    
    RooDataHist H_nll_H("H_nll_H","H_nll_H",y,plot_H_nll);
    RooDataHist H_nll_Z("H_nll_Z","H_nll_Z",y,plot_Z_nll);
    
    
    
    RooRealVar mean("mean","mean",0,-6.,6.);
    RooRealVar sigma("sigma","sigma",1,0.,6.);
    RooGaussian gauss("gauss","gauss",y,mean,sigma);
    
    
    gStyle->SetOptTitle(0);
    
    TCanvas *c=new TCanvas;
    c->Divide(1,2);
    
    c->cd(2);
    RooPlot *frame_H=y.frame();
    H_chi2_H.plotOn(frame_H,LineColor(2),MarkerColor(2));
    gauss.fitTo(H_chi2_H,Hesse(true),Minos(true));
    gauss.plotOn(frame_H,LineColor(2),LineStyle(2));
    
    
    H_nll_H.plotOn(frame_H,LineColor(4),MarkerColor(4));
    gauss.fitTo(H_nll_H,Hesse(true),Minos(true));
    gauss.plotOn(frame_H,LineColor(4),LineStyle(2));
    frame_H->Draw();
    
    gPad->SetGridx();
    gPad->SetGridy();
    
    c->cd(1);
    RooPlot *frame_Z=y.frame();
    H_chi2_Z.plotOn(frame_Z,LineColor(2),MarkerColor(2));
    gauss.fitTo(H_chi2_Z,Hesse(true),Minos(true));
    gauss.plotOn(frame_Z,LineColor(2),LineStyle(2));
    
    
    H_nll_Z.plotOn(frame_Z,LineColor(4),MarkerColor(4));
    gauss.fitTo(H_nll_Z,Hesse(true),Minos(true));
    gauss.plotOn(frame_Z,LineColor(4),LineStyle(2));
    frame_Z->Draw();
    
    gPad->SetGridx();
    gPad->SetGridy();
    
    c->Draw();
    
    Runner.Run(true);
}


void inizializeParameters(){
    
    a0_fit->setVal(a0_gen->getValV());
    a1_fit->setVal(a1_gen->getValV());
    a2_fit->setVal(a2_gen->getValV());
    a3_fit->setVal(a3_gen->getValV());
    a4_fit->setVal(a4_gen->getValV());
    
    n_bkg->setVal(BKG);
    n_H->setVal(H->Integral());
    n_Z->setVal(Z->Integral());
    
}