#include <iostream>
#include <stdio.h>
#include <string>
#include <algorithm>
#include <vector>
#include <numeric>
#include <time.h>
#include <random>
#include <sstream>

#include "TRandom3.h"
#include "TCanvas.h"
#include "TDatime.h"
#include "TStopwatch.h"
#include "TLegend.h"
#include "TIterator.h"
#include "TH3.h"
#include "TLatex.h"

#include "RooChi2Var.h"
#include "RooAbsData.h"
#include "RooRealSumPdf.h"
#include "RooPoisson.h"
#include "RooGaussian.h"
#include "RooRealVar.h"
#include "RooMCStudy.h"
#include "RooMinuit.h"
#include "RooCategory.h"
#include "RooHistPdf.h"
#include "RooSimultaneous.h"
#include "RooExtendPdf.h"
#include "RooDataSet.h"
#include "RooDataHist.h"
#include "RooFitResult.h"
#include "RooMsgService.h"
#include "RooParamHistFunc.h"
#include "RooHist.h"
#include "RooRandom.h"
#include "RooMsgService.h"

#include "RooStats/ModelConfig.h"
#include "RooStats/ToyMCSampler.h"
#include "RooStats/MinNLLTestStat.h"

#include "RooStats/HistFactory/FlexibleInterpVar.h"
#include "RooStats/HistFactory/PiecewiseInterpolation.h"
#include "RooStats/HistFactory/HistFactorySimultaneous.h"
#include "RooStats/HistFactory/Channel.h"
#include "RooStats/HistFactory/MakeModelAndMeasurementsFast.h"
#include "RooStats/HistFactory/Measurement.h"
#include "RooStats/HistFactory/ParamHistFunc.h"
#include "RooStats/HistFactory/HistFactoryModelUtils.h"
#include "RooStats/HistFactory/RooBarlowBeestonLL.h"

#include "RooStats/FrequentistCalculator.h"
#include "RooStats/HypoTestInverter.h"
#include "RooStats/HypoTestInverterPlot.h"

using namespace RooFit;
using namespace RooStats;
using namespace HistFactory;

# define NToys 10000

std::random_device rd;
std::default_random_engine generator(rd());

std::vector<std::string> Sample_names = {
    "L_x_signal_nominal_channel_overallSyst_x_Exp",
    "L_x_background_nominal_channel_overallSyst_x_HistSyst"
};

void MakeTemplate() {

    // make signal template
    TH1D* signal = new TH1D("signal_nominal", "variable;evt", 4, 0, 4);
    signal->SetBinContent(1, 1);
    signal->SetBinContent(2, 1);
    signal->SetBinContent(3, 2);
    signal->SetBinContent(4, 5);

    // make background template
    TH1D* background = new TH1D("background_nominal", "variable;evt", 4, 0, 4);
    background->SetBinContent(1, 50);
    background->SetBinContent(2, 50);
    background->SetBinContent(3, 40);
    background->SetBinContent(4, 30);

    // make background systematic template
    TH1D* background_p = new TH1D("background_p", "variable;evt", 4, 0, 4);
    background_p->SetBinContent(1, 53);
    background_p->SetBinContent(2, 53);
    background_p->SetBinContent(3, 43);
    background_p->SetBinContent(4, 26);

    TH1D* background_m = new TH1D("background_m", "variable;evt", 4, 0, 4);
    background_m->SetBinContent(1, 47);
    background_m->SetBinContent(2, 47);
    background_m->SetBinContent(3, 37);
    background_m->SetBinContent(4, 33);

    // Save histogram templates
    TFile* file = new TFile("PDFandDATA.root", "RECREATE");
    signal->Write();
    background->Write();
    background_p->Write();
    background_m->Write();
    file->Close();

}

void MakeWorkSpace() {
    const double expmu = 1.0;
    const char* fname = "PDFandDATA.root";

    RooStats::HistFactory::Measurement meas("my_measurement", "my measurement");
    meas.SetOutputFilePrefix("results/my_measurement");
    meas.SetExportOnly(kTRUE);

    // setting measurement
    meas.SetPOI("mu");
    meas.SetLumi(1.0);
    meas.AddConstantParam("Lumi");

    // define channel
    HistFactory::Channel channel("channel");
    channel.SetStatErrorConfig(1e-5, "Poisson");

    // read signal template
    RooStats::HistFactory::Sample sig_temp("signal_nominal", "signal_nominal", fname);
    sig_temp.AddOverallSys("efficiency_uncertainty", 0.92, 1.08);
    sig_temp.AddNormFactor("mu", expmu, -100.0, 100.0);
    channel.AddSample(sig_temp);

    // read background template
    RooStats::HistFactory::Sample bkg_temp("background_nominal", "background_nominal", fname);
    bkg_temp.AddHistoSys("some_uncertainty", "background_m", fname, "", "background_p", fname, "");
    bkg_temp.AddOverallSys("mu_bkg", 0.95, 1.05);
    bkg_temp.SetNormalizeByTheory(kFALSE);
    channel.AddSample(bkg_temp);

    // add channel to measurement
    meas.AddChannel(channel);
    meas.CollectHistograms();

    RooWorkspace* w;
    w = RooStats::HistFactory::MakeModelAndMeasurementFast(meas);

    w->Print();
    w->writeToFile("PDFandDATA_workspace.root");
}

double FluctuateNuisanceParameter(RooWorkspace* w) {
    double Nevt = 0.0;

    w->loadSnapshot("NominalParamValues");
    ModelConfig* mc = (ModelConfig*)w->obj("ModelConfig"); // Get model manually
    RooSimultaneous* model = (RooSimultaneous*)mc->GetPdf();
    RooArgSet* obs = (RooArgSet*)mc->GetObservables();
    RooRealVar* x_val = w->var("obs_x_channel");

    std::normal_distribution<double> distribution(0.0, 1.0);

    // fluctuate nuisance parameters
    w->var("alpha_efficiency_uncertainty")->setVal(distribution(generator));
    w->var("alpha_mu_bkg")->setVal(distribution(generator));
    w->var("alpha_some_uncertainty")->setVal(distribution(generator));

    // get the number of event for the current template
    RooAbsBinning const& binning = x_val->getBinning();
    const double oldVal = x_val->getVal();

    for (std::size_t iBin = 0; iBin < binning.numBins(); ++iBin) {
        double binCenter = binning.binCenter(iBin);
        double binWidth = binning.binWidth(iBin);

        *x_val = binCenter; // set x value

        for (unsigned int j = 0; j < Sample_names.size(); j++) {
            RooAbsReal* temp_func = w->function(Sample_names.at(j).c_str());
            Nevt = Nevt + temp_func->getValV();
            if (temp_func->getValV() < 0) {
                printf("[ERROR] negative count!\n");
                exit(1);
            }
        }

    }

    *x_val = oldVal;

    return Nevt;
}

void MyToyMCStudy() {

    std::vector<double> mus;
    std::vector<double> errors;

    const char* fname = "./PDFandDATA_workspace.root";
    TFile* f = TFile::Open(fname);

    RooWorkspace* w = (RooWorkspace*)f->Get("combined");

    ModelConfig* mc = (ModelConfig*)w->obj("ModelConfig"); // Get model manually
    RooSimultaneous* model = (RooSimultaneous*)mc->GetPdf();

    RooArgSet* obs = (RooArgSet*)mc->GetObservables();
    RooRealVar* x = (RooRealVar*)obs->find("obs_x_channel");

    for (int i = 0; i < NToys; i++) {
        double Nevt_total = FluctuateNuisanceParameter(w);

        // generate toys
        RooDataSet* genData = model->generate(RooArgSet(*x, model->indexCat()), Nevt_total, false, true, "", false, true);

        // get nominal values for nuisance parameters
        w->loadSnapshot("NominalParamValues");

        // fit!
        RooFitResult* fitres = model->fitTo(*genData, RooFit::Extended(true), RooFit::SumW2Error(false), PrintLevel(-1), Save());

        // save fitting result
        RooArgSet fitargs = fitres->floatParsFinal();
        TIterator* iter(fitargs.createIterator());

        for (TObject* a = iter->Next(); a != 0; a = iter->Next()) {
            RooRealVar* rrv = dynamic_cast<RooRealVar*>(a);
            std::string name = rrv->GetName();
            double val = rrv->getVal();
            double err = rrv->getError();

            if (name == "mu") {
                mus.push_back(val);
                errors.push_back(err);
            }
        }

    }

    // toy MC study is done. lets plot pull distribution
    TH1F* ToyMCmu = new TH1F("ToyMCmu", ";#mu;Toys", 40, -10, 15);
    TH1F* ToyMCmuerror = new TH1F("ToyMCmuerror", ";error of #mu;Toys", 50, 1, 7);
    TH1F* ToyMCmupull = new TH1F("ToyMCmupull", ";pull of #mu;Toys", 40, -4, 4);

    ToyMCmu->SetMarkerStyle(kFullCircle);
    ToyMCmu->SetLineColor(kBlack);
    ToyMCmu->SetMarkerColor(kBlack);

    ToyMCmuerror->SetMarkerStyle(kFullCircle);
    ToyMCmuerror->SetLineColor(kBlack);
    ToyMCmuerror->SetMarkerColor(kBlack);

    ToyMCmupull->SetMarkerStyle(kFullCircle);
    ToyMCmupull->SetLineColor(kBlack);
    ToyMCmupull->SetMarkerColor(kBlack);

    for (unsigned int j = 0; j < NToys; j++) { // Fill
        ToyMCmu->Fill(mus.at(j));
        ToyMCmuerror->Fill(errors.at(j));
        ToyMCmupull->Fill((mus.at(j) - 1.0) / errors.at(j));
    }

    gStyle->SetOptFit(11);

    TCanvas* c = new TCanvas("canvas_ToyMC_study", "", 800, 800);
    ToyMCmu->Draw("PE1");
    c->SaveAs("TOYMC_mu.png");
    delete c;

    c = new TCanvas("canvas_ToyMC_study", "", 800, 800);
    ToyMCmuerror->Draw("PE1");
    c->SaveAs("TOYMC_muerror.png");
    delete c;

    c = new TCanvas("canvas_ToyMC_study", "", 800, 800);
    ToyMCmupull->Fit("gaus");
    ToyMCmupull->Draw("PE1");
    c->SaveAs("TOYMC_mupull.png");
    delete c;

    delete ToyMCmu;
    delete ToyMCmuerror;
    delete ToyMCmupull;
}

int Histfactory_bias() {

    RooMsgService::instance().setStreamStatus(1, false);
    RooMsgService::instance().setGlobalKillBelow(RooFit::ERROR);

    MakeTemplate();
    MakeWorkSpace();
    MyToyMCStudy();

    return 0;
}