#include "QFramework/TQSystematicsHandler.h"
#include "QFramework/TQHistogramUtils.h"
#include "QFramework/TQUtils.h"
#include "QFramework/TQIterator.h"

// #define _DEBUG_

#include "QFramework/TQLibrary.h"

////////////////////////////////////////////////////////////////////////////////////////////////
//
// TQSystematicsHandler:
//
// With the advent of the TQSystematicsHandler, it is now easily
// possible to include not only normalization, but also shape
// systematics into histograms produced by a plotter or add systematic
// uncertainties to the cutflow uncertainties.
//
////////////////////////////////////////////////////////////////////////////////////////////////

ClassImp(TQSystematicsHandler)

TQSystematicsHandler::TQSystematicsHandler(const TString& name) :
TQSystematicsManager(new TQFolder("histograms")),TQFolder(name),
  _config(new TQFolder("systematicsHandler"))
{
  // constructor with name
}

//__________________________________________________________________________________|___________

TQSystematicsHandler::TQSystematicsHandler(const TString& name, TQFolder* cfg) :
TQSystematicsManager(new TQFolder("histograms")),TQFolder(name),  
  _config(cfg)
{
  // constructor with name and config
}

//__________________________________________________________________________________|___________

TQFolder* TQSystematicsHandler::config(){
  return this->_config;
}

//__________________________________________________________________________________|___________

TQFolder* TQSystematicsHandler::addCut(const TString& id){
  // add a cut to the systematics handler
  if(!TQFolder::isValidName(id)){
    ERRORclass("'%s' is not a valid identifier!",id.Data());
    return NULL;
  }
  TQFolder* cut = config()->getFolder("Cuts+")->getFolder(id+"+");
  cut->setTagString("Cut",id);
  return cut;
}

//__________________________________________________________________________________|___________

TQFolder* TQSystematicsHandler::addSystematic(const TString& id, const TString& tags){
  // add a systematic to the systematics handler
  if(!TQFolder::isValidName(id)){
    ERRORclass("'%s' is not a valid identifier!",id.Data());
    return NULL;
  }
  TQFolder* sys = config()->getFolder("Systematics+")->getFolder(id+"+");
  sys->importTags(tags);
  return sys;
}

//__________________________________________________________________________________|___________

TQFolder* TQSystematicsHandler::addSystematic(const TString& id, const TString& upvar, const TString& dnvar){
  // add a systematic to the systematics handler
  if(!TQFolder::isValidName(id)){
    ERRORclass("'%s' is not a valid identifier!",id.Data());
    return NULL;
  }
  TQFolder* sys = config()->getFolder("Systematics+")->getFolder(id+"+");
  sys->setTagString("Up",upvar);
  sys->setTagString("Down",dnvar);  
  return sys;
}

//__________________________________________________________________________________|___________

TQFolder* TQSystematicsHandler::addVariation(const TString& id, const TString& tags){
  // add a variation to the systematics handler
  if(!TQFolder::isValidName(id)){
    ERRORclass("'%s' is not a valid identifier!",id.Data());
    return NULL;
  }
  TQFolder* var = config()->getFolder("Variations+")->getFolder(id+"+");
  var->importTags(tags);
  return var;
}

//__________________________________________________________________________________|___________

TQFolder* TQSystematicsHandler::addSample(const TString& id, const TString& path){
  // add a sample to the systematics handler
  if(!TQFolder::isValidName(id)){
    ERRORclass("'%s' is not a valid identifier!",id.Data());
    return NULL;
  }
  TQFolder* sample = config()->getFolder("Samples+")->getFolder(id+"+");
  sample->setTagString("Path",path);
  return sample;
}

//__________________________________________________________________________________|___________

TQFolder* TQSystematicsHandler::addHistogram(const TString& id, const TString& path, const TString& options){
  TQFolder* hist = config()->getFolder("Histograms+")->getFolder(id+"+");
  hist->setTagString("Path",path);
  TString tmp(path);
  hist->setTagString("Cut",TQFolder::getPathHead(tmp));
  hist->importTags(options);
  return hist;
}

namespace {
  struct HistOptions{
    TString name;
    TString cut;
    TString path;
    TQTaggable* options;
    HistOptions(TString name){
      this->path = name;
      this->name = TQFolder::getPathTail(name);
      this->cut = name;
      this->options = NULL;
    }
    HistOptions(TQFolder* config){
      this->path = config->getTagStringDefault("Path");
      this->name = config->GetName();
      this->cut = config->getTagStringDefault("Cut");
      this->options = config;
    }
  };
}

void TQSystematicsHandler::collectHistograms(TQSampleDataReader* rd, TQFolder* variation){
  // collect the counters from a reader for one subpath of the type systname/file.root/
  TQFolder* cuts = this->config()->getFolder("Cuts");
  if(!variation->hasTag("Variation")) variation->setTagString("Variation",variation->GetName());  
  TQFolderIterator sampleItr(this->config()->getListOfFolders("Samples/?"));
  while(sampleItr.hasNext()){
    TQFolder* sampleConfig = sampleItr.readNext();
    TString sfolderpath = variation->replaceInText(sampleConfig->getTagStringDefault("Path",""));
    if(sfolderpath.Contains("$")){    
      ERROR("sample '%s' has unexpanded alias in path: '%s'. skipping.",sampleConfig->GetName(),sfolderpath.Data());
      continue;
    }
    std::vector<HistOptions> histograms;
    TQIterator histoItr(rd->getListOfHistogramNames(sfolderpath),true);
    if(histoItr.getCollection() && histoItr.getCollection()->GetSize() > 0){
      while(histoItr.hasNext()){
        TString histo(histoItr.readNext()->GetName());
        auto h = HistOptions(histo);
        if(cuts && !cuts->getFolder(h.cut)) continue;
        histograms.push_back(h);
      }
    }
    TQFolderIterator morehistos(this->config()->getListOfFolders("Histograms/?"),true);
    while(morehistos.hasNext()){
      TQFolder* c = morehistos.readNext();
      histograms.emplace_back(HistOptions(c));
    }
    std::map<TString,std::vector<TString> > histoNames;
    for(const auto& hist:histograms){
      TH1* h = rd->getHistogram(sfolderpath,hist.path,hist.options);
      if(!h){
        ERROR("failed to retrieve histogram '%s' from path '%s' with options '%s'",hist.name.Data(),hist.path.Data(),hist.options ? hist.options->exportTagsAsString().Data() : "");
        continue;
      }
      if(histoNames.find(hist.cut) == histoNames.end()) histoNames[hist.cut] = std::vector<TString>();
      histoNames[hist.cut].push_back(hist.name);
      h->SetName(sampleConfig->getName()+"."+hist.name);
      this->storeVarHisto(h,variation->GetName(),hist.cut);
    }
    for(auto histo:histoNames){
      TQFolder* cut = this->getFolder(histo.first+"+");
      cut->setTagList("Shapes",histo.second);
    }
  }
}

//__________________________________________________________________________________|___________

void TQSystematicsHandler::collectCounters(TQSampleDataReader* rd, TQFolder* variation){
  // collect the counters from a reader for one subpath of the type systname/file.root/
  TQFolder* cuts = this->config()->getFolder("Cuts");  
  TQFolderIterator sampleItr(this->config()->getListOfFolders("Samples/?"));
  while(sampleItr.hasNext()){
    TQFolder* sampleConfig = sampleItr.readNext();
    TString sfolderpath = variation->replaceInText(sampleConfig->getTagStringDefault("Path",""));
    if(sfolderpath.Contains("$")){
      ERROR("sample '%s' has unexpanded alias in path: '%s'. skipping.",sampleConfig->GetName(),sfolderpath.Data());
      continue;
    }    
    DEBUGclass("retrieving counters for sample '%s' from path '%s'",sampleConfig->GetName(),sfolderpath.Data());
    TQIterator cntItr(rd->getListOfCounterNames(sfolderpath),true);
    if(!cntItr.getCollection() || cntItr.getCollection()->GetSize() == 0){
      ERRORclass("unable to retrieve list of counters from sample '%s'",sfolderpath.Data());
    } else {
      while(cntItr.hasNext()){
        TString cutname(cntItr.readNext()->GetName());
        if(cuts && !cuts->getFolder(cutname)) continue;        
        TQCounter* cnt = rd->getCounter(sfolderpath,cutname);
        if(cnt->getRawCounter()>0){
          this->getFolder(cutname+"+")->setTagString("Cut",cutname);
          cnt->SetName(sampleConfig->getName()+".yield");      
          TH1* h = TQHistogramUtils::counterToHistogram(cnt);
          this->storeVarHisto(h,variation->GetName(),cutname);
        }
        delete cnt;
      }
    }
  }
}


//__________________________________________________________________________________|___________

void TQSystematicsHandler::collectVariation(TQFolder* var){
  // collect counters and histograms for one systematic
  //@tag: [verbose] This object tag enables verbosity.
  if(this->getTagBoolDefault("verbose",false)){
    INFOclass("collecting variation '%s'",var->GetName());
  }
  TQSampleFolder* sf = this->getSampleFolder(var);
  TQSampleDataReader* rd = new TQSampleDataReader(sf);
  rd->setApplyStyles(false);
  this->collectCounters  (rd,var);
  this->collectHistograms(rd,var);
  delete rd;
}

//__________________________________________________________________________________|___________

void TQSystematicsHandler::collect(){
  // collect counters and histograms for all registered systematics
  DEBUGclass("starting collection of variations");
  TQFolderIterator itr(this->config()->getListOfFolders("Variations/?"),true);
  while(itr.hasNext()){
    this->collectVariation(itr.readNext());
  }
}

//__________________________________________________________________________________|___________

void TQSystematicsHandler::compute(){
  // compute the systematic uncertainties from the variations
  DEBUGclass("starting collection of variations");  
  TQFolderIterator cuts(this->getListOfFolders("?"));
  while(cuts.hasNext()){
    TQFolder* cut = cuts.readNext();
    DEBUGclass("computing for cut '%s'",cut->GetName());      
    cut->importTags(this->config()->getFolder("Cuts")->getFolder(cut->GetName()));
    TQFolderIterator samples(this->config()->getListOfFolders("Samples/?"),true);
    while(samples.hasNext()){
      TQFolder* sample = cut->getFolder(samples.readNext()->getName()+"+");
      sample->setTagString("Sample",sample->GetName());
      DEBUGclass("    for sample '%s'",sample->GetName());
      TQFolderIterator systematics(this->config()->getListOfFolders("Systematics/?"),true);
      while(systematics.hasNext()){
        TQFolder* sysConfig = systematics.readNext();    
        DEBUGclass("        for systematic '%s'",sysConfig->GetName());              
        TQFolder* sys = sample->getFolder(sysConfig->getName()+"+");
        sys->setTagString("Systematic",sysConfig->getName());
        sys->importTags(sysConfig);
        this->includeOverallSys(sysConfig,sys,cut->GetName(),sample->getName()+".yield");
        for(const auto& shape:cut->getTagVString("Shapes")){
          DEBUGclass("            shape '%s'",shape.Data());                        
          this->includeHistoSys(sysConfig,sys->getFolder(shape+"+"),cut->GetName(),sample->getName()+"."+shape);              
        }
      }
    }
  }
}

//__________________________________________________________________________________|___________

void TQSystematicsHandler::exportObjects(TQFolder* section,TQFolder* target, bool includeNormalizationInShape, bool symmetrizeOneSided){
  // export the created objects for one cut
  if(!section || !target) return;

  // obtain the normalization
  TQFolderIterator systItr(section->getListOfFolders("?"),true);  
  double totalyield = 0.;
  std::vector<double> variations;
  while(systItr.hasNext()){
    TQFolder* syst = systItr.readNext();
    if(!syst) continue;
    double up = 0.;
    double dn = 0.;
    // division by 2 here because of double counting of up/down    
    if(syst->getTagDouble("High",up) && syst->getTagDouble("Low",dn)){
      if(TQUtils::isNum(up) && TQUtils::isNum(dn)){
        totalyield += pow(0.5*(up-dn),2);
      } else if(TQUtils::isNum(up)){
        totalyield += pow(up,2);
      } else if(TQUtils::isNum(dn)){
        totalyield += pow(dn,2);
      }
    }
  }
  double normsys = sqrt(fabs(totalyield));
  target->setTagDouble("yield",normsys);

  // obtain the shape(s)
  target->setTagBool("shapeIncludesNormalization",includeNormalizationInShape);
  std::vector<TString> shapes = section->getTagVString("~Shapes");
  for(const auto& shape:shapes){
    DEBUGclass("exporting histograms '%s' for '%s' to '%s'",shape.Data(),section->getPath().Data(),target->getPath().Data());
    TQFolderIterator systItr(section->getListOfFolders("?/"+shape),true);
    std::vector<TH1*> histos_up;
    std::vector<TH1*> histos_dn;
    int n=0;
    while(systItr.hasNext()){
      TQFolder* syst = systItr.readNext();
      if(!syst) continue;
      TH1* h_up = this->getHisto(syst->getTagStringDefault("HistoRelHigh"));
      TH1* h_dn = this->getHisto(syst->getTagStringDefault("HistoRelLow"));
      if(!h_up || !h_dn){
        ERRORclass("unable to retrieve histograms for '%s'",syst->getTagStringDefault("~Systematic").Data());
        continue;
      }
      histos_up.push_back(h_up);
      histos_dn.push_back(h_dn);
      n++;
    }
    TH1* hist = NULL;
    double up;
    double dn;
    double delta;
    for(size_t i=0; i<histos_up.size(); ++i){
      if(!hist){
        hist = TQHistogramUtils::copyHistogram(histos_up[i],shape);
        hist->Reset();
      }
      for(int b=0; b<hist->GetNbinsX()+1; ++b){
        up = histos_up[i]->GetBinContent(b);
        dn = histos_dn[i]->GetBinContent(b);
        if (symmetrizeOneSided && up < 1. && dn < 1.) {
          delta = (up < dn) ? (2. * (1. - up)) : (2. * (1. - dn));
        }
        else if (symmetrizeOneSided && up > 1. && dn > 1.) {
          delta = (up > dn) ? (2. * (up - 1.)) : (2. * (dn - 1.));
        }
        else {
          delta = up - dn;
        }
        hist->AddBinContent(b,pow(0.5*delta,2));
      }
    }
    if(hist){
      double normComp = includeNormalizationInShape ? normsys : 0;
      for(int i=0; i<hist->GetNbinsX()+1; ++i){
        hist->SetBinContent(i,sqrt(hist->GetBinContent(i) + pow(normComp,2)));
        hist->SetBinError(i,0);
      }
      hist->SetEntries(n);
      hist->SetMinimum(0);
      if(!target->addObject(hist)){
        WARNclass("unable to export histogram '%s'",hist->GetName());
      }
    }
  }
}

//__________________________________________________________________________________|___________

TQFolder* TQSystematicsHandler::exportSystematics(const TString& sample, bool includeNormalizationInShape, bool symmetrizeOneSided){
  // create and export variation histograms and counters for all systematics
  TQFolder* systematics = new TQFolder(this->GetName());
  TQFolderIterator cuts(this->getListOfFolders(TQFolder::concatPaths("?",sample)),true);
  while(cuts.hasNext()){
    TQFolder* cut = cuts.readNext();
    if(!cut) continue;
    DEBUGclass("processing section '%s'",cut->getPath().Data());
    TQFolder* target = systematics->getFolder(cut->getTagStringDefault("~Cut")+"+");
    TQSystematicsHandler::exportObjects(cut, target, includeNormalizationInShape, symmetrizeOneSided);
  }
  return systematics;
}

//__________________________________________________________________________________|___________

void TQSystematicsHandler::printSystematics(){
  // create and export variation histograms and counters for all systematics
  TQFolderIterator folders(this->getListOfFolders("?/?"),true);
  while(folders.hasNext()){
    TQFolder* folder = folders.readNext();
    if(!folder) continue;
    std::cout << folder->GetName() << std::endl;
  }
}

//__________________________________________________________________________________|___________

bool TQSystematicsHandler::writeFolderHook(TDirectory* dir, const TString&, int, bool){
  // redirect internal writing to another folder
  TQFolder* tmpfolder = this->copy();
  bool ok = tmpfolder->writeFolder(dir,-1,false);
  delete tmpfolder;
  return ok;
}

//__________________________________________________________________________________|___________

std::vector<TQFolder*> TQSystematicsHandler::getRanking(const TString& cutname, const TString& sample){
  // return the systematics ranking for one specific cut
  TQFolder* section = this->getFolder(TQFolder::concatPaths(cutname,sample));
  std::vector<TQFolder*> retval;
  if(!section){
    ERRORclass("unable to retrieve ranking for cut '%s' and sample '%s'",cutname.Data(),sample.Data());
    return retval;
  }
  TQFolderIterator systematics(section->getListOfFolders("?"),true);
  while(systematics.hasNext()){
    TQFolder* f = systematics.readNext();
    retval.push_back(f);
  }
  std::sort(retval.begin(),retval.end(),[](TQFolder* a, TQFolder* b) {
    return a->getTagDoubleDefault("Percent",0) > b->getTagDoubleDefault("Percent",0);
  });
  return retval;
}

//__________________________________________________________________________________|___________

TQTable* TQSystematicsHandler::getTable(const TString& sample, const TString& cutname){
  // write the systematics table
  std::vector<TQFolder*> ranking = this->getRanking(sample,cutname);
  if(ranking.size() == 0) return NULL;
  TQTable* table = new TQTable(sample+"-"+cutname);
  table->expand(ranking.size()+1,2);
  table->setEntry(0,0,"Systematic");
  table->setEntry(0,1,"\\% yield","latex");
  table->setEntry(0,1,"% yield","plain");
  table->setEntry(0,1,"% yield","html");
  int idx=0;
  for(auto* systematic:ranking){
    double yield = systematic->getTagDoubleDefault("Percent",0);
    size_t row = idx+1;
    table->setEntry(row,0,systematic->GetName(),"ascii");
    //@tag: [title.latex,title.html] This folder tag is used to set the LaTeX/HTML representations of the systematic's title.
    table->setEntry(row,0,systematic->getTagStringDefault("title.latex",systematic->GetName()),"latex");
    table->setEntry(row,0,systematic->getTagStringDefault("title.html",systematic->GetName()),"html");
    table->setEntryValue(row,1,yield);
    ++idx;
  }
  return table;
}

//__________________________________________________________________________________|___________

TQSystematicsHandler::~TQSystematicsHandler(){
  delete _config;
}

//__________________________________________________________________________________|___________

TQSampleFolder* TQSystematicsHandler::getSampleFolder(const TString& name){
  return this->getSampleFolder(this->config()->getFolder("Variations")->getFolder(name));
}