#include <string>

constexpr const int kEntriesPerTree = 10;

void createMainTree(int idx)
{
  TFile file(std::string("maintree_" + std::to_string(idx) + ".root").c_str(), "RECREATE");
  TTree tree("tree", "Main tree");
  int val;
  tree.Branch("val", &val);
  for (int i = 0; i < kEntriesPerTree; ++i) {
    val = -(i + idx * kEntriesPerTree);
    tree.Fill();
  }
  tree.Write();
}

void createFriendTree(int idx)
{
  TFile file(std::string("friendtree_" + std::to_string(idx) + ".root").c_str(), "RECREATE");
  TTree tree("friend", "Friend tree");
  int entry;
  tree.Branch("entry", &entry);
  for (int i = 0; i < 30/*kEntriesPerTree*/; ++i) {
    entry = i/* + idx * kEntriesPerTree*/;
    tree.Fill();
  }
  tree.Write();
}

std::array<std::unique_ptr<TChain>, 2> getReadChains()
{
  std::unique_ptr<TChain> chainFriend{new TChain("friend")};
  chainFriend->Add("friendtree_*.root");

  std::unique_ptr<TChain> chainMain{new TChain("tree")};
  chainMain->Add("maintree_*.root");
  chainMain->AddFriend(chainFriend.get());
  return {std::move(chainMain), std::move(chainFriend)};
}

void checkValues(int entry, int mainVal, int friendVal, const char* diag) {
  if (mainVal != -entry) {
    std::cerr << diag << ": Unexpected value in main tree: expected "
	      << -entry << ", read " << mainVal << '\n';
    exit(1);
  }
  if (friendVal != entry) {
    std::cerr << "Reader: Unexpected value in friend tree: expected "
	      << entry << ", read " << friendVal << '\n';
    exit(1);
  }
}

void ReadSetBranchAddress()
{
  auto chains = getReadChains();
  TChain *chain = chains[0].get();

  int mainVal;
  int friendVal;
  chain->SetBranchAddress("val", &mainVal);
  chain->SetBranchAddress("entry", &friendVal);
  int entry = -1;
  while (chain->LoadTree(++entry)) {
    checkValues(entry, mainVal, friendVal, "BranchAddress");
  }
}

void ReadReader()
{
  auto chains = getReadChains();
  TChain *chainMain = chains[0].get();

  TTreeReader reader(chainMain);
  TTreeReaderValue<int> mainVal(reader, "val");
  TTreeReaderValue<int> friendVal(reader, "entry");
  int entry = 0;

  reader.SetEntry(3);
  checkValues(3, *mainVal, *friendVal, "Reader3");

  reader.SetEntry(13);
  checkValues(13, *mainVal, *friendVal, "Reader13");

  reader.Restart();

  while (reader.Next()) {
    //if (entry % kEntriesPerTree == 0)
    // std::cout << "&*mainVal=" << &*mainVal << " &*friendVal=" << &*friendVal << '\n';
    checkValues(entry, *mainVal, *friendVal, "Reader");
    ++entry;
  }
}

void TTreeReaderFriend()
{
  for (int i = 0; i < 3; ++i) {
    createMainTree(i);
    //createFriendTree(i);
  }
  for (int i = 0; i < 1; ++i) {
    createFriendTree(i);
  }
  

  ReadSetBranchAddress();
  ReadReader();
}