#ifndef MISID_BETARDF_H
#define MISID_BETARDF_H

#include <vector>
#include <type_traits>
#include <functional>
#include <algorithm>

#include <TLorentzVector.h>
#include <ROOT/RDataFrame.hxx>

typedef std::vector<unsigned int> uivec;
template<typename S> using EnableIfRDFDefinable = typename std::enable_if<std::is_convertible<S,ROOT::RDF::RInterface<ROOT::Detail::RDF::RLoopManager>>::value, void>::type*;

static std::vector<TLorentzVector> init_lvs;

template<typename T, typename ...args>
std::vector<TLorentzVector> lv_lambda(T& px, T& py, T& pz, T& pe, args&&... params) {
  init_lvs.emplace_back(px,py,pz,pe);
  if constexpr (sizeof...(params) > 3)
    lv_lambda(params...);
  return init_lvs;
}

template<typename T>
ROOT::RDF::RInterface<ROOT::Detail::RDF::RLoopManager>
MbetaRDF(std::vector<std::string>& daughters, T&& data_frame, EnableIfRDFDefinable<T> = nullptr) {

  unsigned int n_daughters = daughters.size();

  //get the daughter columns to define the TLorentzVectors
  std::vector<std::string> columns;
  for(const auto& daughter : daughters){
    columns.push_back(daughter+"_PX");
    columns.push_back(daughter+"_PY");
    columns.push_back(daughter+"_PZ");
    columns.push_back(daughter+"_PE");
  }

  //this is a hack... we can't overwrite or append to RInterface<RLoopManager>
  // for this reason, we fill a vector and use RInterface<RLoopManager>::Define on the last element
  std::vector<ROOT::RDF::RInterface<ROOT::Detail::RDF::RLoopManager>> df{data_frame.Define("misID_vlv",lv_lambda<double>,columns)};

  //construct all possible n_body combinations given n_daughters
  //the idea to construct the combinations is that the combination indices are never the same number
  //and rise with the index number. to achieve this, the last index gets incremented until
  //n_daughters is reached, then the second to last is incremented, while building all possible
  //combinations with the last index etc. this makes the function recursive.
  //EXAMPLE (3 body combinations for 5 final state daughters)
  // 123, 124, 125, 134, 135, 145, 234, 235, 245, 345
  std::function<void(unsigned int&, uivec&, unsigned int)> combinatorics = [&n_daughters,&df,&combinatorics] (unsigned int& n_body, uivec& combination_indices, unsigned int index) -> void {

    //add functor for betas
    //very important to make a copy of combination_indices instead of capturing by reference
    auto beta_functor = [combination_indices](std::vector<TLorentzVector>& lvs) {
      double p0 = lvs[combination_indices[0]-1].P(), beta_stub = 0;
      for(decltype(combination_indices.size()) i = 1; i < combination_indices.size(); i++)
        beta_stub += lvs[combination_indices[i]-1].P();
      return (beta_stub-p0)/(beta_stub+p0);
    };

    //add functor for invariant masses
    auto mass_functor = [combination_indices](std::vector<TLorentzVector>& lvs) {
      TLorentzVector temp;
      for(const auto& idx : combination_indices)
        temp += lvs[idx-1];
      return temp.M();
    };

    //function to call Define
    auto call_define = [&df,&combination_indices,&beta_functor,&mass_functor] (std::string&& name) -> void {
      for(const auto& idx : combination_indices)
        name += std::to_string(idx);
      if(name[0] == 'M')
        df.push_back(df.back().Define(name.data(),mass_functor,{"misID_vlv"}));
      else
        df.push_back(df.back().Define(name.data(),beta_functor,{"misID_vlv"}));
    };

    //simple function to get all relevant permutations for beta of (n>3)-body combinations
    auto permutations = [&df,&combination_indices,&call_define] () -> void {
      for(decltype(combination_indices.size()) i = 0; i < combination_indices.size(); i++){
        call_define("beta_");
        //move first element to the end (just convention to write the rotation afterwards.)
        //http://en.cppreference.com/w/cpp/algorithm/rotate
        std::rotate(combination_indices.begin(),combination_indices.begin()+1,combination_indices.end());
      }
    };

    //break condition ultimately fulfilled when function tries
    //to increment the first index to an invalid value
    while(combination_indices[index] <= n_daughters - n_body + index + 1){

      if(index > 0 && combination_indices[index] <= combination_indices[index-1])
        combination_indices[index] = combination_indices[index-1]+1;

      if(index < n_body-1)
        combinatorics(n_body,combination_indices,index+1u);
      else if(index == n_body-1){
        call_define("M_");
        if(n_body > 2)
          permutations();
        else
          call_define("beta_");
      }
      else throw std::runtime_error("How the hell did i end up here?");

      //count up here. reset to lowest possible values before combination index would exceed maximum
      //at_max is a measure for indices being at their maximal value. if this condition is not met,
      //but the current index is at max, this index and all following are set back to their minimally
      //possible value. this value is dictated by the previous index
      bool at_max = true;
      for(decltype(combination_indices.size()) i = 0; i < combination_indices.size(); i++)
        if(combination_indices[i] < n_daughters - n_body + i + 1)
          at_max = false;
      if(combination_indices[index] == n_daughters - n_body + index + 1 && !at_max && index > 0){
        for(decltype(combination_indices.size()) i = index; i < combination_indices.size(); i++)
          combination_indices[i] = combination_indices[i-1] + 1;
        break;
      }
      else
        combination_indices[index] += 1;
    }
  };

  //iterate all possible n_body combinations
  for(auto n_body = n_daughters; n_body > 1; n_body--){//go from full n-body down to two-body combinations
    uivec combination_indices(n_body,0);//init index vector filled with 0s
    for(decltype(combination_indices.size()) i = 0; i < combination_indices.size(); i++)
      combination_indices[i] = i+1;//fill index vector with ascending values
    combinatorics(n_body,combination_indices,0u);
  }
  return df.back();
}
#endif