#include <TROOT.h>
#include <TTree.h>
#include <TFile.h>
#include <TH1D.h>
#include <TStyle.h>
#include <TCanvas.h>
#include <TGraph.h>
#include <TLegend.h>
#include <TTimeStamp.h>
#include <TRandom.h>

#include <RooStats/ModelConfig.h>
#include <RooStats/HistFactory/Measurement.h>
#include <RooStats/HistFactory/MakeModelAndMeasurementsFast.h>

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


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


void make_dm_roostats(string mvaname1="tree.root",string mvaname2="tree.root",
                string outfolder="/Users/fmgaray/Documents/Research/clic_lfv_tau",  // "../../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", // string weight = "reweight",
                //double xmin = 123, double xmax = 129, int nbin = 25  // m_emu_fsr
                //double xmin = -0.5, double xmax = 0.4, int nbin = 250    // BDT
                double xmin = 0.3, double xmax = 0.65, int nbin = 100,    // BDTG // nbin = 250
                double error_bkg = 0.1, double error_sgn = 0.1 // Systematics
                ) {

    bool draw = false;
    bool norm = false;                // Normalize the Histograms (for plot only)
    bool verbose = true; // false
    
    bool cut_variables = true;
    bool cut_bdt = true;
    string cut_br[2] = {"1.0","1.0"};  // Signal, Background   // 6.1e-5

    // Reading our trees
    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;
    }
    if (verbose) {cout << Colors::Verbose; gDirectory->ls();}
    
    
    // Generating Histogram
    cout << Colors::Info << "Generating Histograms" << endl;
    TH1D * samphist[2];
    string hname[2]={"bdt_sig","bdt_bkg"};
    
    ///////////////// FIXING WEIGHT //////////////////
    /*
    string cutstr = "";
    string cutstr_weight_sgn = "*61.38";
    string cutstr_br = "*6.1e-5";
    string cutstr_weight_bkg = "*1.39";
    
    
    
    string cutstr_bdt = varname + varcut;
    string cutstr_mass = "m_emu_fsr<140&&m_emu_fsr>105";
    
    if(cut_mass && cut_bdt) cutstr = "(" + cutstr_bdt + "&&" + cutstr_mass + ")";
    else if(cut_bdt && !cut_mass) cutstr = "(" + cutstr_bdt + ")";
    else if(!cut_bdt && cut_mass) cutstr = "(" + cutstr_mass + ")";
    
    string cutstr_sgn_final;
    if(cut_br) cutstr_sgn_final = cutstr + cutstr_weight_sgn + cutstr_br;
    else cutstr_sgn_final = cutstr + cutstr_weight_sgn;
    
    string cutstr_good[2] = {cutstr_sgn_final, cutstr + cutstr_weight_bkg};   
    */    
    /////////////////////////////////////////////////////////////////////////////////

    string cutstr[2] = {"",""};
    string cutstr_bdt= Cut_varname + varcut;
    string cutstr_var= "costheta_higgs>-0.95 && beta_higgs<0.95 && delta_R>0.25";
    
    for(int i=0;i<2;i++){
        if(cut_variables && cut_bdt) cutstr[i] = weight + "*(" + cutstr_bdt + " && " + cutstr_var + ")";
        else if (cut_variables && !cut_bdt) cutstr[i] = weight + "*(" + cutstr_var + ")";
        else if (!cut_variables && cut_bdt) cutstr[i] = weight + "*(" + cutstr_bdt  + ")";
        else cutstr[i] = weight + "*(" + Analysis_varname + ">-99999999 )";
        cutstr[i] = cutstr[i] + "*" + cut_br[i];
    }
    
    
    for(int j=0;j<2;j++){
        samphist[j]= new TH1D(hname[j].c_str(),"BDT response",nbin,xmin,xmax);
    }
    
    cout << "Cut Selection Conditional String Sgn: " << cutstr[0] << endl;
    cout << "Cut Selection Conditional String Bkg: " << cutstr[1] << endl;
    for(int i=0; i<2; i++){
        samptree[i]->Project(hname[i].c_str(),Analysis_varname.c_str(), cutstr[i].c_str());
    }
    
    // Signal
    TCanvas *c1 = new TCanvas("ch1","BDT histograms",20,20,600,450);
    if( gDirectory->FindObject("rsignal")!=NULL) gDirectory->Delete("rsignal");
    TH1D *rsignal = new TH1D("rsignal","BDT response",nbin,xmin,xmax);
    rsignal->Add(samphist[0]);
    
    rsignal->SetLineWidth(2);
    rsignal->SetLineColor(4);
    cout << "Signal Histogram Created" << endl;
    
    if (verbose){
        cout << Colors::Verbose;
        for(int ibin=0; ibin<nbin; ibin++){
            double bin_content = rsignal->GetBinContent(ibin);
            if (bin_content!=0) cout << "Signal " << ibin << " bin content is: " << bin_content << endl;
        }
    }
    
    
    // Background
    TCanvas *c2 = new TCanvas("ch2","BDT histograms",20,20,600,450);
    if( gDirectory->FindObject("rbg")!=NULL) gDirectory->Delete("rbg");
    TH1D *rbg = new TH1D("rbg","BDT response",nbin,xmin,xmax);
    rbg->Add(samphist[1]);
    
    rbg->SetLineWidth(2);
    rbg->SetLineColor(2);
    cout << "Background Histogram Created" << endl;

    
    if (verbose){
        cout << Colors::Verbose;
        int counter_bins_background = 0;
        for(int ibin=0; ibin<nbin; ibin++){
            double bin_content = rbg->GetBinContent(ibin);
            if (bin_content!=0) {
                cout << "Background " << ibin << " bin content is: " << bin_content << endl;
                // Add a counter here to check the number of "used" background bins
                counter_bins_background++;
            }
        }
        cout << "THE TOTAL NUMBER OF USED BKG BINS IS: " << counter_bins_background << endl;
    }

    
    // Experiment Sample
    TCanvas *cc = new TCanvas("cc","BDT histograms",20,20,600,450);
    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
         }
         cout << "number of observed events= " << nobs << endl;
        rdata->SetBinContent(ibin,nobs);
     }
     //exit(0);

    /*if( gDirectory->FindObject("rsum")!=NULL) gDirectory->Delete("rsum");
    TH1D *rsum = new TH1D("rsum","BDT response",nbin,xmin,xmax);
    rsum->Add(rbg);
    rsum->Add(rsignal);
    for(int ibin=0;ibin<nbin;ibin++){
        double nexp = rsum->GetBinContent(ibin);
        double nerr = rsum->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
        }
        //cout << "number of observed events= " << nobs << endl;
        rdata->SetBinContent(ibin,nobs);
	}*/

    rdata->SetLineWidth(2);
    rdata->SetMarkerStyle(20);
    cout << "Experiment Sample Histogram Created" << endl;
    
    if (verbose){
        cout << Colors::Verbose;
        for(int ibin=0; ibin<nbin; ibin++){
            double bin_content = rdata->GetBinContent(ibin);
            if (bin_content!=0) cout << "Experiment " << ibin << " bin content is: " << bin_content << endl;
        }
    }

    
    if(draw){
        if (norm) {
            rbg->Scale(1.0/rbg->Integral());
            rsignal->Scale(1.0/rsignal->Integral());
            rdata->Scale(1.0/rdata->Integral());
        }
        
        cout << "Number of background events = " << rbg->Integral()<<endl;
        cout << "Number of signal events = "     << rsignal->Integral()<<endl;
        cout << "Number of data (bkg) events = " << rdata->Integral()<<endl;
        
        //auto legend = new TLegend(0.65,0.5,0.95,0.7);
        auto legend = new TLegend(0.15,0.5,0.4,0.7);
        legend->AddEntry(rbg,"background","f");
        legend->AddEntry(rsignal,"signal","l");
        legend->AddEntry(rdata,"data","lep");
    
        rsignal->Draw("hist");
        rdata ->Draw("E same");
        rbg->Draw("hist same");      
        
        

        legend->Draw("same");
        cc->Draw();
        cc->Print((outfolder+"r_Data.png").c_str());

    }
  
    // ====================================================================================
    // RooStats starts here
    // ====================
    
    // Model building based on Exercise 8 from
    //     https://twiki.cern.ch/twiki/bin/view/RooStats/RooStatsExercisesMarch2015
    if (!draw){
        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);
    
        RooStats::HistFactory::Sample signal("signal");
        //signal.AddNormFactor("BR",0.0001, 0.0,0.0008);
        signal.AddNormFactor("BR",0.3,0.,5.);
        signal.SetHisto(rsignal); 
        //signal.AddOverallSys("sgn_unc",1.0 - error_sgn, 1.0 + error_sgn); // Signal systematics
        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);
    
        //BRlimit.Print();
    
        // make the model 
        cout << endl << Colors::Info << "Making the Workspace" << endl;
        
        RooWorkspace * wbr = RooStats::HistFactory::MakeModelAndMeasurementFast(BRlimit);
        
        wbr->SetName("wbr");
        wbr->SetTitle("Workspace for LFV Higgs BR limit");
    
        RooStats::ModelConfig* wbr_config = (RooStats::ModelConfig *) wbr->obj("ModelConfig");
        if(!wbr_config){
            std::cout << "Error: no ModelConfig found in the created workspace ?! "<<  std::endl;
            return;
        }
        if (verbose) wbr->Print();
    
        string outroot;
        outroot = outfolder+outname+".root";
        wbr->writeToFile(outroot.c_str());
        
//////////////////////////////////////////////////////////////////////////////////////////
/// root> StandardHypoTestFCNC("fileName","workspace name","S+B modelconfig name","B model name","data set name",calculator type, test statistic type, use CLS,
///                                number of points, xmin, xmax, number of toys, use number counting)
///
/// type = 0 Freq calculator
/// type = 1 Hybrid calculator
/// type = 2 Asymptotic calculator
/// type = 3 Asymptotic calculator using nominal Asimov data sets (not using fitted parameter values but nominal ones)
///
/// testStatType = 0 LEP
///              = 1 Tevatron
///              = 2 Profile Likelihood two sided
///              = 3 Profile Likelihood one sided (i.e. = 0 if mu < mu_hat)
///              = 4 Profile Likelihood signed ( pll = -pll if mu < mu_hat)
///              = 5 Max Likelihood Estimate as test statistic
///              = 6 Number of observed event as test statistic

        cout << endl << Colors::Info << "Running the Hypothesis test" << endl;
        StandardHypoTestFCNC(outroot.c_str(),"wbr","ModelConfig","","obsData",2,3,true,200,0.,0.00006,10000,true);
    }
}

void callforLimits(){

    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");
    
}