// Uses the run field to divide a tree (chain) into smaller ones
#include "TROOT.h"
#include "TTree.h"
#include "TFile.h"
#include "TLeaf.h"
#include "TString.h"
#include "TEventList.h"
#include <iostream>
#include <set>
#include <vector>
//#include <list>
//#include <map> // for std::pair
#include <algorithm>
#include <stdio.h>


class TRunDivider {
 public:
  TRunDivider ();
  // shallow copy OK: ~TRunDivider ();
  
  void SetFieldNames (TString& fieldName);

  // for all functions: - negative returns are errors, otherwise a count
  //                    - tree may be a chain
  void Clear ();
  int Read (TTree &tree); // tree is the input
  int WritePart (TTree &tree, const int index, const int Nparts, // tree is the input
		 TString newFileName = "/") const; // a trailing / is replaced with part number
  int WriteRuns (TTree &tree,
		 TString newFileName = "/") const; // a trailing / is replaced with part number
  int Size() const;

  int Dump (int N=10);

  // shallow copy OK: TRunDivider (const TRunDivider& that);
  // shallow copy OK: TRunDivider& operator = (const TRunDivider& that);

  typedef std::set  <int> KeySet;
  typedef std::vector<int> KeyVector;
  typedef std::vector<int>::iterator KeyIter;

 private:
  KeySet _keySet;
  KeyVector _keyVector;
  TString _name;
};

TRunDivider::TRunDivider ()
{
  _name ="run";
}

void TRunDivider::Clear ()
{
  _keySet.clear();
  _keyVector.clear();
}

void TRunDivider::SetFieldNames (TString& fieldName)
{
  _name = fieldName;
}

int TRunDivider::Size () const
{
  return _keyVector.size();
}

int TRunDivider::Dump (int N)
{
  if (N <= 0) N=_keyVector.size();
  int M=N;
  std::cout<<"    "<<_name<<endl;
  for (KeyIter it=_keyVector.begin(); it != _keyVector.end() && N>0; ++it, --N) {
    std::cout<<M-N<<") "<<*it<<std::endl;
  }
  return this->Size();
}


int TRunDivider::Read (TTree &tree)
{
  TEventList* list = tree.GetEventList();
  TLeaf *plr;
  Int_t prevTree = -1;

  _keySet.clear();
  _keyVector.clear();

  //read everything into the list as before, then copy it into the vector and sort it

  const Long64_t N = list ? Long64_t (list->GetN()) : Long64_t (tree.GetEntries());
  for (int i=0; i<N; ++i) {
    Long64_t ientry = list ? list->GetEntry (i) : Long64_t (i);
    if (ientry < 0) {std::cerr<<"bug! ientry: "<<ientry<<std::endl; return -2;}

    //if (i<10 || 0==i%1000 || i>N-10) cout<<"N: "<<N<<" i: "<<i<<" ientry: "<<ientry<<" list? "<<((bool) list)<<endl;

    Int_t itmp = tree.LoadTree (ientry); // switch to the right tree in the chain
    if (itmp < 0) {std::cerr<<"bug? what's going on with itmp: "<<itmp<<std::endl; return -3;}

    //if (i<10 || 0==i%1000 || i>N-10) cout<<"Loader itmp: "<<itmp<<endl;
    // leaves need to be found after loading the tree (happens several times in a chain)
    if (tree.GetTreeNumber () != prevTree) {
      prevTree = tree.GetTreeNumber ();
      //cout<<"DBG: getting leaf."<<endl;
      plr = tree.GetLeaf (_name);
      if (plr == 0) {std::cerr<<"Can't find key (run) in tree."<<std::endl; return -1;}
    }
    //if (i<10 || 0==i%1000 || i>N-10) cout<<"checked tree num"<<endl;

    tree.GetEntry(ientry); 

    Int_t run = (Int_t) plr -> GetValue();
    //if (i<10 || 0==i%1000 || i>N-10) cout<<"run: "<<run<<endl;

    _keySet.insert (run);
    //if (i<10 || 0==i%1000 || i>N-10) cout<<"Inserted"<<endl;
  }

  cout<<"DBG: copying"<<endl;
  std::copy (_keySet.begin(), _keySet.end(), inserter(_keyVector, _keyVector.begin())); //,_keyVector.begin());
  cout<<"DBG: sorting"<<endl;
  std::sort (_keyVector.begin(), _keyVector.end());
  cout<<"DBG: done"<<endl;

  return this->Size();
}




int TRunDivider::WritePart (TTree &tree, const int index, const int Nparts, 
		 TString newFileName) const
{
  if (index >= Nparts || index < 0 || Nparts <1) return 0;

  if (newFileName.EndsWith("/")) {
    newFileName.Chop(); // removes /
    newFileName.Append (Form ("_%d", index));
  }

  // figure out the requested range
  int nTot = this->Size();
  cout<<"DBG: Will look at "<<nTot*(index-1)/Nparts<<" and "<<nTot*index/Nparts - 1<<". nTot: "<<nTot<<flush;
  int minRun = _keyVector [nTot*index/Nparts];
  int maxRun = _keyVector [nTot*(index+1)/Nparts - 1];
  cout<<" --> "<<minRun<<" - "<<maxRun<<endl;
  
  // just get iterators to the right indices in the vector, and check that
  TFile f(newFileName, "recreate");
  TTree *newTree = tree.GetTree()->CloneTree(0);
  if (newTree == 0) {std::cerr<<"Can't clone tree."<<std::endl; return -1;}

  TLeaf* plr;
  Int_t prevTree = -1;
  TEventList* list = tree.GetEventList();
  const Long64_t N = list ? Long64_t (list->GetN()) : Long64_t (tree.GetEntriesFast());
   
  for (Int_t i=0; i<N; ++i) {
    Long64_t ientry = list ? list->GetEntry (i) : Long64_t (i);
    if (ientry < 0) {std::cerr<<"bug! ientry: "<<ientry<<std::endl; return -2;}

    Int_t itmp = tree.LoadTree (ientry); // switch to the right tree in the chain
    if (itmp < 0) {std::cerr<<"bug? what's going on with itmp: "<<itmp<<std::endl; return -3;}

    // leaves need to be found after loading the tree (happens several times in a chain)
    if (tree.GetTreeNumber () != prevTree) {
      prevTree = tree.GetTreeNumber ();
      plr = tree.GetLeaf (_name);
      if (plr == 0) {std::cerr<<"Can't find key (run) in tree."<<std::endl; return -1;}
    }

    tree.GetEntry(ientry); 

    Int_t run = (Int_t) plr -> GetValue();

    if (run >= minRun && run <= maxRun) {
      newTree->Fill();
    }
  }
  newTree->AutoSave();
  return newTree->GetEntriesFast();
}

int TRunDivider::WriteRuns (TTree &tree, TString newFileName) const
{
  int Nkeys = _keyVector.size();
  TFile *files [Nkeys];
  TTree *trees [Nkeys];

  for (int ikey=0; ikey<Nkeys; ++ikey) {
    TString curName = newFileName;
    int curRun = _keyVector [ikey];
    
    if (curName.EndsWith("/")) {
      curName.Chop(); // removes /
      curName.Append (Form ("_%d", curRun));
    }

    files [ikey] = new TFile (curName, "RECREATE");
    if (files [ikey] == 0) {cerr<<"Can't allocate file #"<<ikey<<" called: "<<curName<<endl; return -11;}
    if (! files [ikey] -> IsWritable()) 
      {cerr<<"Can't open for writing file #"<<ikey<<" called: "<<curName<<endl; return -12;}
    trees [ikey] = tree.GetTree()->CloneTree(0);
    if (trees [ikey] == 0) {std::cerr<<"Can't clone tree for clone #"<<ikey<<std::endl; return -13;}
  }

  TLeaf* plr;
  Int_t prevTree = -1;
  TEventList* list = tree.GetEventList();
  const Long64_t N = list ? Long64_t (list->GetN()) : Long64_t (tree.GetEntriesFast());
   
  for (Int_t i=0; i<N; ++i) {
    Long64_t ientry = list ? list->GetEntry (i) : Long64_t (i);
    if (ientry < 0) {std::cerr<<"bug! ientry: "<<ientry<<std::endl; return -2;}

    Int_t itmp = tree.LoadTree (ientry); // switch to the right tree in the chain
    if (itmp < 0) {std::cerr<<"bug? what's going on with itmp: "<<itmp<<std::endl; return -3;}

    // leaves need to be found after loading the tree (happens several times in a chain)
    if (tree.GetTreeNumber () != prevTree) {
      prevTree = tree.GetTreeNumber ();
      plr = tree.GetLeaf (_name);
      if (plr == 0) {std::cerr<<"Can't find key (run) in tree."<<std::endl; return -1;}
    }

    tree.GetEntry(ientry); 

    Int_t run = (Int_t) plr -> GetValue();

    bool found = false;
    for (int ikey=0; ikey<N && !found; ++ikey) {
      if (_keyVector [ikey] == run) {
	trees [ikey] -> Fill();
	found = true;
      }
    }
    if (!found) {std::cerr<<"Unknown run ("<<run<<"). Aborting!"<<std::endl; return -4;}
  }

  int sum = 0;
  for (int ikey=0; ikey<Nkeys; ++ikey) {
    sum += trees [ikey]->GetEntriesFast();
    cout<<"Writing tree #"<<ikey<<" with "<<trees [ikey]->GetEntriesFast()<<" entries."<<endl;
    trees [ikey]->AutoSave();
    delete trees [ikey]; // calls Close and writes.
  }
  return sum;
}