//C++
#include <type_traits>
#include <vector>
#include <map>
#include <exception>
#include <string>
#include <sstream>
#include <functional>
#include <chrono>
#include <thread>
//root
#include <TROOT.h>
#include <TObject.h>
#include <TFile.h>
#include <TString.h>
#include <TTree.h>
#include <TLeaf.h>
#include <TBranch.h>
#include <TFriendElement.h>
#include <TObjArray.h>
#include <TObjString.h>
#include <TEntryList.h>
#include <TEnv.h>
#include <TStopwatch.h>
//new RDF-stuff
#include <ROOT/RDataFrame.hxx>
#include <ROOT/RDF/RInterface.hxx>
//local
#include "ProgressBar.h"
#include "IOjuggler.h"
#include "misID_betaRDF.h"

#ifndef NTUPLE_GIZMO_GIT_HASH
#define NTUPLE_GIZMO_GIT_HASH " "
#endif

int main(int argc, char** argv){

  TStopwatch clock;
  clock.Start();

  //parse command line options
  const auto options = IOjuggler::parse_options(argc, argv, "c:d:i:o:t:v:h","nonoptions: any number of <file:friendtree> combinations");
  const auto wd = options.get<std::string>("workdir");
  const auto ofn = options.get<std::string>("outfilename");
  const auto ifn = options.get<std::string>("infilename");
  const auto itn = options.get<std::string>("treename");

  //init MessageService for easy and nice printout
  MessageService msgsvc("tree_trimmerRDF",static_cast<MSG_LVL>(options.get<int>("verbosity")));
  msgsvc.debugmsg("Current ntuple-gizmo git hash: " + static_cast<std::string>(NTUPLE_GIZMO_GIT_HASH));

  //get config-file
  auto configtree = IOjuggler::get_ptree(options.get<std::string>("config"));
  //manipulations of the config file
  IOjuggler::auto_replace_in_ptree(configtree);
  IOjuggler::auto_append_in_ptree(configtree);

  //get variables we want to address, store and/or transform from config file
  std::vector<std::string> branches_in = {};
  std::vector<std::string> default_branches_out = {};
  std::vector<std::pair<std::string,std::string>> trafos;
  if(const auto trim_vars = configtree.get_child_optional("variables")){
    msgsvc.infomsg("Running with " + std::to_string((*trim_vars).size()) + " input variables");
    for(const auto& var : *trim_vars){
      if(var.first != "dummy" || var.first != "dummy_i")
        branches_in.push_back(var.first);
      if(auto nn = var.second.get_optional<std::string>("nn")){
        // decide whether to add the new brach to the output tree
        if(!var.second.get("to_output",1) || var.first == "dummy_i")
          msgsvc.debugmsg("Using " + *nn + " only as input branch");
        else
          default_branches_out.push_back(*nn);
        // register a transformation
        if(auto tfm = var.second.get_optional<std::string>("tf"))
          trafos.emplace_back(*nn,*tfm);
        else
          trafos.emplace_back(*nn,var.first);
      }
      else if(var.first.find("[0]") != std::string::npos){
        default_branches_out.push_back(boost::algorithm::replace_all_copy(var.first, "[0]", "Z"));
        trafos.emplace_back(default_branches_out.back(),var.first);
      }
      else
        default_branches_out.push_back(var.first);
    }
  }
  else msgsvc.infomsg("Running with full set of input variables");

  //now we can get the input file and set up the dataframe
  //first, we need to know how many events there are and how many threads we use:
  ROOT::EnableImplicitMT();//need to call this first to be able to call ROOT::GetImplicitMTPoolSize()
  const auto nSlots = std::min(configtree.get("threads",ROOT::GetImplicitMTPoolSize()),ROOT::GetImplicitMTPoolSize());
  msgsvc.debugmsg("Running with " + std::to_string(nSlots) + " parallel threads");
  ROOT::EnableImplicitMT(nSlots);

  //get input as TChain
  auto chain = IOjuggler::get_chain(itn,ifn,wd,msgsvc);
  //attach friend-files, if any
  msgsvc.debugmsg("Trying to invite some friends");
  try{
    [[maybe_unused]] auto fes = IOjuggler::get_friends(options,chain,msgsvc);
    if(fes[0] != nullptr){
      msgsvc.errormsg("Adding friend trees is currently not supported. Returning...");
      return 0;
    }
  }
  catch (std::exception& e){
    msgsvc.warningmsg("Caught exception " + static_cast<std::string>(e.what()));
    throw;
  }

  ROOT::RDataFrame d(*chain.get(),branches_in);
  //use ROOT::RDF::RNode to chain commands on the dataframe
  ROOT::RDF::RNode df = d;
  if(!trafos.empty())
    for(const auto& tf : trafos){
      msgsvc.debugmsg("transforming " + tf.second + " to a new branch: " + tf.first);
      df = df.Define(tf.first,tf.second);
    }
  if(const auto misID = configtree.get_child_optional("misIDMbeta")){
    std::vector<std::string> daughters;
    for(const auto& daughter : *misID)
      daughters.push_back(daughter.first);
    MbetaRDF(daughters,df,default_branches_out,msgsvc);
    msgsvc.debugmsg("Added masses and betas for misID studies");
  }

  //apply cuts
  if(const auto cut = configtree.get_optional<std::string>("Filter")){
    msgsvc.infomsg("Applying Filter action: " + *cut);
    df = df.Filter(*cut);
  }

  //lazy action that enables to draw the progress bar (the action doesn't take additional time)
  auto lac = df.Count();

  //keep us entertained with a progress bar
  ProgressBar pg;
  pg.start(chain->GetEntries(),msgsvc,nSlots);
  std::mutex barMutex;
  unsigned long long evt = 0ull;
  lac.OnPartialResultSlot(1ull, [&pg,&evt,&barMutex](unsigned int, auto& ) {
    std::lock_guard<std::mutex> l(barMutex); // lock_guard locks the mutex at construction, releases it at destruction
    pg.update(evt);
    evt++;
  });

  //options for the output file the following options are currently (Dec 18) being updated by the root developers.
  //Check some time later here: https://root.cern.ch/doc/master/Compression_8h_source.html
  ROOT::RDF::RSnapshotOptions rso(configtree.get("outfopt","RECREATE"),
                                  static_cast<ROOT::ECompressionAlgorithm>(configtree.get<unsigned int>("outcompalg",ROOT::RCompressionSetting::EAlgorithm::kInherit)),
                                  configtree.get<unsigned int>("outcomplvl",ROOT::RCompressionSetting::ELevel::kInherit),0,99,false);

  IOjuggler::dir_exists(wd + "/" + ofn,msgsvc);
  df.Snapshot(configtree.get("outtreename",itn),wd + "/" + ofn,default_branches_out,rso);

  const int time = pg.stop();
  msgsvc.infomsg("Elapsed time for event loop: "+std::to_string(time)+" s");

  clock.Stop();
  clock.Print();

  return 0;
}