#include <TROOT.h>
#include <TTree.h>
#include <TFile.h>
#include <TH1D.h>
#include <TTimeStamp.h>
#include <TRandom.h>
#include <TMVA/Reader.h>

#include <RooStats/ModelConfig.h>
#include <RooStats/HistFactory/Measurement.h>
#include "RooStats/HistFactory/HistRef.h"
#include "RooStats/HistFactory/Systematics.h"
#include <RooStats/HistFactory/MakeModelAndMeasurementsFast.h>
#include <RooStats/HistFactory/HistoToWorkspaceFactoryFast.h>

#include <iostream>
#include <fstream>
#include <sstream>
#include <string>

#include "StandardHypoTestFCNC.C"
#include "../utils/Colors.h"


void make_dm_roostats(string mvaname1="tree.root",string mvaname2="tree.root",
                string outfolder="../../../../BDT_Samples_3TeV/emu/Results_AnalysisLimits/",  // "../../Samples/Results/Less_overtraining/" ; "../../Samples/Results/Test_BDT/"
                string outname= "higgs_lfv_stat_model",
                string Analysis_varname="BDT_Response",
                string Cut_varname="BDT_Response",
                string varcut = ">-0.5",
                string weight = "reweight",
                double xmin = 0.3, double xmax = 0.65, int nbin = 250,    // BDTG
                double error_bkg = 0.2 // Systematics
                ) {

    string cut_br[2] = {"1.0","1.0"};  // Signal, Background   // 6.1e-5

    // Reading our trees and checking
    cout << Colors::Info << "Reading Files" << endl;
    TTree *samptree[2];
    TFile *f[2];
    string filenames[2] = {mvaname1, mvaname2};
    try {
        int nentries=0;
        for(int i=0;i<2;i++){
            f[i] = new TFile(filenames[i].c_str());
            f[i]->GetObject("HiggsLFV_all",samptree[i]);  
            if(!samptree[i]) cout << "NO TREEEE" << endl;
            nentries = samptree[i]->GetEntries();
            
            cout << "File " << filenames[i] << " was succesfully opened." << endl;
            cout << " MVA tree " << samptree[i]->GetTitle() << " with " << nentries << " entries opened\n\n";
        }
    } catch(...){
        cout << Colors::Error <<"Could not access tree, please verify filename" << endl;
        return;
    }
    
    // Generating Histogram and getting them from tree
    cout << Colors::Info << "Generating Histograms" << endl;
    TH1D * samphist[2];
    string hname[2]={"bdt_sig","bdt_bkg"};

    for(int j=0;j<2;j++){
        samphist[j]= new TH1D(hname[j].c_str(),"BDT response",nbin,xmin,xmax);
    }
    for(int i=0; i<2; i++){
        samptree[i]->Project(hname[i].c_str(),Analysis_varname.c_str(), "reweight"/*cutstr[i].c_str()*/);
    }
    
    // Signal
    if( gDirectory->FindObject("rsignal")!=NULL) gDirectory->Delete("rsignal");
    TH1D *rsignal = new TH1D("rsignal","BDT response",nbin,xmin,xmax);
    rsignal->Add(samphist[0]);

    cout << "Signal Histogram Created" << endl;
    
    
    // Background
    if( gDirectory->FindObject("rbg")!=NULL) gDirectory->Delete("rbg");
    TH1D *rbg = new TH1D("rbg","BDT response",nbin,xmin,xmax);
    rbg->Add(samphist[1]);

    cout << "Background Histogram Created" << endl;

    
    // Experiment Sample
    if( gDirectory->FindObject("rdata")!=NULL) gDirectory->Delete("rdata");
    TH1D *rdata = new TH1D("rdata","BDT response",nbin,xmin,xmax);
    
    TTimeStamp ttsmp;
    TRandom *RG = new TRandom(ttsmp.GetNanoSec()/1000);
    for(int ibin=0;ibin<nbin;ibin++){
        double nexp = rbg->GetBinContent(ibin);
        double nerr = rbg->GetBinError(ibin);
        double addvar = nexp - nerr*nerr; //Impose full statistical fluctuations on data (assuming no MC fluctuations)
        int  nobs = nexp + 0.5;

        if(addvar > 0.){
            addvar=sqrt(addvar);
            nobs = RG->Gaus(nexp+0.5, addvar); // ~Poission distribution
        }

        rdata->SetBinContent(ibin,nobs);
    }

    cout << "Experiment Sample Histogram Created" << endl;
    std::cout <<"Background events from MC= "<< rbg->Integral() << std::endl;
    std::cout << "Signal events from MC = "<< rsignal->Integral() << std::endl;
    std::cout << "Simulated data from background events = "<< rdata->Integral() << std::endl;

    //exit(0);
    //rbg->Scale(1.0/rbg->Integral());
    //rsignal->Scale(1.0/rsignal->Integral());
    //rdata->Scale(1.0/rdata->Integral());
  

    std::cout <<"Background events from MC= "<< rbg->Integral() << std::endl;
    std::cout << "Signal events from MC = "<< rsignal->Integral() << std::endl;
    std::cout << "Simulated data from background events = "<< rdata->Integral() << std::endl;

    //exit(0);
    // ====================
    // RooStats starts here
    // ====================
    // Model building based on Exercise 8 from
    //     https://twiki.cern.ch/twiki/bin/view/RooStats/RooStatsExercisesMarch2015

    cout << endl << Colors::Info << "Initializing Roostats HistFactory" << endl;
    RooStats::HistFactory::Measurement BRlimit("BRlimit","Limit on LFV Higgs branching ratio");

    BRlimit.SetOutputFilePrefix(outfolder+"RooStats");
    BRlimit.SetPOI("BR"); // Parameter of interest
    
    BRlimit.SetLumi(1.0);
    BRlimit.SetLumiRelErr(0.001);  // not relevant (?!) but can not be set to zero  (NAN in calculations!)// this does not make lumi varying
    BRlimit.AddConstantParam("Lumi");

    RooStats::HistFactory::Channel BDTresp("BDTresp"); 
    BDTresp.SetData(rdata); // rdata

    RooStats::HistFactory::Sample signal("signal");
    // signal.AddNormFactor("BR",0.0001, 0.0,0.0008);
    signal.AddNormFactor("BR",0.3,0.,5.);
    signal.SetHisto(rsignal);

    BDTresp.AddSample(signal);

    RooStats::HistFactory::Sample backg("background");
    backg.SetHisto(rbg);
    backg.AddOverallSys("bkg_unc", 1.0 - error_bkg, 1.0 + error_bkg); // Background systematics

    BDTresp.AddSample(backg);

    BRlimit.AddChannel(BDTresp);

    // make the model 
    cout << endl << Colors::Info << "Making the Workspace" << endl;

    // ======================================== //
    
    RooStats::HistFactory::HistoToWorkspaceFactoryFast factory{BRlimit};

    HistFactory::Channel& channel = BRlimit.GetChannels().at( 0 );
    std::unique_ptr<RooWorkspace> ws_single{factory.MakeSingleChannelModel( BRlimit, channel )};

    auto combined_config = static_cast<ModelConfig *>(ws_single->obj("ModelConfig"));

    if (ws_single->data("obsData")) {
        RooAbsData* simData = ws_single->data("obsData");
        RooAbsPdf* model = combined_config->GetPdf();
        model->fitTo(*simData, Minos(true), PrintLevel(0));
    }
    else {
        RooAbsData* simData = ws_single->data("asimovData");
        RooAbsPdf* model = combined_config->GetPdf();
        model->fitTo(*simData, Minos(true), PrintLevel(0));
    }

    // ======================================== //

}

void callforLimits_basic(){

    make_dm_roostats("/Users/fmgaray/Documents/Research/clic_lfv_tau/treeLFV_sgn_ntuple.root",
                     "/Users/fmgaray/Documents/Research/clic_lfv_tau/treeLFV_bkg_ntuple.root");
    
}