#include <TSystem.h>

#include "IRooTrackerVtxModuleBase.hxx"
#include "IGeomInfo.hxx"
#include "IP0DGeom.hxx"
#include "IOADatabase.hxx"
#include "ICOMETLog.hxx"

ClassImp(COMET::IRooTrackerVtxModuleBase);

using std::string;
//
//------------------------------------------------------------------------------
COMET::IRooTrackerVtxModuleBase::IRooTrackerVtxModuleBase() :
fNVtx(0) {
  fRooTrackerTree = NULL; 
  fInputFileTree = NULL; 
  fInputKinemTree = NULL; 
  this->ResetFileInfo(); 
}
//
//------------------------------------------------------------------------------
COMET::IRooTrackerVtxModuleBase::~IRooTrackerVtxModuleBase() { 
  if(fInputKinemTree != NULL){ delete fInputKinemTree; fInputKinemTree = NULL;}  
  if(fInputFileTree != NULL) { delete fInputFileTree; fInputFileTree = NULL;}
}
//
//------------------------------------------------------------------------------
bool COMET::IRooTrackerVtxModuleBase::FillTree(COMET::ICOMETEvent& event) {

  fNVtx = 0;
  fVtx->Clear();

  /// Make sure not Dereferencing a NULL IHandle to stop the module being
  /// disabled.
  if(!event.Has<COMET::IG4PrimaryVertexContainer>("truth/G4PrimVertex00")){
    COMETVerbose("Skipping event (filling with TClonesArray of size 0)"<<
               " as can not find any primary vertices!");
    return true;
  }
 
  COMET::IHandle<COMET::IG4PrimaryVertexContainer> g4PrimVertex00
    = event.Get<COMET::IG4PrimaryVertexContainer>("truth/G4PrimVertex00");

  /// Loop over the vertices and fill genie or neut roo tracker event objects.
  int i_vertex = -1;
  for(std::vector<COMET::IG4PrimaryVertex>::const_iterator v_iter
    = g4PrimVertex00->begin();
    v_iter != g4PrimVertex00->end();
    ++v_iter){
  i_vertex++;

    COMETVerbose("\nGot G4Vertex -  filename: "<<this->GetFileName((*v_iter))<<
             "\n             -  treename: "<< this->GetTreeName((*v_iter))   <<
             "\n             - origentry: "<< this->GetEntryNum((*v_iter))   <<
             "\n             - currentry: "<< v_iter->GetInteractionNumber() <<
             "\n             -   genname: "<< v_iter->GetGeneratorName()     <<
               "             - COMETTime: "<< v_iter->GetPosition().T());

    /// Load the trees if necessary.
    if(this->LoadPassThroughInput((*v_iter)) == false){
      throw COMET::ENoKinematicsPassThroughTrees();
      return false;
    }

    /// Set the bookkeeping info from the current vertex and pass-through info.
    /// This done here as is common to all generators.
    if(this->UpdateBookKeepingInfo((*v_iter)) == false){
      throw COMET::ENoKinematicsPassThroughTrees();
      return false;
    }

    /// Fill the vertex.
    this->FillVtx(fInputTreeEntryNumber);
    
    /// Check basic pass-through quantities match those of G4PrimaryVertex.
    if(this->CheckMismatch((*v_iter)) == false){ 
      throw COMET::EPassThroughMismatch();
      return false;
    }
  }
  return true;
}
//
//------------------------------------------------------------------------------
bool COMET::IRooTrackerVtxModuleBase::LoadPassThroughInput(
                                              const COMET::IG4PrimaryVertex & vtx){
  /// First check if we need to do anything since the last time.
  if(this->AreTreesLoaded(vtx)){
    COMETDebug("No need to load the pass-through trees again"); 
    return true;
  }
  COMETDebug("Pass-through trees are not loaded for this vertex. Trying to load"<<
           " them now.");
  /// Reset bookkeeping info as need to load the pass-through info again.
  this->ResetVtxInfo();
  
  /// Try to get the pass-through info from the SIMG4 pass-through first.
  if(this->GetTreesFromFile(vtx)){
    COMETDebug("Loaded pass-through input trees from SIMG4 pass-through"<<
             " directory"); 
  }
  /// Else try to recover pass-through from the current working directory or a
  /// user specified input directory. 
  else if(this->GetTreesFromDirectory(vtx)){
    COMETDebug("Loaded input trees from user defined or current working "<<
             "directory"); 
  }
  else {
    /// Else cannot find input files so disabling modlule.
    COMETVerbose("Unable to load pass-through input trees. The module will be " <<
               "disabled and there will be no pass-through information in the"<<
               " analysis output file!");
    return false;
  }

  /// Now set the input tree addresses as implemented in the derived classes.
  this->SetGeneratorTreeAddresses();
  return true;
}
//
//------------------------------------------------------------------------------
bool COMET::IRooTrackerVtxModuleBase::UpdateBookKeepingInfo(
                                              const COMET::IG4PrimaryVertex & vtx){

  /// First get what we can directly from the G4Primary vertex.
  fOrigTreeEntryNumber  = this->GetEntryNum(vtx); 
  fInputTreeEntryNumber = fOrigTreeEntryNumber;
  fInputTreeName        = this->GetTreeName(vtx); 
  fOrigInputFileName    = this->GetFileName(vtx); 
  fGeneratorName        = string(vtx.GetGeneratorName());
  fTimeInSpill          = vtx.GetPosition().T();
  fTruthVertexID        = vtx.GetInteractionNumber();

  /// Now only need to get the input file POT and the entry number of the input 
  /// rootracker tree we have. This is done differently depending on whether we
  /// have the pass-through info or if we are recovering it from a folder.

  /// If have pass-through info use book-keeping trees.
  if(fPassThroughPresent == true){
    int input_file_number = -1;
    fInputKinemTree->SetBranchAddress("inputEntryNum", &fOrigTreeEntryNumber);
    fInputKinemTree->SetBranchAddress("inputFileNum", &input_file_number);
    fInputFileTree->SetBranchAddress("filePOT", &fOrigInputTreePOT);
    fInputFileTree->SetBranchAddress("fileEntries", &fOrigInputTreeEntries); 
    fInputTreeEntryNumber = vtx.GetInteractionNumber();
    fInputKinemTree->GetEntry(fInputTreeEntryNumber);
    fInputFileTree->GetEntry(input_file_number);
  }
  /// otherwise get the information from the original input files.
  else {
    fOrigInputTreePOT     = fRooTrackerTree->GetWeight();
    fOrigInputTreeEntries = fRooTrackerTree->GetEntries();
  }
  return true;
}
//
//------------------------------------------------------------------------------
bool COMET::IRooTrackerVtxModuleBase::AreTreesLoaded(
                                              const COMET::IG4PrimaryVertex & vtx){
  // First check if any trees.
  bool inputfile = (fCurrInputFile != NULL);
  bool rootracker_tree = (fRooTrackerTree != NULL);
  if(!inputfile || !rootracker_tree){
    COMETDebug("Trees are not loaded yet!");
    return false;
  } 

  bool same_name = false;
  // If have pass-through see if have different comet input file.
  if(fPassThroughPresent && fCurrInputFile){
    TFile * curr_inputfile = fLastBeginFile;//IOADatabase::Get().CurrentInputFile();
     same_name = this->ComparePaths(fCurrInputFile->GetName(),
                       curr_inputfile->GetName());  
  }
  // Else if recovering info from directory see if the current input file is 
  // still the correct one.
  else if (fCurrInputFile){
    string filename_vtx = this->GetFileName(vtx);
    same_name=this->ComparePaths(fCurrInputFile->GetName(), filename_vtx.c_str());
  }

  if(!same_name){this->ResetVtxInfo();}
  return same_name;
}
//
//------------------------------------------------------------------------------
bool COMET::IRooTrackerVtxModuleBase::GetTreesFromFile(
                                              const COMET::IG4PrimaryVertex & vtx){
  this->ResetFileInfo();
  fPassThroughPresent = false;

  TFile * curr_infile = fLastBeginFile; //IOADatabase::Get().CurrentInputFile();  
  if(curr_infile != NULL){ 
    if(!curr_infile->IsZombie()){
      /// Make copies of the trees rather than just getting them from the file 
      /// as otherwise get problems when both neut and genie modules later try
      /// to access and delete the same objects. This is because using the 
      /// TFile::Get method we get two pointers to same object in memory.
      TTree * tmpKinemTree = (TTree*) curr_infile->Get("SIMG4/InputKinem");  
      TTree * tmpFileTree = (TTree*) curr_infile->Get("SIMG4/InputFiles");  
      if(!tmpKinemTree || !tmpFileTree){return false;}

      fInputKinemTree = (TTree*)tmpKinemTree->CloneTree();
      fInputFileTree = (TTree*)tmpFileTree->CloneTree();
      delete tmpKinemTree; tmpKinemTree = NULL;
      delete tmpFileTree; tmpFileTree = NULL;
      if(fInputKinemTree == NULL || fInputFileTree == NULL){return false;}

      fCurrInputFile = curr_infile; 
      fPassThroughPresent = true;

      string generator_name = vtx.GetGeneratorName();
      string tree_name = this->GetTreeName(vtx); 
      /// Check if we recognise the generator. This can be by name or treename  and 
      /// is implemented differently in the derived classes. The GENIE module knows
      /// about gRooTracker and the NEUT module knows about nRooTracker.
      if(!this->IsKnownGenerator(generator_name.c_str()) &&
         !this->IsKnownGenerator(tree_name.c_str())     ){
        throw COMET::EIncorrectInputGenerator();
        return false;
      }

      /// And most importantly get the rootracker tree.
      std::string tree_path = string("SIMG4/")+tree_name;
      fRooTrackerTree = (TTree*) fCurrInputFile->Get(tree_path.c_str());
      COMETDebug("Looking for pass-through info in file:"<<
                fCurrInputFile->GetName());
      if(fRooTrackerTree == NULL){ 
        this->ClearInputTrees();
        COMETDebug("Could not get pass-through tree "<< tree_path.c_str() <<
           " from file:"<< fCurrInputFile->GetName());
        throw COMET::ENoKinematicsPassThroughTrees();
        return false;
      }     
      COMETDebug("Successfully found pass-through tree from "<< tree_path.c_str() <<
               "in comet output file: "<< fCurrInputFile->GetName());
      return true;
    }
  }
  return false;
}
//
//------------------------------------------------------------------------------
bool COMET::IRooTrackerVtxModuleBase::GetTreesFromDirectory(
                                              const COMET::IG4PrimaryVertex & vtx){
  this->ResetFileInfo();
  /// Get the input file name from the pass-through tree and see if it exists
  /// in the directory to search. If no directory specified check in the
  /// current working directory.
  std::string directorytosearch(gSystem->WorkingDirectory());
  if(fInputDirectory.size() != 0){
    directorytosearch = std::string(fInputDirectory.c_str());
  } 

  /// First look for the file.
  TString filename( this->GetFileName(vtx) );
  COMETDebug("Looking for input rootracker file: "<< filename.Data() << 
          "\nIn path: \n" << directorytosearch.c_str());
  const char * filepath = gSystem->FindFile(directorytosearch.c_str(), filename);
  if(filepath == 0){
    COMETDebug("Could not find pass-through file in directory!");
    return false;
  }
 
  COMETDebug("Found filepath: "<< filepath);
  /// Now open it and make clones of the relevant files.
  TFile * infile = new TFile(filepath, "OPEN");
  if(infile == 0){
    COMETDebug("Could not open pass-through file!");
    return false;
  }
  /// Now set this to the file reference and flag that we are using a 
  /// directory to recover the pass-through info.
  fCurrInputFile = infile;
  fPassThroughPresent = false;

  string generator_name = vtx.GetGeneratorName();
  string tree_name = this->GetTreeName(vtx); 
  /// Check if we recognise the generator. This can be by name or treename  and 
  /// is implemented differently in the derived classes. The GENIE module knows
  /// about gRooTracker and the NEUT module knows about nRooTracker.
  if(!this->IsKnownGenerator(generator_name.c_str()) &&
     !this->IsKnownGenerator(tree_name.c_str())     ){
    throw COMET::EIncorrectInputGenerator();
    return false;
  }

  /// And most importantly get the rootracker tree.
  std::string tree_path = tree_name;
  fRooTrackerTree = (TTree*) fCurrInputFile->Get(tree_path.c_str());
  COMETDebug("Looking for pass-through info in file:"<<
            fCurrInputFile->GetName());
  if(fRooTrackerTree == NULL){ 
    this->ClearInputTrees();
    COMETDebug("Could not get pass-through tree "<< tree_path.c_str() <<
       " from file:"<< fCurrInputFile->GetName());
    throw COMET::ENoKinematicsPassThroughTrees();
    return false;
  }     
  COMETDebug("Successfully found input tree from external directory in file: "<<
            fCurrInputFile->GetName());
  return true;
}
//
//------------------------------------------------------------------------------
void COMET::IRooTrackerVtxModuleBase::ClearInputTrees(){
  if(fInputKinemTree != NULL){ delete fInputKinemTree; fInputKinemTree = NULL;}
  if(fInputFileTree != NULL) { delete fInputFileTree; fInputFileTree = NULL;}
  if(fRooTrackerTree != NULL){ delete fRooTrackerTree; fRooTrackerTree = NULL;}
}
//
//------------------------------------------------------------------------------
void COMET::IRooTrackerVtxModuleBase::InitializeBranches() {
  fOutputTree->Branch("NVtx", &fNVtx, "NVtx/I", fBufferSize);
  fOutputTree->Branch("Vtx","TClonesArray", &fVtx   , fBufferSize, 1);
}
//
//------------------------------------------------------------------------------
string COMET::IRooTrackerVtxModuleBase::GetFileName(
                                         const COMET::IG4PrimaryVertex & prim_vtx){
  std::string name(prim_vtx.GetFilename());
  std::string::size_type start_pos = name.rfind("/");
  std::string::size_type finish_pos = name.find(":");
  if (start_pos == std::string::npos) start_pos = 0; else ++start_pos;
  std::string filename(name,start_pos, finish_pos -start_pos); 
  COMETVerbose("Got filename " << filename.c_str()  <<
               " from filestring: "<<name.c_str() );
  return filename;
}
//
//------------------------------------------------------------------------------
string COMET::IRooTrackerVtxModuleBase::GetTreeName(
                                         const COMET::IG4PrimaryVertex & prim_vtx){
  std::string name(prim_vtx.GetFilename());
  std::string::size_type start_pos = name.find(":");
  std::string::size_type finish_pos = name.rfind(":");
  if (start_pos == std::string::npos) start_pos = 0; else ++start_pos;
  std::string treename(name,start_pos, finish_pos - start_pos); 
  COMETVerbose("Got treename " << treename.c_str()  << 
               " from filestring: "<<name.c_str() );
  return treename;
}
//
//------------------------------------------------------------------------------
int COMET::IRooTrackerVtxModuleBase::GetEntryNum(
                                         const COMET::IG4PrimaryVertex & prim_vtx){
  int entrynum = -1;
  std::string name(prim_vtx.GetFilename());
  std::string::size_type start_pos = name.rfind(":");
  if (start_pos == std::string::npos) start_pos = 0; else ++start_pos;
  std::string filename(name,start_pos);
  entrynum = std::atoi(filename.c_str()); 
  COMETVerbose("Got entry " << entrynum << " from filestring: "<<name.c_str() );
  return entrynum;
}
//
//------------------------------------------------------------------------------
string COMET::IRooTrackerVtxModuleBase::GetFileFromPath(const char * path){
string pathstr(path);
std::string::size_type start_pos = pathstr.rfind("/");
if (start_pos == std::string::npos) start_pos = 0; else ++start_pos;
string substring(pathstr, start_pos);
return substring;
} 
//
//------------------------------------------------------------------------------
bool COMET::IRooTrackerVtxModuleBase::ComparePaths(const char* path1, 
                                                    const char* path2){
string file1name = this->GetFileFromPath(path1);
string file2name = this->GetFileFromPath(path2);
return (file1name.find(file2name) != file1name.npos); 
} 
//
//------------------------------------------------------------------------------
void COMET::IRooTrackerVtxModuleBase::ResetFileInfo(){
  fCurrInputFile = NULL; /// Potential small mem leak but ok as may reuse
  fPassThroughPresent = false;
  this->ResetVtxInfo(); 
}
//
//------------------------------------------------------------------------------
void COMET::IRooTrackerVtxModuleBase::ResetVtxInfo(){
  fOrigInputFileName = string("");
  fOrigTreeEntryNumber = -1;
  fOrigInputTreeEntries = -1;
  fOrigInputTreePOT = -1.0;
  fInputTreeName = string("");
  fGeneratorName = string("");
  fInputTreeEntryNumber = -1;
  fTimeInSpill = -1.0;
  fTruthVertexID = -1;
}
//
//------------------------------------------------------------------------------
Bool_t COMET::IRooTrackerVtxModuleBase::ProcessFirstEvent(COMET::ICOMETEvent& event) {
  bool IsDataFile = event.GetContext().IsDetector();
  if(IsDataFile) {
      throw COMET::EDataFile();
  }
  return true;
}
//------------------------------------------------------------------------------
void COMET::IRooTrackerVtxModuleBase::SetFluxTreeAddresses(COMET::IJNuBeamFlux * flux){
  // Set flux branch addresses common to both generators. 
  // This method should be called from the derived class's definition of 
  // SetGeneratorTreeAddresses() and should be passed a ptr to an object
  // inheriting from COMET::IJNuBeamFlux  
  if(!fRooTrackerTree){ 
    COMETWarn("Cannot set flux tree addresses as no rootracker tree loaded!");
    return;
  }
  // Need to decide whether need to put warning messages if cannot find a branch
  TBranch * fBr = 0;
  if( (fBr = fRooTrackerTree->GetBranch("NuFluxEntry"))     ) fBr->SetAddress(&flux->NuFluxEntry); 
  if( (fBr = fRooTrackerTree->GetBranch("NuFileName"))      ) fBr->SetAddress(&flux->NuFileName); 
  if( (fBr = fRooTrackerTree->GetBranch("NuParentPdg"))     ) fBr->SetAddress(&flux->NuParentPdg);
  if( (fBr = fRooTrackerTree->GetBranch("NuParentDecMode")) ) fBr->SetAddress(&flux->NuParentDecMode); 
  if( (fBr = fRooTrackerTree->GetBranch("NuParentDecX4"))   ) fBr->SetAddress( flux->NuParentDecX4); 
  if( (fBr = fRooTrackerTree->GetBranch("NuParentDecP4"))   ) fBr->SetAddress( flux->NuParentDecP4); 
  if( (fBr = fRooTrackerTree->GetBranch("NuCospibm"))       ) fBr->SetAddress(&flux->NuCospibm); 
  if( (fBr = fRooTrackerTree->GetBranch("NuNorm"))          ) fBr->SetAddress(&flux->NuNorm); 
  if( (fBr = fRooTrackerTree->GetBranch("NuParentProP4"))   ) fBr->SetAddress( flux->NuParentProP4); 
  if( (fBr = fRooTrackerTree->GetBranch("NuParentProX4"))   ) fBr->SetAddress( flux->NuParentProX4); 
  if( (fBr = fRooTrackerTree->GetBranch("NuCospi0bm"))      ) fBr->SetAddress(&flux->NuCospi0bm); 
  if( (fBr = fRooTrackerTree->GetBranch("NuRnu"))           ) fBr->SetAddress(&flux->NuRnu); 
  if( (fBr = fRooTrackerTree->GetBranch("NuXnu"))           ) fBr->SetAddress( flux->NuXnu); 
  if( (fBr = fRooTrackerTree->GetBranch("NuIdfd"))          ) fBr->SetAddress(&flux->NuIdfd); 
  if( (fBr = fRooTrackerTree->GetBranch("NuGipart"))        ) fBr->SetAddress(&flux->NuGipart); 
  if( (fBr = fRooTrackerTree->GetBranch("NuGpos0"))         ) fBr->SetAddress( flux->NuGpos0); 
  if( (fBr = fRooTrackerTree->GetBranch("NuGvec0"))         ) fBr->SetAddress( flux->NuGvec0); 
  if( (fBr = fRooTrackerTree->GetBranch("NuGamom0"))        ) fBr->SetAddress(&flux->NuGamom0); 
  if( (fBr = fRooTrackerTree->GetBranch("NuNg"))            ) fBr->SetAddress(&flux->NuNg); 
  if( (fBr = fRooTrackerTree->GetBranch("NuGp"))            ) fBr->SetAddress( flux->NuGp); 
  if( (fBr = fRooTrackerTree->GetBranch("NuGcosbm"))        ) fBr->SetAddress( flux->NuGcosbm); 
  if( (fBr = fRooTrackerTree->GetBranch("NuGv"))            ) fBr->SetAddress( flux->NuGv); 
  if( (fBr = fRooTrackerTree->GetBranch("NuGpid"))          ) fBr->SetAddress( flux->NuGpid); 
  if( (fBr = fRooTrackerTree->GetBranch("NuGmec"))          ) fBr->SetAddress( flux->NuGmec); 
  if( (fBr = fRooTrackerTree->GetBranch("NuEnusk"))         ) fBr->SetAddress(&flux->NuEnusk); 
  if( (fBr = fRooTrackerTree->GetBranch("NuNormsk"))        ) fBr->SetAddress(&flux->NuNormsk); 
  if( (fBr = fRooTrackerTree->GetBranch("NuAnorm"))         ) fBr->SetAddress(&flux->NuAnorm); 
  if( (fBr = fRooTrackerTree->GetBranch("NuGmat"))          ) fBr->SetAddress( flux->NuGmat); 
  if( (fBr = fRooTrackerTree->GetBranch("NuGdistc"))        ) fBr->SetAddress( flux->NuGdistc); 
  if( (fBr = fRooTrackerTree->GetBranch("NuGdistal"))       ) fBr->SetAddress( flux->NuGdistal); 
  if( (fBr = fRooTrackerTree->GetBranch("NuGdistti"))       ) fBr->SetAddress( flux->NuGdistti); 
  if( (fBr = fRooTrackerTree->GetBranch("NuGdistfe"))       ) fBr->SetAddress( flux->NuGdistfe); 
  if( (fBr = fRooTrackerTree->GetBranch("NuVersion"))       ) fBr->SetAddress(&flux->NuVersion); 
  if( (fBr = fRooTrackerTree->GetBranch("NuTuneid"))        ) fBr->SetAddress(&flux->NuTuneid); 
  if( (fBr = fRooTrackerTree->GetBranch("NuNtrig"))         ) fBr->SetAddress(&flux->NuNtrig); 
  if( (fBr = fRooTrackerTree->GetBranch("NuPint"))          ) fBr->SetAddress(&flux->NuPint); 
  if( (fBr = fRooTrackerTree->GetBranch("NuBpos"))          ) fBr->SetAddress( flux->NuBpos); 
  if( (fBr = fRooTrackerTree->GetBranch("NuBtilt"))         ) fBr->SetAddress( flux->NuBtilt); 
  if( (fBr = fRooTrackerTree->GetBranch("NuBrms"))          ) fBr->SetAddress( flux->NuBrms); 
  if( (fBr = fRooTrackerTree->GetBranch("NuEmit"))          ) fBr->SetAddress( flux->NuEmit); 
  if( (fBr = fRooTrackerTree->GetBranch("NuAlpha"))         ) fBr->SetAddress( flux->NuAlpha); 
  if( (fBr = fRooTrackerTree->GetBranch("NuHcur"))          ) fBr->SetAddress( flux->NuHcur); 
  if( (fBr = fRooTrackerTree->GetBranch("NuRand"))          ) fBr->SetAddress(&flux->NuRand); 
//  if( (fBr = fRooTrackerTree->GetBranch("NuRseed"))         ) fBr->SetAddress(flux->NuRseed);
  fBr = 0;
}
//------------------------------------------------------------------------------