#include <ROOT/RNTupleWriter.hxx>
#include <ROOT/RNTupleReader.hxx>
#include <ROOT/RNTupleModel.hxx>

#define FILEPATH "ignore.root"
#define TUPLENAME "tuple"

void print_nbytes_read(ROOT::RNTupleReader& reader) {
  int nbytes_read = (
      reader.GetMetrics().GetCounter("RNTupleReader.RPageSourceFile.szReadPayload")->GetValueAsInt()
      +
      reader.GetMetrics().GetCounter("RNTupleReader.RPageSourceFile.szReadOverhead")->GetValueAsInt()
  );
  std::cout << nbytes_read << "B" << std::endl;
}

void try_reading_method_impl(
    std::string_view method_name,
    std::function<std::unique_ptr<ROOT::RNTupleReader>(std::unique_ptr<ROOT::RNTupleReader>)> method
) {
  try {
    std::cout << method_name << ": " << std::flush;
    auto reader = ROOT::RNTupleReader::Open(TUPLENAME, FILEPATH);
    reader->EnableMetrics();
    reader = method(std::move(reader));
    print_nbytes_read(*reader);
  } catch (const ROOT::RException& e) {
    std::cout << "\n  Failed: " << e.what() << std::endl;
  }
}

#define try_reading_method(name) try_reading_method_impl(#name, name)

std::unique_ptr<ROOT::RNTupleReader> naive(std::unique_ptr<ROOT::RNTupleReader> reader) {
  auto keep = reader->GetModel().GetDefaultEntry().GetPtr<int>("to_keep");
  for (auto i: *reader) {
    reader->LoadEntry(i);
  }
  return std::move(reader);
}

std::unique_ptr<ROOT::RNTupleReader> own_model(std::unique_ptr<ROOT::RNTupleReader> reader) {
  auto model = ROOT::RNTupleModel::Create();
  auto to_keep = model->MakeField<int>("to_keep");
  reader = ROOT::RNTupleReader::Open(std::move(model), TUPLENAME, FILEPATH);
  reader->EnableMetrics();
  for (auto i : *reader) {
    reader->LoadEntry(i);
  }
  return std::move(reader);
}

std::unique_ptr<ROOT::RNTupleReader> skimmed_model(std::unique_ptr<ROOT::RNTupleReader> reader) {
  auto model = ROOT::RNTupleModel::Create();
  for (auto& value : reader->GetModel().GetDefaultEntry()) {
    if (value.GetField().GetFieldName().find("ignore") != std::string::npos) {
      continue;
    }
    const auto& name = value.GetField().GetFieldName();
    model->AddField(value.GetField().Clone(name));
    model->GetDefaultEntry().BindValue<void>(name, value.GetPtr<void>());
    //model->GetDefaultEntry().EmplaceNewValue(name);
  }

  reader = ROOT::RNTupleReader::Open(std::move(model), TUPLENAME, FILEPATH);
  auto keep = reader->GetModel().GetDefaultEntry().GetPtr<int>("to_keep");
  reader->EnableMetrics();
  for (auto i: *reader) {
    reader->LoadEntry(i);
  }
  return std::move(reader);
}

std::unique_ptr<ROOT::RNTupleReader> view(std::unique_ptr<ROOT::RNTupleReader> reader) {
  auto view = reader->GetView<int>("to_keep");
  for (auto i : reader->GetEntryRange()) {
    view(i);
  }
  return std::move(reader);
}

int main() {
  {
    std::cout << "writing" << std::endl;

    auto model = ROOT::RNTupleModel::Create();
    auto to_ignore = model->MakeField<float>("to_ignore");
    auto to_keep = model->MakeField<int>("to_keep");
    auto writer = ROOT::RNTupleWriter::Recreate(std::move(model), TUPLENAME, FILEPATH);
    for (int i{0}; i < 10; i++) {
      *to_ignore = 10.5*i;
      *to_keep = i;
      writer->Fill();
    }
  }

  std::cout << "reading" << std::endl;
  try_reading_method(naive);
  try_reading_method(own_model);
  try_reading_method(skimmed_model);
  try_reading_method(view);
}