#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;

void run_diagnosis_fullmodel(TString workspace_filename = TString("./debug_workspace.root"), TString my_model = TString("sb_model"), bool tryfit = false);

void run_diagnosis_fullmodel(TString workspace_filename, TString my_model, bool tryfit){

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

  RooWorkspace *myworkspace = (RooWorkspace*) f1->Get("debug_WS");
  if(!myworkspace){
    cout << "Error: workspace not found" << endl;
    exit(0);
  }

  UseNLLOffset(true);

  myworkspace->defineSet("poi", "sig_xs_fb");

  ModelConfig sb_config("sb_config", myworkspace);
  sb_config.SetPdf(*(myworkspace->pdf(my_model.Data())));
  sb_config.SetProtoData(*(myworkspace->data("CMS_data_VV")));
  sb_config.SetObservables("mJJ");
  sb_config.SetParametersOfInterest("sig_xs_fb");
  sb_config.SetNuisanceParameters("P0_WZ_HP,P0_WZ_LP,P2_WZ_HP,P1_WZ_LP,P2_WZ_LP,alpha_tau21");
  myworkspace->var("sig_xs_fb")->setVal(1.);
  sb_config.SetSnapshot(*(myworkspace->set("poi")));

  ModelConfig b_config("b_config", myworkspace);
  b_config.SetPdf(*(myworkspace->pdf(my_model.Data())));
  b_config.SetProtoData(*(myworkspace->data("CMS_data_VV")));
  b_config.SetObservables("mJJ");
  b_config.SetParametersOfInterest("sig_xs_fb");
  b_config.SetNuisanceParameters("P0_WZ_HP,P0_WZ_LP,P2_WZ_HP,P1_WZ_LP,P2_WZ_LP,alpha_tau21");
  myworkspace->var("sig_xs_fb")->setVal(0.);
  b_config.SetSnapshot(*(myworkspace->set("poi")));

  AsymptoticCalculator *myCalc = new AsymptoticCalculator(*myworkspace->data("CMS_data_VV"), b_config, sb_config);
  myCalc->SetPrintLevel(2);
  myCalc->Initialize();

  myworkspace->saveSnapshot("asimovfit",myworkspace->allVars());
  const RooArgSet * asimovfit = myworkspace->getSnapshot("asimovfit");

  // ---- The following commented-out block of code will plot two plots for each category.
  //
  // Plot 1: The Asimov data and the best fit pdf to the Asimov data
  // Plot 2: The real data and the best fit pdf to the signal data.
  //
  // It requires the following changes to AsymptoticCalculator.h for it to work:
  //
  // Change 1: replace
  //  122       const RooArgSet & GetBestFitParams() const { return fBestFitPoi; }
  // with:
  //  122       const RooArgSet & GetBestFitParams() const { return fBestFitParams; }
  //
  // Change 2: Add new method:
  //       RooAbsData & GetAsimovData() const { return *fAsimovData; }


  /*
  RooSimultaneous* simPdf = dynamic_cast<RooSimultaneous*>(myworkspace->pdf("sim_sb_model"));
  if (!simPdf){
    cout << "Not Sim" << endl;
    exit(0);
  }

  RooCategory& channelCat = (RooCategory&)simPdf->indexCat();
  int nrIndices = channelCat.numTypes();
  cout << nrIndices << endl;

  for(int i = 0; i < nrIndices; i++){

    myworkspace->allVars().assignValueOnly(*asimovfit);
    channelCat.setIndex(i);
    TString data_cut = TString("signal_cats==signal_cats::") + TString(channelCat.getLabel());

    TCanvas *c2 = new TCanvas();
    c2->SetLogy();
    RooPlot* mJJframe2 = myworkspace->var("mJJ")->frame();
    TString title = TString(channelCat.getLabel()) + " Asimov my binning";
    mJJframe2->SetTitle(title.Data());
    myCalc->GetAsimovData().plotOn(mJJframe2,Cut(data_cut.Data()),Binning(""));
    myworkspace->pdf("sim_sb_model")->plotOn(mJJframe2,Slice(*myworkspace->cat("signal_cats"),channelCat.getLabel()),ProjWData(*myworkspace->cat("signal_cats"),myCalc->GetAsimovData()));
    mJJframe2->SetAxisRange(0.01,1e5,"Y");
    mJJframe2->Draw();

    TCanvas *c1 = new TCanvas();
    c1->SetLogy();
    RooPlot* mJJframe = myworkspace->var("mJJ")->frame();
    title = TString(channelCat.getLabel()) + " real data";
    mJJframe->SetTitle(title.Data());
    RooArgSet testing(myCalc->GetBestFitParams());
    myworkspace->allVars().assignValueOnly(testing);
    myworkspace->data("CMS_data_VV")->plotOn(mJJframe,Cut(data_cut.Data()));
    myworkspace->pdf("sim_sb_model")->plotOn(mJJframe,Slice(*myworkspace->cat("signal_cats"),channelCat.getLabel()),ProjWData(*myworkspace->cat("signal_cats"),*myworkspace->data("CMS_data_VV")));
    mJJframe->SetAxisRange(0.01,1e5,"Y");
    mJJframe->Draw();

  }
  */


  if(tryfit) myworkspace->pdf("sb_model")->fitTo(*myworkspace->data("CMS_data_VV"));

}