//STL
#include <iostream>
#include <map>
#include <math.h>
#include <string>
//ROOT
#include "TCanvas.h"
#include "TFile.h"
#include "TH1.h"
#include "TString.h"
#include "TStopwatch.h"
#include "TTree.h"
//Roofit
#include "RooAddPdf.h"
#include "RooArgSet.h"
#include "RooCategory.h"
#include "RooChebychev.h"
#include "RooChi2Var.h"
#include "RooDataSet.h"
#include "RooFitResult.h"
#include "RooFormulaVar.h"
#include "RooGaussian.h"
#include "RooMinimizer.h"
#include "RooMinuit.h"
#include "RooNLLVar.h"
#include "RooPlot.h"
#include "RooRealVar.h"
#include "RooSimultaneous.h"
#include "RooUnblindPrecision.h"
#include "RooWorkspace.h"

int nthreads(4);
double mDp_pdg(1.86962);//, mDs_pdg(1.96849);

void makeModel(RooWorkspace &w);
void fit( RooAbsPdf *model, RooDataSet *data, bool kNLL = true, bool kChi = false, bool kMinos = false);
void plot( RooRealVar *x, RooAbsPdf *model, RooDataSet *data );

int testBlind()
{

  TFile *WS = new TFile("res/testBlind_WS.root");
  RooWorkspace *w = new RooWorkspace("testBlind_WS");
  if (WS->IsZombie()) {
    makeModel(*w);
    w->writeToFile("res/testBlind_WS.root");
  } else
    w = (RooWorkspace*)WS->Get("testBlind_WS");

  // generate Dataset
  int nev = w->var("nD")->getVal() + w->var("nbkg")->getVal();
  RooDataSet *data = w->pdf("model")->generate(RooArgSet(*w->var("mD")),nev);

  bool blind = false;
  w->cat("blindCat")->setIndex(int(blind));

  //  fit( w->pdf("model"), data );
  //  plot( w->var("mD"), w->pdf("model"), data );


  // Print Results
  std::cout << "\n\n mean = " << w->var("m")->getVal() << " +- " << w->var("m")->getError()
	    << "\n sigma = " << w->var("s")->getVal() << " +- " << w->var("s")->getError()
	    << "\n\n nD = " << w->var("nD")->getVal() << " +- " << w->var("nD")->getError()
	    << "\n nbkg = " << w->var("nbkg")->getVal() << " +- " << w->var("nbkg")->getError()
	    << std::endl << std::endl;

  return 0;
}

void makeModel(RooWorkspace &w)
{

  RooRealVar mD("mD","m(D)",1.82,1.92);
  // Sig Gau
  RooRealVar m("m","mean D+",mDp_pdg,mDp_pdg-0.02,mDp_pdg+0.02,"GeV/c^{2}");
  RooRealVar s("s","sigma D+",0.004,0.001,0.01,"GeV/c^{2}");
  RooGaussian g("g","gaussian D+",mD,m,s);
  // poly Bkg
  RooRealVar p0("p0","poly0",.29,-1.,1.); //p0.setConstant(1);
  RooRealVar p1("p1","poly1",.029,-1.,1.); //p1.setConstant(1); 
  RooChebychev poly("poly","polynomial",mD,RooArgList(p0,p1));
  // yields
  RooRealVar nD("nD","D events",0.29*1e4,0.,1e+8);
  nD.setError(1);
  RooRealVar nbkg("nbkg","bkg events D",(1-0.29)*1e4,0.,1e+8);
  nbkg.setError(1);
  // Blinding
  RooCategory blindCat("blindCat","blind state Category");
  blindCat.defineType("unblind",0);
  blindCat.defineType("blind",1);
  std::string blindString("blinded");
  //  RooUnblindPrecision nD_Blind("nD_Blind","nD blind",(std::string(nD.GetName())+blindString).c_str(),100.,10.,nD,blindCat,0);
  RooUnblindPrecision nD_Blind("nD_Blind","nD blind","nDblinded",100.,10.,nD,blindCat,1);
  
  // sum
  RooAddPdf *model = 
    new RooAddPdf("model","model",RooArgSet(g,poly),RooArgList(nD_Blind,nbkg));
  
  w.import(*model);
  w.Print();
}

void fit( RooAbsPdf *model, RooDataSet *data, bool kNLL , bool kChi , bool kMinos)
{
  // Timer
  TStopwatch timer;
  timer.Start();

  double chi2val(0);
  // Fit
  RooMinimizer *minuit(0);
  int simplexStatus(0);
  int migradStatus(0);
  int hesseStatus(0);
  int minosStatus(0);

  if(kChi) { //Chi2
    RooDataHist *dataBinned = data->binnedClone();
    RooChi2Var chi2("chi2","chi2",*model,*dataBinned,RooFit::DataError(RooDataHist::Poisson),
		    RooFit::Extended(1),RooFit::NumCPU(nthreads));
    
    minuit = new RooMinimizer(chi2);    
    minuit->optimizeConst(true);
    minuit->setProfile(true);
    minuit->setStrategy(1);
    minuit->setPrintLevel(1);

    //      simplexStatus = minuit->simplex();// raw search for minimum
    migradStatus = minuit->migrad();//minimization
    hesseStatus = minuit->hesse();//covariance matrix
    chi2val = chi2.getVal();
      
  } 
  if (kNLL) {  // Likelihood
  RooNLLVar nll("nll","nll",*model,*data,RooFit::Extended(1),RooFit::NumCPU(nthreads));
    
  minuit = new RooMinimizer(nll);    
    
  minuit->optimizeConst(true);
  minuit->setProfile(true);
  minuit->setStrategy(1); 
  minuit->setPrintLevel(1);
  
    //    simplexStatus = minuit->simplex();// raw search for minimum
  migradStatus = minuit->migrad(); // minimization
  hesseStatus = minuit->hesse(); // covariance matrix
  /*
    conf.fixPars();
    migradStatus = minuit->migrad(); // minimization
    hesseStatus = minuit->hesse(); // covariance matrix
  */
  if(kMinos) minosStatus = minuit->minos();//RooArgSet(conf.GetYields(),conf.GetAsyms()));
  }

  timer.Stop(); 

  if (kChi || kNLL) {

    // Fit Result 
    RooFitResult *res = minuit->save();
    cout << endl;
    cout << "Fit Result"<<endl;
    res->Print("V");
    cout << "Migrad status = " << migradStatus << endl;
    cout << "Hesse status = " << hesseStatus << endl;
    cout << "Minos status = " << minosStatus << endl;
    cout << endl;
    if(kChi) cout << "Chi2 = " << chi2val << endl << endl;

    cout << "Real Time  = " << timer.RealTime() << endl;
    cout << "CPU Time  = " << timer.CpuTime() << endl;
    cout << endl;
  }
}


void plot( RooRealVar *x, RooAbsPdf *model, RooDataSet *data ) 
{
  TCanvas *c = new TCanvas("c_test", "c_test", 600, 600);
  RooPlot *frame= x->frame();
  data->plotOn(frame);
  model->plotOn(frame);
  model->plotOn(frame, RooFit::Components("poly"), RooFit::LineStyle(kDashed));
  frame->Draw();
}