#include "RooWorkspace.h"
#include "RooHistPdf.h"
#include "RooRealVar.h"
#include "RooSimultaneous.h"
#include "RooExtendPdf.h"
#include "RooDataHist.h"
#include "RooDataSet.h"
#include "RooCategory.h"
#include "RooConstVar.h"


using namespace RooFit;
using std::cout;
using std::endl;

void ROOT_10093(){

  // workspace
  RooWorkspace* w = new RooWorkspace("Workspace");

  cout << "workspace created" << endl;
  
  // variables
  RooRealVar x("x", "x", -1, 1);
  x.setBins(2);

  RooRealVar y("y", "y", -1, 1);
  y.setBins(2);

  RooArgSet vars(x, y, "vars");
  
  // categories
  RooCategory cat1("cat1", "cat1");
  cat1.defineType("cat1Negative", 0);
  cat1.defineType("cat1Positive", 1);
  
  RooCategory cat2("cat2", "cat2");
  cat2.defineType("cat2Negative", 0);
  cat2.defineType("cat2Positive", 1);
  
  RooArgSet varsAndCats(x, y, cat1, cat2, "VarsAndCats");

  cout << "variables set" << endl;


  //-------------- Create Pdfs ------------------
  // create negative x negative y pdf
  RooDataSet xNegYNegDS("xNegYNegDS", "xNegYNegDS", vars);
  x.setVal(-0.5);
  y.setVal(-0.5);
  for(int i = 0; i < 1; i++){xNegYNegDS.add(vars);}
  RooDataHist xNegYNegDH("xNegYNegDH", "xNegYNegDH", vars,
                         xNegYNegDS);
  RooHistPdf xNegYNegPdf("xNegYNegPdf", "xNegYNegPdf", vars,
                         xNegYNegDH);

  // create negative x positive y pdf
  RooDataSet xNegYPosDS("xNegYPosDS", "xNegYPosDS", vars);
  x.setVal(-0.5);
  y.setVal(0.5);
  for(int i = 0; i < 2; i++){xNegYPosDS.add(vars);}
  RooDataHist xNegYPosDH("xNegYPosDH", "xNegYPosDH", vars,
                         xNegYPosDS);
  RooHistPdf xNegYPosPdf("xNegYPosPdf", "xNegYPosPdf", vars,
                         xNegYPosDH);

    // create positive x negative y pdf
  RooDataSet xPosYNegDS("xPosYNegDS", "xPosYNegDS", vars);
  x.setVal(0.5);
  y.setVal(-0.5);
  for(int i = 0; i < 3; i++){xPosYNegDS.add(vars);}
  RooDataHist xPosYNegDH("xPosYNegDH", "xPosYNegDH", vars,
                         xPosYNegDS);
  RooHistPdf xPosYNegPdf("xPosYNegPdf", "xPosYNegPdf", vars,
                         xPosYNegDH);

    // create positive x positive y pdf
  RooDataSet xPosYPosDS("xPosYPosDS", "xPosYPosDS", vars);
  x.setVal(0.5);
  y.setVal(0.5);
  for(int i = 0; i < 4; i++){xPosYPosDS.add(vars);}
  RooDataHist xPosYPosDH("xPosYPosDH", "xPosYPosDH", vars,
                         xPosYPosDS);
  RooHistPdf xPosYPosPdf("xPosYPosPdf", "xPosYPosPdf", vars,
                         xPosYPosDH);

  cout << "sub pdfs created" << endl;


  //----------- Extend the PDFs to fix their relative fractions --------

  RooExtendPdf xNegYNegExt("xNegYNegExt", "xNegYNegExt", xNegYNegPdf,
                           RooConst(1));

  RooExtendPdf xNegYPosExt("xNegYPosExt", "xNegYPosExt", xNegYPosPdf,
                           RooConst(2));

  RooExtendPdf xPosYNegExt("xPosYNegExt", "xPosYNegExt", xPosYNegPdf,
                           RooConst(3));

  RooExtendPdf xPosYPosExt("xPosYPosExt", "xPosYPosExt", xPosYPosPdf,
                           RooConst(4));


  //--------------- Combine Pdfs into RooSimultaneous ----------------
  // create RooSimultaneous
  RooSimultaneous simulXNegPdf("simulXNegPdf", "simulXNegPdf",
    RooArgList(xNegYNegExt, xNegYPosExt), cat2);

  RooSimultaneous simulXPosPdf("simulXPosPdf", "simulXPosPdf",
    RooArgList(xPosYNegExt, xPosYPosExt), cat2);


  // create top-level simultaneous
  RooSimultaneous simulPdf("simulPdf", "simulPdf",
      RooArgList(simulXNegPdf, simulXPosPdf), cat1);
      
  auto& superCat = simulPdf.indexCat();
  
  // Check result:
  std::cout << "\n----------------------------\nThe simultaneous is:\n";
  simulPdf.Print("T");
  std::cout << "\n----------------------------\nThe categories are:\n";
  superCat.Print("V");


  //------------------ Generate Events --------------------------

  // generate the actual data set
  cout << "The current super cat is: " << superCat.getLabel() << endl;
  RooDataSet* genNegDS = simulPdf.generate(
    RooArgSet(x, y, cat1, cat2), AutoBinned(0), NumEvents(20), Verbose(1));
  
  std::cout << "\n----------------------------\nThe dataset is:\n";  
  genNegDS->Print("V");

  // look at events
  std::cout << "\n----------------------------\nThe first 20 events:\n";  
  for(int i = 0; i < std::min(20, genNegDS->numEntries()); i++){
    const RooArgSet* tmpSet = genNegDS->get(i);
    std::cout << i << ":\n";
    tmpSet->Print("V");
  }
  
}