#include <TCanvas.h>
#include <TH1.h>
#include <stdio.h>

#include "RooStats/HistFactory/Measurement.h"
#include "RooStats/HistFactory/MakeModelAndMeasurementsFast.h"


#define USE_STANDARDHYPOTESTDEMO
#define USE_MODELINSPECTOR
#undef  SHOW_HISTOGRAMS

#define LUMINOSITY_DEFAULT 35000.0


// Data sets (note: sum of backgrounds must be nonzero)

#if 1

const int BinCount = 10;

const int SignalBinContent[BinCount + 2] = { 0, 3,  27,  53,  82,  58,  33,  28,  26,  15,   5, 0 };
const int Bkg1BinContent[BinCount + 2]   = { 0, 1,  14,  61, 106,  68,  51,  18,   8,   4,   0, 0 };
const int Bkg2BinContent[BinCount + 2]   = { 0, 5,  24,  29,  23,   9,   6,   1,   0,   1,   1, 0 };

const float SignalCrossSection = 3.1154;
const float Bkg1CrossSection   = 0.87961;
const float Bkg2CrossSection   = 0.26008;

const float SignalEfficiency   = 0.065436;
const float Bkg1Efficiency     = 0.055157;
const float Bkg2Efficiency     = 0.020988;

#elif 1

const int BinCount = 20;

const int SignalBinContent[BinCount + 2] = { 0,  1,  2,  3,  4,  5,  6,  7,  8,  9, 10, 11, 12, 13, 14, 15, 16, 17, 18, 19, 20, 0 };
const int Bkg1BinContent[BinCount + 2]   = { 0, 20, 19, 18, 17, 16, 15, 14, 13, 12, 11, 10,  9,  8,  7,  6,  5,  4,  3,  2,  1, 0 };
const int Bkg2BinContent[BinCount + 2]   = { 0, 20, 19, 18, 17, 16, 15, 14, 13, 12, 11, 10,  9,  8,  7,  6,  5,  4,  3,  2,  1, 0 };

const float SignalCrossSection = 1.0;
const float Bkg1CrossSection   = 0.5;
const float Bkg2CrossSection   = 0.5;

const float SignalEfficiency   = 1.0;
const float Bkg1Efficiency     = 1.0;
const float Bkg2Efficiency     = 1.0;

#elif 1

const int BinCount = 20;

const int SignalBinContent[BinCount + 2] = { 0,  0,  0,  0,  1,  2,  4,  8, 16, 32, 16,  8,  4,  2,  1,  0,  0,  0,  0,  0,  0, 0 };
const int Bkg1BinContent[BinCount + 2]   = { 0,  1,  1,  2,  2,  3,  3,  4,  4,  5,  5,  6,  6,  7,  7,  8,  8,  9,  9, 10, 10, 0 };
const int Bkg2BinContent[BinCount + 2]   = { 0,  1,  1,  2,  2,  3,  3,  4,  4,  5,  5,  6,  6,  7,  7,  8,  8,  9,  9, 10, 10, 0 };

const float SignalCrossSection = 1.0;
const float Bkg1CrossSection   = 0.5;
const float Bkg2CrossSection   = 0.5;

const float SignalEfficiency   = 1.0;
const float Bkg1Efficiency     = 1.0;
const float Bkg2Efficiency     = 1.0;

#elif 1

const int BinCount = 20;

const int SignalBinContent[BinCount + 2] = { 0,  1,  2,  5, 10, 20, 40, 80, 160, 320, 160, 80, 40, 20, 10,  5,  2,  1,  0,   0,   0, 0 };
const int Bkg1BinContent[BinCount + 2]   = { 0, 10, 15, 20, 25, 30, 35, 40,  45,  50,  55, 60, 65, 70, 75, 80, 85, 90, 95, 100, 105, 0 };
const int Bkg2BinContent[BinCount + 2]   = { 0, 10, 15, 20, 25, 30, 35, 40,  45,  50,  55, 60, 65, 70, 75, 80, 85, 90, 95, 100, 105, 0 };

const float SignalCrossSection = 1.0;
const float Bkg1CrossSection   = 0.5;
const float Bkg2CrossSection   = 0.5;

const float SignalEfficiency   = 1.0;
const float Bkg1Efficiency     = 1.0;
const float Bkg2Efficiency     = 1.0;

#elif 1

const int BinCount = 6;

const int SignalBinContent[BinCount + 2] = { 0,  1000, 10000, 100000, 10000,  1000,   100,   0 };
const int Bkg1BinContent[BinCount + 2]   = { 0, 10000, 20000,  30000, 40000, 50000, 60000,   0 };
const int Bkg2BinContent[BinCount + 2]   = { 0, 10000, 20000,  30000, 40000, 50000, 60000,   0 };

const float SignalCrossSection = 1.0;
const float Bkg1CrossSection   = 0.5;
const float Bkg2CrossSection   = 0.5;

const float SignalEfficiency   = 1.0;
const float Bkg1Efficiency     = 1.0;
const float Bkg2Efficiency     = 1.0;

#else
#error No data set selected, can't continue!
#endif


#define FILE_HISTOGRAMS "/tmp/repro-significance-lumi.root"

#define WORKSPACE_NAME "combined"
#define MEASUREMENT_NAME "MyMeasurement"

#define FILE_HISTFACTORYPREFIX "/tmp/repro-significance-lumi"

#define FILE_ROOWORKSPACE "/tmp/repro-significance-lumi-RooWorkspace.root"
#define FILE_HISTFACTORYMODEL FILE_HISTFACTORYPREFIX "_" WORKSPACE_NAME "_" MEASUREMENT_NAME "_model.root"


#ifdef USE_STANDARDHYPOTESTDEMO
// forward declaration of main function of StandardHypoTestDemo.C
void StandardHypoTestDemo(const char* infile = "",
                          const char* workspaceName = "combined",
                          const char* modelSBName = "ModelConfig",
                          const char* modelBName = "",
                          const char* dataName = "obsData",
                          int calcType = 0, /* 0 freq 1 hybrid, 2 asymptotic */
                          int testStatType = 3,   /* 0 LEP, 1 TeV, 2 LHC, 3 LHC - one sided*/
                          int ntoys = 5000,
                          bool useNC = false,
                          const char * nuisPriorName = 0);
#endif

#ifdef USE_MODELINSPECTOR
// forward declaration of main function of ModelInspector.C
void ModelInspector(
    const char * infile = "",
    const char * workspaceName = "combined",
    const char * modelConfigName = "ModelConfig",
    const char * dataName = "obsData");
#endif


TH1D * CreateInitHistogram(const char * szName, const char * szTitle, const int * rgnContent)
{
#if 1
    TH1D * phist = new TH1D(szName, szTitle, BinCount, 0.0, 1000.0);
    phist->Sumw2();

    if (rgnContent)
    {
        for (size_t i = 0; i < BinCount + 2; i++)
        {
            int ibin = phist->GetBin(i);
            double x = phist->GetBinCenter(ibin);

            // fill bin as many times as indicated
            // (this should be an uncontroversial way to regenerate histogram from counts, as it produces correct errors and # of entries)
            for (int n = rgnContent[i]; n > 0; n--)
                phist->Fill(x);
        }
    }
#else
    TH1D * phist = new TH1D(szName, szTitle, 1, 0.0, 1000.0);
    phist->Sumw2();

    if (rgnContent)
    {
        for (size_t i = 0; i < cContentCount; i++)
        {
            for (int n = rgnContent[i]; n > 0; n--)
                phist->Fill(555.555);
        }
    }
#endif

    return phist;
}


void ScaleHistogram(TH1D * pHistogram, float fLuminosity, float fCrossSection, float fEfficiency)
{
    pHistogram->Scale(fLuminosity * fCrossSection * fEfficiency / pHistogram->Integral());
}


void Repro(float fLuminosity = LUMINOSITY_DEFAULT)
{
    // initialize histograms from raw events (condensed to counts for simplicity)

    TH1D * phistsig  = CreateInitHistogram("h_sig", "Signal",        SignalBinContent);
    TH1D * phistbkg1 = CreateInitHistogram(0,       "Background #1", Bkg1BinContent);
    TH1D * phistbkg2 = CreateInitHistogram(0,       "Background #2", Bkg2BinContent);

#ifdef SHOW_HISTOGRAMS
    {
        TCanvas * pcanvas = new TCanvas("RawHistograms", "Raw Histograms", 800, 600);
        pcanvas->Divide(2, 2);
        pcanvas->cd(1); phistsig->Draw();
        pcanvas->cd(2); phistbkg1->Draw();
        pcanvas->cd(4); phistbkg2->Draw();
        pcanvas->Update();
        pcanvas->WaitPrimitive();
        printf("Press Enter>"); getchar();
        pcanvas->Close();
        delete pcanvas;
    }
#endif

    // scale histograms

    ScaleHistogram(phistsig,  fLuminosity, SignalCrossSection, SignalEfficiency);
    ScaleHistogram(phistbkg1, fLuminosity, Bkg1CrossSection,   Bkg1Efficiency);
    ScaleHistogram(phistbkg2, fLuminosity, Bkg2CrossSection,   Bkg2Efficiency);

    // create composite background and data (signal+background) histograms

    TH1D * phistbkg = CreateInitHistogram("h_bkg", "Background", 0);
    phistbkg->Add(phistbkg1, phistbkg2);

    TH1D * phistdat = CreateInitHistogram("h_dat", "Data", 0);
    phistdat->Add(phistsig, phistbkg);

#ifdef SHOW_HISTOGRAMS
    {
        TCanvas * pcanvas = new TCanvas("ScaledHistograms", "Scaled Histograms", 800, 600);
        pcanvas->Divide(2, 2);
        pcanvas->cd(1); phistsig->Draw();
        pcanvas->cd(2); phistbkg->Draw();
        pcanvas->cd(3); phistdat->Draw();
        pcanvas->Update();
        pcanvas->WaitPrimitive();
        printf("Press Enter>"); getchar();
        pcanvas->Close();
        delete pcanvas;
    }
#endif

    // write signal, background, and data histograms to file

    {
        TFile f(FILE_HISTOGRAMS, "recreate");
        phistsig->Write("signal");
        phistbkg->Write("background");
        phistdat->Write("data");
        f.Close();
    }

    // create HistFactory prerequisites

    unlink(FILE_HISTFACTORYMODEL);  // make sure we don't accidentally use an old one if it fails to generate a new one

        // example.xml layout:
        //   Combination
        //     Input = "example_channel.xml"
        //     Measurement: Name, Lumi, LumiRelErr
        //       POI = "SigXsecOverSM"
        //       ParamSetting = "Lumi alpha_syst1"

        // example_channel.xml layout:
        //   Channel: Name, InputFile
        //     Data: HistoName
        //     StatErrorConfig: RelErrorThreshold, ConstraintType
        //     Sample: Name, HistoName
        //       OverallSys: Name, High, Low
        //       NormFactor: Name, Val, Low, High
        //     Sample: Name, NormalizeByTheory, HistoName
        //       StatError: Activate
        //       OverallSys: Name, Low, High

    RooStats::HistFactory::Measurement meas(MEASUREMENT_NAME, "My Measurement");

    meas.SetOutputFilePrefix(FILE_HISTFACTORYPREFIX);
    meas.SetExportOnly(false);

    meas.SetLumi(fLuminosity);
    meas.SetLumiRelErr(0.1);

    //meas.SetBinLow(0);
    //meas.SetBinHigh(999);

    meas.SetPOI("SigXsecOverSM");

    meas.AddConstantParam("Lumi");
    meas.AddConstantParam("alpha_syst1");

    RooStats::HistFactory::Channel chan("channel");

    chan.SetData("data", FILE_HISTOGRAMS);

    chan.SetStatErrorConfig(0.05, "Poisson");

    RooStats::HistFactory::Sample samplesig("signal", "signal", FILE_HISTOGRAMS);
    samplesig.AddOverallSys("syst1", 0.95, 1.05);
    samplesig.AddNormFactor("SigXsecOverSM", 10, 0.0, 50.0);
    chan.AddSample(samplesig);

    RooStats::HistFactory::Sample samplebkg("background", "background", FILE_HISTOGRAMS);
    samplebkg.ActivateStatError();
    samplebkg.AddOverallSys("syst2", 0.95, 1.05);
    chan.AddSample(samplebkg);

    meas.AddChannel(chan);

    meas.CollectHistograms();

    // create workspace and model

    RooWorkspace * prws = MakeModelAndMeasurementFast(meas);

    //prws->Print();
    //prws->writeToFile(FILE_ROOWORKSPACE);

#ifdef USE_STANDARDHYPOTESTDEMO
    StandardHypoTestDemo(
        FILE_HISTFACTORYMODEL,
        WORKSPACE_NAME,
        "ModelConfig",
        "",
        "obsData",
        2,
        3,
        5000,
        false,
        0);

    printf("\n\nThe above output is from: StandardHypoTestDemo(\"%s\", ...)\n", FILE_HISTFACTORYMODEL);
#endif

#ifdef USE_MODELINSPECTOR
    printf("\n\nAbout to run: ModelInspector(\"%s\")\nPress Enter>", FILE_HISTFACTORYMODEL);
    getchar();

    ModelInspector(FILE_HISTFACTORYMODEL);
#endif
}


#ifdef USE_STANDARDHYPOTESTDEMO
#include "/...YOUR_PATH_HERE.../root-6.10.00/tutorials/roostats/StandardHypoTestDemo.C"
#endif

#ifdef USE_MODELINSPECTOR
#include "/...YOUR_PATH_HERE.../root-6.10.00/tutorials/roostats/ModelInspector.C"
#endif