#include "TCanvas.h"
#include "TROOT.h"
#include "RooPlot.h"
#include "RooDataHist.h"
#include "TH1.h"
#include "RooRealVar.h"
#include "RooStats/AsymptoticCalculator.h"
#include "RooWorkspace.h"
#include "RooProfileLL.h"
#include "RooStats/ProfileLikelihoodCalculator.h"
#include <string>
#include <cstdlib>
#include <vector>
#include <map>
#include "TVectorT.h"
#include <iostream>
#include <fstream>
#include <sstream>
#include <math.h>
#include "RooStats/HistFactory/FlexibleInterpVar.h"

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

struct exclusion {
  double mWp;
  double limit;
  double exp_limit;
  double exp_limit_p1;
  double exp_limit_m1;
  double exp_limit_p2;
  double exp_limit_m2;
};

RooWorkspace *readWorkspace(TString name);
RooWorkspace *createBasicCMSICHEPWorkspace();
RooWorkspace *completeWorkspace(TString signal_filename = TString("./WZ_2000/CMShadronic.root"));
void saveWorkspace(TString output_filename = TString("debug_workspace.root"), TString signal_filename = TString("./WZ_2000/CMShadronic.root"));

void saveWorkspace(TString output_filename, TString signal_filename){

  RooWorkspace *myworkspace = completeWorkspace(signal_filename);

  TFile* fout = new TFile(output_filename.Data(),"RECREATE");
  if(!fout){
    cout << "Error: file " << output_filename << " could not be created." << endl;
    exit(0);
  }
  
  myworkspace->Write();
  fout->Close();

}

RooWorkspace *completeWorkspace(TString signal_filename){

  string WZ_HP_string("WZ_HP");
  string WZ_LP_string("WZ_LP");

  RooWorkspace *myworkspace = createBasicCMSICHEPWorkspace();

  RooCategory signal_cats = *myworkspace->cat("signal_cats");
  RooRealVar mJJ_arg = *myworkspace->var("mJJ");

  TFile* f1 = new TFile(signal_filename.Data());
  if(!f1){
    cout << "Error: file " << signal_filename << " not found." << endl;
    exit(0);
  }

  vector<TString> variable_variables;
  variable_variables.emplace_back("tau21");

  vector<TString> postscripts;
  postscripts.emplace_back("p");
  postscripts.emplace_back("m");

  vector<TString> mass_categories;
  mass_categories.emplace_back("WZ");

  vector<TString> purity_categories;
  purity_categories.emplace_back("HP");
  purity_categories.emplace_back("LP");

  TString twopar_form("( @0 / 13000 )^@1");
  TString threepar_form("(1 - ( @0 / 13000 ))**@1 * ( @0 / 13000 )**@2");

  RooRealVar P0_WZ_HP("P0_WZ_HP", "P0_WZ_HP", 4400., 440.,44000.);
  RooRealVar P0_WZ_LP("P0_WZ_LP", "P0_WZ_LP", 26000., 2600.,260000.);

  RooRealVar P2_WZ_HP("P2_WZ_HP","P2_WZ_HP",-8,-50,50);

  RooRealVar P1_WZ_LP("P1_WZ_LP","P1_WZ_LP",16.,-50,50);
  RooRealVar P2_WZ_LP("P2_WZ_LP","P2_WZ_LP",-7.,-50,50);

  RooGenericPdf WZ_HP_b_model_proto("WZ_HP_b_model_proto", "WZ_HP_b_model_proto", twopar_form.Data(), RooArgList(mJJ_arg, P2_WZ_HP));
  RooGenericPdf WZ_LP_b_model_proto("WZ_LP_b_model_proto", "WZ_LP_b_model_proto", threepar_form.Data(), RooArgList(mJJ_arg, P1_WZ_LP, P2_WZ_LP));

  RooExtendPdf WZ_HP_b_model("WZ_HP_b_model","WZ_HP_b_model",WZ_HP_b_model_proto,P0_WZ_HP);
  RooExtendPdf WZ_LP_b_model("WZ_LP_b_model","WZ_LP_b_model",WZ_LP_b_model_proto,P0_WZ_LP);

  map<string,RooAbsPdf*> bg_pdf_map;
  bg_pdf_map[WZ_HP_string] = &WZ_HP_b_model;
  bg_pdf_map[WZ_LP_string] = &WZ_LP_b_model;

  RooRealVar sig_xs_fb("sig_xs_fb","sig_xs_fb",10.,0.,5000.);

  map<TString, TH1*> signal_nom_histmap;
  map<TString, RooDataHist*> signal_nom_datahist;
  map<TString, RooHistFunc*> signal_nom_histpdf;
  double total_nom = 0;
  map<TString, RooRealVar*> normfactor_vars;

  vector<map<TString,TH1*>> signal_p_histmaps(1);
  vector<map<TString,RooDataHist*>> signal_p_datahists(1);
  vector<map<TString,RooHistFunc*>> signal_p_histpdfs(1);
  vector<map<TString,TH1*>> signal_m_histmaps(1);
  vector<map<TString,RooDataHist*>> signal_m_datahists(1);
  vector<map<TString,RooHistFunc*>> signal_m_histpdfs(1);

  map<TString,PiecewiseInterpolation*> signal_pw;
  map<TString,RooRealSumPdf*> signal_pwpdf;

  RooRealVar alpha_tau21("alpha_tau21","alpha_tau21",0.,-1,1);

  RooArgList alpha_set(alpha_tau21);

  for(int i = 0; i < mass_categories.size(); i++){
    for(int j = 0; j < purity_categories.size(); j++){
      TString category = mass_categories[i] + "_" + purity_categories[j];
      signal_nom_histmap[category] = (TH1F*)f1->Get(category.Data());
      if(!signal_nom_histmap[category]){
        cout << "Error: histogram " << category << " not found." << endl;
        exit(0);
      }
      
      TString datahist_name = category + "_nom_datahist";
      signal_nom_datahist[category] = new RooDataHist(datahist_name.Data(),datahist_name.Data(),RooArgList(mJJ_arg),signal_nom_histmap[category]);
      total_nom += signal_nom_datahist[category]->sum(false);

      TString histpdf_name = category + "_nom_histpdf";
      signal_nom_histpdf[category] = new RooHistFunc(histpdf_name.Data(),histpdf_name.Data(),RooArgList(mJJ_arg),*signal_nom_datahist[category]);

      for(int k = 0; k < variable_variables.size(); k++){
	TString histname_p = category + "_" + variable_variables[k] + "_p";
	signal_p_histmaps[k][category] = (TH1F*)f1->Get(histname_p.Data());
	if(!signal_p_histmaps[k][category]){
	  cout << "Error: histogram " << category << " not found." << endl;
	  exit(0);
	}
        TString histname_m = category + "_" + variable_variables[k] + "_m";
        signal_m_histmaps[k][category] = (TH1F*)f1->Get(histname_m.Data());
        if(!signal_m_histmaps[k][category]){
          cout << "Error: histogram " << category << " not found." << endl;
          exit(0);
        }

	TString datahistname_p =  category + "_" + variable_variables[k] + "_p_datahist";
	TString datahistname_m =  category + "_" + variable_variables[k] + "_m_datahist";
	signal_p_datahists[k][category] = new RooDataHist(datahistname_p.Data(),datahistname_p.Data(),RooArgList(mJJ_arg),signal_p_histmaps[k][category]);
	signal_m_datahists[k][category] = new RooDataHist(datahistname_m.Data(),datahistname_m.Data(),RooArgList(mJJ_arg),signal_m_histmaps[k][category]);

	TString histpdfname_p = category + "_" + variable_variables[k] + "_p_histpdf";
	TString histpdfname_m = category + "_" + variable_variables[k] + "_m_histpdf";
	signal_p_histpdfs[k][category] = new RooHistFunc(histpdfname_p.Data(),histpdfname_p.Data(),RooArgList(mJJ_arg),*signal_p_datahists[k][category]);
	signal_m_histpdfs[k][category] = new RooHistFunc(histpdfname_m.Data(),histpdfname_m.Data(),RooArgList(mJJ_arg),*signal_m_datahists[k][category]);
      }

      RooArgSet *lowset = new RooArgSet(*signal_m_histpdfs[0][category]);
      RooArgSet *highset = new RooArgSet(*signal_p_histpdfs[0][category]);
      TString pw_name = category + "_pw";
      signal_pw[category] = new PiecewiseInterpolation(pw_name, pw_name, *signal_nom_histpdf[category], *lowset, *highset, alpha_set);
      signal_pw[category]->setAllInterpCodes(4);
      signal_pw[category]->setPositiveDefinite(kTRUE);
      TString pwpdf_name = pw_name + "pdf";
      signal_pwpdf[category] = new RooRealSumPdf(pwpdf_name.Data(),pwpdf_name.Data(),*signal_pw[category],RooArgList());
    }
  }

  RooRealVar norm_factor_nom("norm_factor_nom","norm_factor_nom",total_nom);

  total_nom = 10000;
  map<TString,FlexibleInterpVar*> interpvar;
  vector<double> event_totals(1);
  map<TString,double> nom_eff;
  map<TString,vector<double>> var_effs_m;
  map<TString,vector<double>> var_effs_p; 
  map<TString,RooFormulaVar*> XS_formula;
  map<TString,RooExtendPdf*> extended_component;
  map<string,RooAbsPdf*> added_component;

  for(int i = 0; i < mass_categories.size(); i++){
    for(int j = 0; j < purity_categories.size(); j++){
      TString category = mass_categories[i] + "_" + purity_categories[j];
      TString interpvar_name = category + "_interpvar";
      nom_eff[category] = signal_nom_datahist[category]->sum(false)/total_nom;
      vector<double> temp_var_effs_m;
      vector<double> temp_var_effs_p;
      for(int k = 0; k < variable_variables.size(); k++){
	temp_var_effs_m.emplace_back(signal_m_datahists[k][category]->sum(false)/total_nom);
	temp_var_effs_p.emplace_back(signal_p_datahists[k][category]->sum(false)/total_nom);
      }
      var_effs_m[category] = temp_var_effs_m;
      var_effs_p[category] = temp_var_effs_p;

      interpvar[category] = new FlexibleInterpVar(interpvar_name.Data(),interpvar_name.Data(),alpha_set,nom_eff[category],var_effs_m[category],var_effs_p[category]);
      interpvar[category]->setAllInterpCodes(4);

      TString formulavar_name = category + "_XSform";
      TString formula = "sig_xs_fb * 12.9 * 0.69 * 0.67 * " + interpvar_name;
      XS_formula[category] = new RooFormulaVar(formulavar_name.Data(),formulavar_name.Data(),formula,RooArgList(sig_xs_fb,*interpvar[category]));

      TString extended_name = category + "_extendedpdf";
      extended_component[category] = new RooExtendPdf(extended_name.Data(),extended_name.Data(),*signal_pwpdf[category],*XS_formula[category]);

      TString added_name = category + "_sb";
      string category_string(category.Data());
      added_component[category_string] = new RooAddPdf(added_name.Data(),added_name.Data(),RooArgList(*extended_component[category],*bg_pdf_map[category_string]));
    }
  }


  RooSimultaneous *sim_sb_model = new RooSimultaneous("sim_sb_model","sim_sb_model", added_component, signal_cats);
  

  RooRealVar alpha_m("alpha_m","alpha_m",0.);
  RooRealVar alpha_sd("alpha_sd","alpha_sd",1.);
  RooGaussian gauss_tau21("gauss_tau21","gauss_tau21",alpha_tau21,alpha_m,alpha_sd);
  RooArgSet gaussians(gauss_tau21);

  RooProdPdf sb_model("sb_model","sb_model",gaussians,Conditional(*sim_sb_model,mJJ_arg));
  myworkspace->import(sb_model);

  RooFIter iter = myworkspace->components().fwdIterator();
  RooAbsArg* myarg;
  while ((myarg = iter.next())) {
    if (myarg->IsA() == RooRealSumPdf::Class()) {
      myarg->setAttribute("BinnedLikelihood");
      cout << "Activating binned likelihood attribute for " << myarg->GetName() << endl;
    }
  }


  f1->Close();
  return myworkspace;
}

RooWorkspace *createBasicCMSICHEPWorkspace(){
  RooRealVar mJJ_arg("mJJ","mJJ",955,7000,"GeV");

  int nbins_VV = 32;

  double VV_bins[] = {955, 1000, 1058, 1118, 1181, 1246, 1313, 1383, 1455, 1530, 1607,
                         1687, 1770, 1856, 1945, 2037, 2132, 2231, 2332, 2438, 2546, 2659,
                         2775, 2895, 3019, 3147, 3279, 3416, 3558, 3704, 3854, 4010, 4171};

  RooBinning VV_binning(nbins_VV,VV_bins);
  mJJ_arg.setBinning(VV_binning);

  int WZ_HP_datapoints_list[] = {1335, 1017, 701, 431, 311, 193, 122, 102, 67, 33, 32, 19, 10, 11, 9, 
				 2, 5, 2, 1, 0, 0, 2};
  int WZ_LP_datapoints_list[] = {7219, 6393, 4146, 2846, 1916, 1281, 832, 577, 415, 287, 197, 144, 98,
				 49, 42, 44, 27, 16, 11, 8, 7, 2, 3, 4, 0, 3, 1, 1, 0, 0, 0, 1};

  TH1I *WZ_HP_data_TH1 = new TH1I("WZ_HP_data_TH1","WZ_HP_data_TH1", nbins_VV,VV_bins);
  TH1I *WZ_LP_data_TH1 = new TH1I("WZ_LP_data_TH1","WZ_LP_data_TH1", nbins_VV,VV_bins);

  vector<int> WZ_HP_datapoints(WZ_HP_datapoints_list, WZ_HP_datapoints_list + sizeof(WZ_HP_datapoints_list) / sizeof(int));
  for(int i = 0; i < WZ_HP_datapoints.size(); i++){
    WZ_HP_data_TH1->AddBinContent(i+1,WZ_HP_datapoints[i]);
  }

  vector<int> WZ_LP_datapoints(WZ_LP_datapoints_list, WZ_LP_datapoints_list + sizeof(WZ_LP_datapoints_list) / sizeof(int));
  for(int i = 0; i < WZ_LP_datapoints.size(); i++){
    WZ_LP_data_TH1->AddBinContent(i+1,WZ_LP_datapoints[i]);
  }

  RooCategory signal_cats("signal_cats","signal_cats");
  signal_cats.defineType("WZ_HP");
  signal_cats.defineType("WZ_LP");

  map<string, TH1*> histmap;
  string WZ_HP_string("WZ_HP");
  string WZ_LP_string("WZ_LP");
  histmap[WZ_HP_string] = WZ_HP_data_TH1;
  histmap[WZ_LP_string] = WZ_LP_data_TH1;

  RooDataHist CMS_data_VV("CMS_data_VV","CMS_data_VV", RooArgList(mJJ_arg), signal_cats, histmap);

  RooWorkspace *myworkspace = new RooWorkspace("debug_WS");
  myworkspace->import(CMS_data_VV);

  return myworkspace;
}

RooWorkspace *readWorkspace(TString name){
  TString path("/nfs/thry/drpv/triboson_scans/roottests/");
  TString file_name = name + ".root";
  TFile *myfile = new TFile(file_name.Data());
  RooWorkspace *temp = (RooWorkspace*) myfile->Get(name.Data());
  myfile->Close();
  return temp;
}

/*
void addSignalModeltoWorkspace(RooWorkspace *ws, vector<RooDataHist> datavect){
  
}
*/