//////////////////////////////////////
// RooStats Model Inspector
// Author Kyle Cranmer <cranmer@cern.ch>
// Version 1, October 2011
//   - based on tutorial macro by Bertrand Bellenot, Ilka Antcheva
// Version 2, November 2011
//   - fixes from Bertrand Bellenot <Bertrand.Bellenot@cern.ch> for scrolling window for many parameters
//
//
// Usage:
// The usage is the same as the StandardXxxDemo.C macros.
// The macro expects a root file containing a workspace with a ModelConfig and a dataset
// $ root
// .L ModelInspector.C+
// ModelInspector(fileName, workspaceName, modelConfigName, dataSetName);
//
// Drag the sliders to adjust the parameters of the model.
// the min and max range of the sliders are used to define the upper & lower variation
// the pointer position of the slider is the central blue curve.
//
// Click the FIT button to 
//
// To Do:
//  - check boxes to specify which nuisance parameters used in making variation
//  - a button to make the profile inspector plots
//  - a check button to use MINOS errors
//  - have fit button show the covariance matrix from the fit
//  - a button to make the log likelihood plots
//  - a dialog to open the desired file
//  - ability to see teh signal and background contributions?
//   
   
#include "TGButton.h"
#include "TRootEmbeddedCanvas.h"
#include "TGLayout.h"
#include "TF1.h"
#include "TMath.h"
#include "TCanvas.h"
#include "TGTextEntry.h"
#include "TGLabel.h"
#include "TGTripleSlider.h"
#include "RooWorkspace.h"
#include "RooStats/ModelConfig.h"
#include "TFile.h"
#include "RooArgSet.h"
#include "TList.h"
#include "RooAbsPdf.h"
#include "RooRealVar.h"
#include "RooPlot.h"
#include "TGButton.h"
#include <map>
#include "RooFitResult.h"
#include "TROOT.h"
#include <TApplication.h>
#include "RooSimultaneous.h"
#include "RooCategory.h"

enum ETestCommandIdentifiers {
   HId1,
   HId2,
   HId3,
   HCId1,
   HCId2,

   HSId1
};

using namespace RooFit;
using namespace RooStats;
class ModelInspectorGUI : public TGMainFrame {

private:
   TRootEmbeddedCanvas *fCanvas;
   TGLayoutHints       *fLcan;
   TF1                 *fFitFcn;
  RooPlot* fPlot;
  RooWorkspace* fWS;  
  TFile* fFile;
  ModelConfig* fMC;
  RooAbsData* fData;
  RooFitResult* fFitRes;

  TList fSliderList;
  TList fFrameList;
  vector<RooPlot*> fPlotList;
  map<TGTripleHSlider*,const char*> fSliderMap;
  map<TGTripleHSlider*,TGLabel*> fLabelMap;

  TGButton* fFitButton;
  TGTextButton *fExitButton; 

   // BB: a TGCanvas and a vertical frame are needed for using scrollbars
   TGCanvas *fCan;
   TGVerticalFrame *fVFrame;
  
   TGHorizontalFrame   *fHframe0, *fHframe1, *fHframe2;
   TGLayoutHints       *fBly, *fBfly1, *fBfly2, *fBfly3;
   TGTripleHSlider     *fHslider1;
  //   TGTextEntry         *fTeh1, *fTeh2, *fTeh3;
   TGTextBuffer        *fTbh1, *fTbh2, *fTbh3;
   TGCheckButton       *fCheck1, *fCheck2;

public:
  ModelInspectorGUI(RooWorkspace*, ModelConfig*, RooAbsData*);
   virtual ~ModelInspectorGUI();

   void CloseWindow();
   void DoText(const char *text);
   void DoSlider();
   void DoSlider(const char*);
   void DoFit();
   void DoExit();
   void HandleButtons();

   ClassDef(ModelInspectorGUI, 0)
};

//______________________________________________________________________________
ModelInspectorGUI::ModelInspectorGUI(RooWorkspace* w, ModelConfig* mc, RooAbsData* data) 
  : TGMainFrame(gClient->GetRoot(), 100, 100)
{
  
  RooMsgService::instance().getStream(1).removeTopic(RooFit::NumIntegration);
  fWS = w;
  fMC = mc;
  fData = data;
  RooSimultaneous* simPdf = NULL;
  Int_t numCats = 1;
  if(strcmp(fMC->GetPdf()->ClassName(),"RooSimultaneous")==0){
    cout <<"Is a simultaneous PDF"<< endl;     
    simPdf = (RooSimultaneous*)(fMC->GetPdf());
    RooCategory* channelCat = (RooCategory*) (&simPdf->indexCat());
    cout <<" with " << channelCat->numTypes() << " categories"<<endl;
    numCats = channelCat->numTypes();
  } else {
    cout <<"Is not a simultaneous PDF"<<endl;
  }
  fFitRes=0;
  

  //char buf[32];
   SetCleanup(kDeepCleanup);

   // Create an embedded canvas and add to the main frame, centered in x and y
   // and with 30 pixel margins all around
   fCanvas = new TRootEmbeddedCanvas("Canvas", this, 600, 400);
   fLcan = new TGLayoutHints(kLHintsExpandX | kLHintsExpandY, 10, 10, 10, 10);
   AddFrame(fCanvas, fLcan);
   fPlotList.resize(numCats);
   if(numCats>1){
     fCanvas->GetCanvas()->Divide(numCats);
     for(int i=0; i<numCats; ++i){
       //   fCanvas->GetCanvas()->SetFillColor(33);
       //   fCanvas->GetCanvas()->SetFrameFillColor(41);
       fCanvas->GetCanvas()->cd(i+1)->SetBorderMode(0);
       fCanvas->GetCanvas()->cd(i+1)->SetGrid();
       //   fCanvas->GetCanvas()->SetLogy();
     }
   }

   fHframe0 = new TGHorizontalFrame(this, 0, 0, 0);

   fCheck1 = new TGCheckButton(fHframe0, "&Constrained", HCId1);
   fCheck2 = new TGCheckButton(fHframe0, "&Relative", HCId2);
   fCheck1->SetState(kButtonUp);
   fCheck2->SetState(kButtonUp);
   fCheck1->SetToolTipText("Pointer position constrained to slider sides");
   fCheck2->SetToolTipText("Pointer position relative to slider position");

   fHframe0->Resize(200, 50);


   fHframe2 = new TGHorizontalFrame(this, 0, 0, 0);

   fFitButton = new TGTextButton(fHframe2,"Fit");
   fFitButton->Connect("Clicked()","ModelInspectorGUI",this,"DoFit()");
   fExitButton = new TGTextButton(fHframe2, "Exit ");
   fExitButton->Connect("Clicked()", "ModelInspectorGUI", this, "DoExit()");


   fCheck1->Connect("Clicked()", "ModelInspectorGUI", this,
                    "HandleButtons()");
   fCheck2->Connect("Clicked()", "ModelInspectorGUI", this,
                    "HandleButtons()");

   fHframe2->Resize(100, 25);

   //--- layout for buttons: top align, equally expand horizontally
   fBly = new TGLayoutHints(kLHintsTop | kLHintsExpandX, 5, 5, 5, 5);

   //--- layout for the frame: place at bottom, right aligned
   fBfly1 = new TGLayoutHints(kLHintsTop | kLHintsCenterX, 5, 5, 5, 5);
   fBfly2 = new TGLayoutHints(kLHintsTop | kLHintsLeft,    5, 5, 5, 5);
   fBfly3 = new TGLayoutHints(kLHintsTop | kLHintsRight,   5, 5, 5, 5);

   //   fHframe0->AddFrame(fCheck1, fBfly2);
   //   fHframe0->AddFrame(fCheck2, fBfly2);
   fHframe2->AddFrame(fFitButton, fBfly2);
   fHframe2->AddFrame(fExitButton, fBfly3);

   AddFrame(fHframe0, fBly);
   AddFrame(fHframe2, fBly);


   // Loop over POI & NP, create slider
   // need maps of NP->slider? or just slider->NP
   RooArgSet parameters;
   parameters.add(*fMC->GetParametersOfInterest());
   parameters.add(*fMC->GetNuisanceParameters());
   TIterator* it = parameters.createIterator();
   RooRealVar* param=NULL;
   
   // BB: This is the part needed in order to have scrollbars 
   fCan = new TGCanvas(this, 100, 100, kFixedSize);
   AddFrame(fCan, new TGLayoutHints(kLHintsExpandY | kLHintsExpandX));
   fVFrame = new TGVerticalFrame(fCan->GetViewPort(), 10, 10);
   fCan->SetContainer(fVFrame);
   // And that't it!
   // Obviously, the parent of other subframes is now fVFrame instead of "this"...

   while( (param=(RooRealVar*)it->Next()) ){
     cout <<"Adding Slider for "<<param->GetName()<<endl;
     TGHorizontalFrame* hframek = new TGHorizontalFrame(fVFrame, 0, 0, 0);

     TGLabel* hlabel = new TGLabel(hframek,Form("%s = %.3f +%.3f",param->GetName(),param->getVal(),param->getError()));
     TGTripleHSlider* hsliderk = new TGTripleHSlider(hframek, 190, kDoubleScaleBoth, HSId1,
				     kHorizontalFrame,
				     GetDefaultFrameBackground(),
				     kFALSE, kFALSE, kFALSE, kFALSE);
     hsliderk->Connect("PointerPositionChanged()", "ModelInspectorGUI", 
			this, "DoSlider()");
     hsliderk->Connect("PositionChanged()", "ModelInspectorGUI", 
			this, "DoSlider()");
     hsliderk->SetRange(param->getMin(),param->getMax());
     
     hframek->Resize(200, 25);
     fSliderList.Add(hsliderk);
     fFrameList.Add(hframek);

     hsliderk->SetPosition(param->getVal()-param->getError(),param->getVal()+param->getError());
     hsliderk->SetPointerPosition(param->getVal());

     hframek->AddFrame(hlabel, fBly);
     hframek->AddFrame(hsliderk, fBly);
     fVFrame->AddFrame(hframek, fBly);
     fSliderMap[hsliderk]=param->GetName();
     fLabelMap[hsliderk]=hlabel;
   }

   // Set main frame name, map sub windows (buttons), initialize layout
   // algorithm via Resize() and map main frame
   SetWindowName("RooFit/RooStats Model Inspector");
   MapSubwindows();
   Resize(GetDefaultSize());
   MapWindow();

   DoSlider();
}

//______________________________________________________________________________
ModelInspectorGUI::~ModelInspectorGUI()
{
   // Clean up

   Cleanup();
}

//______________________________________________________________________________
void ModelInspectorGUI::CloseWindow()
{
   // Called when window is closed via the window manager.

   delete this;
}

//______________________________________________________________________________
void ModelInspectorGUI::DoText(const char * /*text*/)
{
   // Handle text entry widgets.

   TGTextEntry *te = (TGTextEntry *) gTQSender;
   Int_t id = te->WidgetId();

   switch (id) {
      case HId1:
         fHslider1->SetPosition(atof(fTbh1->GetString()),
                                fHslider1->GetMaxPosition());
         break;
      case HId2:
         fHslider1->SetPointerPosition(atof(fTbh2->GetString()));
         break;
      case HId3:
         fHslider1->SetPosition(fHslider1->GetMinPosition(),
                                atof(fTbh1->GetString()));
         break;
      default:
         break;
   }
   DoSlider();
}

//______________________________________________________________________________
void ModelInspectorGUI::DoFit()
{
  fFitRes = fMC->GetPdf()->fitTo(*fData,Save());
  map<TGTripleHSlider*,const char*>::iterator it;;
  it = fSliderMap.begin();
  for(; it!=fSliderMap.end(); ++it){
    RooRealVar* param=fWS->var(it->second);
    param = (RooRealVar*) fFitRes->floatParsFinal().find(it->second);
    it->first->SetPosition(param->getVal()-param->getError(),param->getVal()+param->getError());
    it->first->SetPointerPosition(param->getVal());
  }
  DoSlider();

}

//______________________________________________________________________________
void ModelInspectorGUI::DoSlider(const char* text)
{
  cout << "." << text <<endl;
}

//______________________________________________________________________________
void ModelInspectorGUI::DoSlider()
{
   // Handle slider widgets.

   //char buf[32];

  RooSimultaneous* simPdf = NULL;
  Int_t numCats = 0;
  if(strcmp(fMC->GetPdf()->ClassName(),"RooSimultaneous")==0){
    simPdf = (RooSimultaneous*)(fMC->GetPdf());
    RooCategory* channelCat = (RooCategory*) (&simPdf->indexCat());
    numCats = channelCat->numTypes();
  } else {
  }


  /////////////////////////////////////////////
  if(!simPdf){
    /////////////////////////////////////////////
    // if not SimPdf
    /////////////////////////////////////////////

    // pre loop      
    map<TGTripleHSlider*,const char*>::iterator it;;
    delete fPlot;
    fPlot = ((RooRealVar*)fMC->GetObservables()->first())->frame();
    fData->plotOn(fPlot);
    double normCount;

    // high loop
    it = fSliderMap.begin();
    for(; it!=fSliderMap.end(); ++it){
      const char* name = it->second;
      fWS->var(name)->setVal(it->first->GetMaxPosition());
      RooRealVar* param = fWS->var(name);
      fLabelMap[it->first]->SetText(Form("%s = %.3f [%.3f,%.3f]",param->GetName(),it->first->GetPointerPosition(),it->first->GetMinPosition(),it->first->GetMaxPosition()));
    }
    normCount = fMC->GetPdf()->expectedEvents(*fMC->GetObservables());
    fMC->GetPdf()->plotOn(fPlot,LineColor(kRed),Normalization(normCount,RooAbsReal::NumEvent));
    
    // low loop
    it = fSliderMap.begin();
    for(; it!=fSliderMap.end(); ++it){
      const char* name = it->second;
      fWS->var(name)->setVal(it->first->GetMinPosition());
      }
    normCount = fMC->GetPdf()->expectedEvents(*fMC->GetObservables());
    fMC->GetPdf()->plotOn(fPlot,LineColor(kGreen),Normalization(normCount,RooAbsReal::NumEvent));
    
    // central oop
    it = fSliderMap.begin();
    for(; it!=fSliderMap.end(); ++it){
      const char* name = it->second;
      fWS->var(name)->setVal(it->first->GetPointerPosition());
    }
    normCount = fMC->GetPdf()->expectedEvents(*fMC->GetObservables());
    fMC->GetPdf()->plotOn(fPlot,LineColor(kBlue),Normalization(normCount,RooAbsReal::NumEvent));
    fPlot->Draw();
    
    fCanvas->GetCanvas()->Modified();
    fCanvas->GetCanvas()->Update();
    ////////////////////////////////////////////////////////////////////////////
  } else {
    ////////////////////////////////////////////////////////////////////////////
    // else (is simpdf)
    ////////////////////////////////////////////////////////////////////////////
    RooCategory* channelCat = (RooCategory*) (&simPdf->indexCat());
    //    TIterator* iter = simPdf->indexCat().typeIterator() ;
    TIterator* iter = channelCat->typeIterator() ;
    RooCatType* tt = NULL;
    Int_t frameIndex = 0;
    while((tt=(RooCatType*) iter->Next())) {
      ++frameIndex;
      fCanvas->GetCanvas()->cd(frameIndex);

      // pre loop
      RooAbsPdf* pdftmp = simPdf->getPdf(tt->GetName()) ;
      RooArgSet* obstmp = pdftmp->getObservables(*fMC->GetObservables()) ;
      RooRealVar* obs = ((RooRealVar*)obstmp->first());

      fPlot = fPlotList.at(frameIndex-1);
      if(fPlot) delete fPlot;
      fPlot = obs->frame();
      fPlotList.at(frameIndex-1) = fPlot;

      RooFit::MsgLevel msglevel = RooMsgService::instance().globalKillBelow();
      RooMsgService::instance().setGlobalKillBelow(RooFit::WARNING);
      fData->plotOn(fPlot,MarkerSize(1),Cut(Form("%s==%s::%s",channelCat->GetName(),channelCat->GetName(),tt->GetName())),DataError(RooAbsData::None));
      RooMsgService::instance().setGlobalKillBelow(msglevel);


      map<TGTripleHSlider*,const char*>::iterator it;;
      double normCount;

      // high loop
      it = fSliderMap.begin();
      for(; it!=fSliderMap.end(); ++it){
	const char* name = it->second;
	fWS->var(name)->setVal(it->first->GetMaxPosition());
	RooRealVar* param = fWS->var(name);
	fLabelMap[it->first]->SetText(Form("%s = %.3f [%.3f,%.3f]",param->GetName(),it->first->GetPointerPosition(),it->first->GetMinPosition(),it->first->GetMaxPosition()));
      }
      normCount = pdftmp->expectedEvents(*obs);
      pdftmp->plotOn(fPlot,LineColor(kRed),LineWidth(2.),Normalization(normCount,RooAbsReal::NumEvent)) ;


      // low loop
      it = fSliderMap.begin();
      for(; it!=fSliderMap.end(); ++it){
	const char* name = it->second;
	fWS->var(name)->setVal(it->first->GetMinPosition());
	RooRealVar* param = fWS->var(name);
	fLabelMap[it->first]->SetText(Form("%s = %.3f [%.3f,%.3f]",param->GetName(),it->first->GetPointerPosition(),it->first->GetMinPosition(),it->first->GetMaxPosition()));
      }
      normCount = pdftmp->expectedEvents(*obs);
      pdftmp->plotOn(fPlot,LineColor(kGreen),LineWidth(2.),Normalization(normCount,RooAbsReal::NumEvent)) ;

      // central loop
      it = fSliderMap.begin();
      for(; it!=fSliderMap.end(); ++it){
	const char* name = it->second;
	fWS->var(name)->setVal(it->first->GetPointerPosition());
	RooRealVar* param = fWS->var(name);
	fLabelMap[it->first]->SetText(Form("%s = %.3f [%.3f,%.3f]",param->GetName(),it->first->GetPointerPosition(),it->first->GetMinPosition(),it->first->GetMaxPosition()));
      }
      normCount = pdftmp->expectedEvents(*obs);
      if(!fFitRes)
	pdftmp->plotOn(fPlot,LineColor(kBlue),LineWidth(2.),Normalization(normCount,RooAbsReal::NumEvent)) ;
      else{
	pdftmp->plotOn(fPlot,Normalization(normCount,RooAbsReal::NumEvent),VisualizeError(*fFitRes,*fMC->GetNuisanceParameters()),FillColor(kYellow)) ;
	pdftmp->plotOn(fPlot,LineColor(kBlue),LineWidth(2.),Normalization(normCount,RooAbsReal::NumEvent)) ;
      msglevel = RooMsgService::instance().globalKillBelow();
      RooMsgService::instance().setGlobalKillBelow(RooFit::WARNING);
      fData->plotOn(fPlot,MarkerSize(1),Cut(Form("%s==%s::%s",channelCat->GetName(),channelCat->GetName(),tt->GetName())),DataError(RooAbsData::None));
      RooMsgService::instance().setGlobalKillBelow(msglevel);
      }
      fPlot->Draw();
    }    
    fCanvas->GetCanvas()->Modified();
    fCanvas->GetCanvas()->Update();
    ///////////////////////////////////////////
    // end if(simPdf)
  }


}

//______________________________________________________________________________
void ModelInspectorGUI::HandleButtons()
{
   // Handle different buttons.

   TGButton *btn = (TGButton *) gTQSender;
   Int_t id = btn->WidgetId();

   switch (id) {
      case HCId1:
         fHslider1->SetConstrained(fCheck1->GetState());
         break;
      case HCId2:
         fHslider1->SetRelative(fCheck2->GetState());
         break;
      default:
         break;
   }
}
void ModelInspectorGUI::DoExit()
{
   printf("Exit application...");
   gApplication->Terminate(0);
}


void ModelInspector(const char* infile = "",
		      const char* workspaceName = "combined",
		      const char* modelConfigName = "ModelConfig",
		      const char* dataName = "obsData"){

#ifdef __CINT__
  cout <<"You must use ACLIC for this.  Use ModelInspector.C+"<<endl;
  return;
#endif

  /////////////////////////////////////////////////////////////
  // First part is just to access a user-defined file 
  // or create the standard example file if it doesn't exist
  ////////////////////////////////////////////////////////////
  const char* filename = "";
  if (!strcmp(infile,""))
    filename = "results/example_combined_GaussExample_model.root";
  else
    filename = infile;
  // Check if example input file exists
  TFile *file = TFile::Open(filename);

  // if input file was specified byt not found, quit
  if(!file && strcmp(infile,"")){
    cout <<"file not found" << endl;
    return;
  } 

  // if default file not found, try to create it
  if(!file ){
    // Normally this would be run on the command line
    cout <<"will run standard hist2workspace example"<<endl;
    gROOT->ProcessLine(".! prepareHistFactory .");
    gROOT->ProcessLine(".! hist2workspace config/example.xml");
    cout <<"\n\n---------------------"<<endl;
    cout <<"Done creating example input"<<endl;
    cout <<"---------------------\n\n"<<endl;
  }

  // now try to access the file again
  file = TFile::Open(filename);
  if(!file){
    // if it is still not there, then we can't continue
    cout << "Not able to run hist2workspace to create example input" <<endl;
    return;
  }

  
  /////////////////////////////////////////////////////////////
  // Tutorial starts here
  ////////////////////////////////////////////////////////////

  // get the workspace out of the file
  RooWorkspace* w = (RooWorkspace*) file->Get(workspaceName);
  if(!w){
    cout <<"workspace not found" << endl;
    return;
  }

  // get the modelConfig out of the file
  ModelConfig* mc = (ModelConfig*) w->obj(modelConfigName);

  // get the modelConfig out of the file
  RooAbsData* data = w->data(dataName);

  // make sure ingredients are found
  if(!data || !mc){
    w->Print();
    cout << "data or ModelConfig was not found" <<endl;
    return;
  }

  new ModelInspectorGUI(w,mc,data);
}