////////////////////////////////////////////////////////////////////////////////
//
// modifytree_x
// ------------
//
////////////////////////////////////////////////////////////////////////////////


#include <TSystem.h>
#include <TFile.h>
#include <TTree.h>
#include <TRandom.h>

#include <stdexcept>
#include <iostream>
#include <iomanip>
#include <fstream>
#include <sstream>
#include <string>
#include <vector>
#include <set>
#include <map>


using namespace std;


////////////////////////////////////////////////////////////////////////////////
// declare global functions and typedefs
////////////////////////////////////////////////////////////////////////////////

typedef pair<int,int>   cid_t;
typedef map<cid_t,int>  map_t;
typedef map_t::iterator mapiter_t;

bool cmdline(int argc, char** argv,map<string,string>& opts);
bool splitstring(const string& istring, vector<int>& values, string separator);
bool splitstring(const string& istring, vector<string>& values, string separator);

////////////////////////////////////////////////////////////////////////////////
// main
////////////////////////////////////////////////////////////////////////////////

//______________________________________________________________________________
int main(int argc, char** argv)
{
    // interpret command line
    map<string,string> opts;
    if (!cmdline(argc,argv,opts)) return 0;
    
    string input       = opts["input"];
    string output      = opts["output"];
    string skipjt5     = opts["skipjt5"];
    
    if (input.empty()) {
	cout<<"USAGE: modifytree_x -input <tree.root> [opts]  "
	    <<endl;
	return 0;
    }
    // print parameters
    cout<<endl
        <<"------------------------------------------------------------\n"
        <<" modifytree_x parameters\n"
        <<"------------------------------------------------------------\n"
        <<"| input:      "<<input      <<"\n"
	<<"| output:     "<<output     <<"\n"
	<<"| skipjt5:    "<<skipjt5    <<"\n"
	<<"------------------------------------------------------------\n"
        <<endl;
     
    bool skipjt = 0;    
    if(!skipjt5.empty()){   
      stringstream ss;
      ss<<skipjt5;
      ss>>skipjt;
      ss.clear();
    }
//    cout << "skip lowest pt jet ? " << skipjt <<endl;
        
    // open tree-file
    TFile* ifile = new TFile(input.c_str(),"READ");
    if(!ifile->IsOpen()) { cout<<"Can't open "<<ifile<<endl; return 0; }
    TTree* itree=(TTree*)ifile->Get("METree");
    
    if(skipjt){   
	itree->SetBranchStatus("njt",0);
	itree->SetBranchStatus("jte",0);
	itree->SetBranchStatus("jteta",0);
	itree->SetBranchStatus("jtphi",0);
	itree->SetBranchStatus("jtpt",0);
    }
    TFile* ofile=new TFile("tmp.root","RECREATE");
    TTree* otree=itree->CloneTree();
    if(skipjt){   
	itree->SetBranchStatus("njt",1);
	itree->SetBranchStatus("jte",1);
	itree->SetBranchStatus("jteta",1);
	itree->SetBranchStatus("jtphi",1);
	itree->SetBranchStatus("jtpt",1);
    }
        
    //modify tree 
    //setup input tree 
    int njt;        
    int jtmcid[20]; 
    double jte[20]; 
    double jteta[20]; 
    double jtphi[20]; 
    double jtpt[20]; 
    itree->SetBranchAddress("njt",     &njt);    
    itree->SetBranchAddress("jtmcid",jtmcid);
    itree->SetBranchAddress("jte",jte);
    itree->SetBranchAddress("jteta",jteta);
    itree->SetBranchAddress("jtphi",jtphi);
    itree->SetBranchAddress("jtpt",jtpt);
    
    //setup output tree 
    int njtnew;        
    double jtenew[20]; 
    double jtetanew[20]; 
    double jtphinew[20]; 
    double jtptnew[20]; 
    TBranch* b_njt;    
    TBranch* b_jte;    
    TBranch* b_jteta;    
    TBranch* b_jtphi;    
    TBranch* b_jtpt;    
    if(skipjt){   
        b_njt   = otree->Branch("njt",    &njtnew,    "njt/I");    
        b_jte   = otree->Branch("jte",    jtenew,    "jte[njt]/D");    
        b_jteta = otree->Branch("jteta",  jtetanew,  "jteta[njt]/D");    
        b_jtphi = otree->Branch("jtphi",  jtphinew,  "jtphi[njt]/D");    
        b_jtpt  = otree->Branch("jtpt",   jtptnew,   "jtpt[njt]/D");    
    }
    else{   
        b_njt   = NULL;    
        b_jte   = NULL;    
        b_jteta = NULL;    
        b_jtphi = NULL;    
        b_jtpt  = NULL;    
    }
    //loop over events and modify required values (e.g.: leptype)
    for (int ievt=0;ievt<itree->GetEntries();ievt++) {
	itree->GetEntry(ievt);
	//skip lowest pt jt 
        if(skipjt){  
            double skippt = 100000;
            int iskip = 0;
	    for(int jj=0;jj<njt;jj++) 
            {
              if (jtpt[jj]< skippt) 
              {
//                cout<<"jet to skip: "<<jj<<", "<<jtpt[jj]<<"<"<<skippt<<endl;
                iskip = jj;
                skippt = jtpt[jj];
               }
	    }
            cout<<"evt "<<ievt<<", jet to skip: "<<iskip<<", jtpt "<<skippt<<endl;
            njtnew = njt-1;
	    for(int i=0;i<njt;i++){
              if(i!=iskip) {
                jtenew[i]   = jte[i];    
                jtetanew[i] = jteta[i];    
                jtphinew[i] = jtphi[i];    
                jtptnew[i]  = jtpt[i];
              }
            }    
            b_njt->Fill();	
            b_jte->Fill();	
            b_jteta->Fill();	
            b_jtphi->Fill();	
            b_jtpt->Fill();	
        }
    }
    //save and clean up
    if(output.empty()) output = input;
    cout<<"Save new tree to "<<output<<" ..."<<flush;
    ifile->Close();
    delete ifile;
    ofile->Write();
    ofile->Close();
    delete ofile;
    string cmd("mv tmp.root "+output);
    gSystem->Exec(cmd.c_str());
    cout<<" done."<<endl;
    return 0;
}


////////////////////////////////////////////////////////////////////////////////
// implement global functions
////////////////////////////////////////////////////////////////////////////////

//______________________________________________________________________________
bool cmdline(int argc, char** argv, map<string,string>& opts)
{
    opts.clear();

    // defaults
    opts["input"]       = "";
    opts["output"]      = "";
    opts["skipjt5"]     = "";
    
    for (int i=1;i<argc;i++) {
	string opt=argv[i];
	if(0!=opt.find("-")) {
	    cout<<"ERROR: options start with '-'!"<<endl;
	    return false;
	}
	opt.erase(0,1);
	if(opts.find(opt)==opts.end()) {
	    cout<<"ERROR: invalid option '"<<opt<<"'!"<<endl;
	    return false;
	}
	string nxtopt=argv[i+1];
	if(0==nxtopt.find("-")||i+1>=argc) {
	    cout<<"ERROR: option '"<<opt<<"' requires value!"<<endl;
	    return false;
	}
	opts[opt]=nxtopt;
	i++;
    }

    return true;
}