#include <iostream>
#include <string>
#include <unordered_map>
#include <vector>

#include <ROOT/RDataFrame.hxx>
#include <ROOT/RDFHistoModels.hxx>
#include <ROOT/RDFInterface.hxx>
#include <ROOT/RResultPtr.hxx>
#include <TH2D.h>

#include <libconfig.hh>

#include <toolbox_RnS/args.h>

struct HistogramDefinition {
   std::string variable;
   int nbins;
   double xlow, xup;
};

/**
 * expects a list of lists in the form of a Setting object that contains the histogram definitions.
 * e.g. ["JetPt", 25, 0, 1000] will create a histogram called JetPt, where later the variable JetPt
 * shall be filled into 25 bins between 0 and 1000.
 */
std::vector<HistogramDefinition> ReadHistogramDefinitions(const libconfig::Setting& settings)
{
   if (settings.getType() != libconfig::Setting::Type::TypeList) {
      std::cerr << settings.getPath() << " constains the histogram definitons and should be a list of lists" << std::endl;
      throw libconfig::SettingTypeException(settings);
   }
   std::vector<HistogramDefinition> ret;
   for (const auto& setting : settings) {
      if (setting.getType() != libconfig::Setting::Type::TypeList) {
	 continue;
      }
      HistogramDefinition hd;
      hd.nbins = setting[1];
      hd.xlow = setting[2];
      hd.xup = setting[3];
      hd.variable = setting[0].c_str();
      ret.push_back(hd);
   }

   return ret;
}

/**
 ***********************************************************************
 * NOT TESTED YET!!!
 * This function creates the filters in a sort of recursive way by reading them from
 * a libconfig::Settings object.
 * @param parent_frame is the data frame onto which the filters here should be applied
 * @param parent_name is the name associated with parent_frame
 * @param frames is the final map in which all data frames are stored.
 *    passed by reference to avoid copying multiple times
 *    this is the result to use from this method
 */
void CreateFilters(const libconfig::Setting& settings,
		   ROOT::RDF::RInterface<ROOT::Detail::RDF::RJittedFilter>& parent_frame,
		   std::string parent_name,
		   std::unordered_map<std::string, ROOT::RDF::RInterface<ROOT::Detail::RDF::RJittedFilter, void>>& frames)
{
   std::string lastsettingname = parent_name;
   for (const libconfig::Setting& setting : settings) {
      if (setting.getType() == libconfig::Setting::Type::TypeString) {
	 lastsettingname =  parent_name + ":" + setting.getName();
	 std::string cutstring = setting.c_str();
	 auto df = parent_frame.Filter(cutstring, lastsettingname);
	 frames.insert({lastsettingname, df});
      } else if (setting.getType() == libconfig::Setting::Type::TypeGroup) {
	 CreateFilters(setting, frames.at(lastsettingname), lastsettingname, frames);
      }
   }
}

/**
 * Creates all histograms for a given dataframe and returns the vector of them
 * These are 2D histograms, the y axis showing Ntries
 * @param frame Dataframe for the histograms
 * @param frame_name key of the dataframe in the map, for unique naming of histograms
 * @param histograms vector of the histogram definitions from earlier
 * @param ntriesrange the upper and lower bound of ntries. After applying many filters, we cannot guarantee that
 *    all Ntries values are in the DataFrame, thus Min/Max cannot be used
 */
std::vector<ROOT::RDF::RResultPtr<TH2D>> CreateHistograms(ROOT::RDF::RInterface<ROOT::Detail::RDF::RJittedFilter>& frame,
					       std::string frame_name,
					       const std::vector<HistogramDefinition>& histograms,
					       const std::pair<int, int> ntriesrange = {-2, 20}
					      )
{
   std::vector<ROOT::RDF::RResultPtr<TH2D>> ret;
   for (const auto& hd : histograms) {
      ROOT::RDF::TH2DModel histmodel((frame_name + "_" + hd.variable).c_str(), (frame_name + "_" + hd.variable).c_str(),
				     hd.nbins, hd.xlow, hd.xup,
				     ntriesrange.second - ntriesrange.first, ntriesrange.first, ntriesrange.second
				    );
      ret.push_back(
	 frame.Histo2D(histmodel,
	    hd.variable,
	    "Ntries",
	    "Weight"
	    )
      );
   }
   return ret;
}

void CalcPrediction(std::vector<ROOT::RDF::RResultPtr<TH2D>> histos,
		    std::string outpath,
		    bool store_truth = false,
		    bool store_rebalanced = false
		   )
{
   TFile outputfile(outpath.c_str(), "RECREATE");

   for (auto& rawprediction : histos) {
      rawprediction->Write();
   }
}
   

int main(int argc, char* argv[]) {
   args::ArgumentParser parser("Process the RandS output with RDataFrames");
   args::HelpFlag help(parser, "help", "show this help menu and exit", {'h', "help"});
   args::ValueFlag<std::string> cutfile(parser, "cutstring file",
					"config file containing the cutstrings and histogram definitions",
					{'c', "cutstrings"});
   args::ValueFlag<std::string> outfile(parser, "output file", "path of the root file to write the histograms to, will be overwritten if it exists", {'o', "output"});
   args::PositionalList<std::string> infiles(parser, "input", "input root files");

   try {
      parser.ParseCLI(argc, argv);
   } catch (args::Help&) {
      std::cout << parser <<std::endl;
      return 1;
   }

   libconfig::Config cfg;
   
   try {
      cfg.readFile(args::get(cutfile).c_str());
      cfg.setAutoConvert(true);
   } catch (libconfig::FileIOException&) {
      std::cerr << "could not read cutstring file. check the -c argument" << std::endl;
      throw;
   } catch (libconfig::ParseException& e) {
      std::cerr << e.what() << std::endl;
   }

   

   auto hdefs = ReadHistogramDefinitions(cfg.lookup("histograms"));
   ROOT::RDataFrame df("PredictionTree", args::get(infiles));

   std::unordered_map<std::string, ROOT::RDF::RInterface<ROOT::Detail::RDF::RJittedFilter, void>> alltheframes;

   auto firstfilter = df.Filter("true");
   CreateFilters(cfg.lookup("cuts"), firstfilter, "", alltheframes);

   std::vector<ROOT::RDF::RResultPtr<TH2D>> hvec;
   hvec.reserve(alltheframes.size()*hdefs.size());

   for (auto& kvp : alltheframes) {
      auto newhistos = CreateHistograms(kvp.second, kvp.first, hdefs);
      hvec.insert(hvec.end(), newhistos.begin(), newhistos.end());
   }

   CalcPrediction(hvec, args::get(outfile));
   
   return 0;
}