// 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 <ROOT/RDataFrame.hxx>
#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>
// BOOST
#include <boost/property_tree/ptree.hpp>
#include <boost/property_tree/info_parser.hpp>
#include <boost/algorithm/string.hpp>
#include <boost/optional.hpp>

template <typename STR = std::string> // template because we want to use r- and l-value std::strings
std::vector<std::string> parse_filenames(STR &&input_file_names, const std::string &workdir = ".")
{
   std::vector<std::string> input_files;
   // if the input string starts with '@', we have to parse a whole input file
   if (boost::algorithm::starts_with(input_file_names, "@")) {
      std::string line;
      std::ifstream myfile;
      myfile.open(input_file_names.substr(1).data()); // substr returns the string without the '@'
      if (!myfile)
         myfile.open((workdir + "/" + input_file_names.substr(1)).data());
      while (std::getline(myfile, line))
         input_files.push_back(line); // TODO: could be that there is a trailing character that needs to be removed
      myfile.close();
   } else // the input string doesn't start with '@', boost does the work for us :)
      boost::algorithm::split(input_files, input_file_names, boost::algorithm::is_any_of(";"));

   if (input_files.empty())
      throw std::runtime_error("Is \"input_file_names\" empty?");
   return input_files;
}

template <typename STR1 = std::string,
          typename STR2 = std::string> // template because we want to use r- and l-value std::strings
std::unique_ptr<TChain>
get_chain(STR1 &&treename, STR2 &&input_file_names, const std::string &workdir = ".",
          const std::string &chain_title = "Chain")
{
   auto input_files = parse_filenames(input_file_names, workdir);
   std::unique_ptr<TChain> ch{new TChain(treename.data(), chain_title.data())};
   for (auto &infn : input_files)
      ch->Add(infn.data());
   if (!ch->GetFile()) {
      std::cout << "Cannot find input files. Trying to prepend workdir." << std::endl;
      for (auto &infn : input_files)
         ch->Add((workdir + "/" + infn).data());
      if (!ch->GetFile())
         throw std::runtime_error("TChain is empty. Check inputs!");
   }
   return std::move(ch);
}

int tree_trimmerTDF(const std::string &ifn, const std::string &itn, const std::string &ofn, const std::string &cfn,
                    const std::string &wd = ".")
{

   TStopwatch clock;
   clock.Start();

   // get input as TChain
   auto chain = get_chain(itn, ifn, wd);
   chain->SetCacheSize();

   // get config-file
   boost::property_tree::ptree configtree;
   boost::property_tree::read_info(cfn, configtree, configtree);

   auto trim_vars = configtree.get_child_optional("variables");

   auto tmp_msg = "Trimming input tree(s) with cut " + configtree.get<std::string>("basiccuts", "");
   if (trim_vars)
      printf("%s and writing out %lu variables\n", tmp_msg.data(), configtree.get_child("variables").size());
   else
      printf("%s and writing out full tree\n", tmp_msg.data());

   std::vector<std::string> branches_in = {};
   std::vector<std::string> default_branches_out = {};
   std::vector<std::pair<std::string, std::string>> transformations;
   if (trim_vars)
      for (const auto &var : *trim_vars) {
         branches_in.push_back(var.first);
         if (var.first.find("[0]") != std::string::npos) {
            default_branches_out.push_back(boost::algorithm::replace_all_copy(var.first, "[0]", "Z"));
            transformations.emplace_back(default_branches_out.back(), var.first);
         } else if (auto nn = var.second.get_optional<std::string>("new_name")) {
            default_branches_out.push_back(*nn);
            if (auto tfm = var.second.get_optional<std::string>("transformation"))
               transformations.emplace_back(default_branches_out.back(), *tfm);
            else
               transformations.emplace_back(default_branches_out.back(), var.first);
         } else if (!var.second.get("to_output", 1))
            printf("Using %s only as input branch\n", var.first.data());
         else
            default_branches_out.push_back(var.first);
      }

   // first, we need to know how many events there are and how many threads we use:
   ROOT::EnableImplicitMT();
   ROOT::RDataFrame d(*chain.get(), branches_in);
   std::vector<ROOT::RDF::RInterface<ROOT::Detail::RDF::RLoopManager>> df{d};
   if (!transformations.empty())
      for (const auto &tf : transformations) {
         printf("transforming %s to a new branch: %s\n", tf.second.data(), tf.first.data());
         df.push_back(df.back().Define(tf.first, tf.second));
      }
   df.back().Snapshot(configtree.get("outtreename", itn), wd + "/" + ofn, default_branches_out); //,tsno

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

   return 0;
}

int main()
{
   //"/afs/cern.ch/work/m/mstahl/public/ForEnrico/XcPi_2011MagUp.root"
   // "/home/eguiraud/Scratchpad/marian_network_reading/XcPi_2011MagUp.root"
   tree_trimmerTDF("/afs/cern.ch/work/m/mstahl/public/ForEnrico/XcPi_2011MagUp.root", "Xibm2XiczPiDTTuple/DecayTree",
                   "test.root", "XiczPi_TDF.txt");
   return 0;
}