// OutROOTProc.cc
// Contact person: P G Jones <p.g.jones@qmul.ac.uk>
// See OutROOTProc.hh for more details.
//—————————————————————---------------------------------------------——//

#include <RAT/OutROOTProc.hh>
#include <RAT/MetaInformation.hh>
#include <RAT/Log.hh>
#include <RAT/DB.hh>
#include <RAT/DS/Entry.hh>
#include <RAT/DS/Run.hh>
using namespace RAT;

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

#include <string>
using namespace std;

OutROOTProc::OutROOTProc() : Processor("OutROOTProc")
{
  fOverrideFileName = "";
  fFileNameSuffix = "";
  fFile = NULL;
  fDSTree = NULL;
  fRunTree = NULL;
  fDSBranch = new DS::Entry();
  fRunBranch = new DS::Run();
  fUpdate = false;

  fAutoFlush =  1000000;
  fAutoSave  = 10000000;
}

OutROOTProc::~OutROOTProc()
{
  if( fFile != NULL )
    {
      // Store the current ROOT global directory (to switch back too) and switch to the file's directory
      TDirectory* const oldDirectory = gDirectory;
      fFile->cd();

      // Write both trees to the file
      if(!fDSTree->Write())
        Log::Die("Could not write DS tree to file!");

      if(!fRunTree->Write())
        Log::Die("Could not write run tree to file!");

      // Write the log and macro to a string in the file
      TObjString* log = new TObjString( Log::GetLogBuffer().c_str() );
      if(!log->Write( "log" ))
        Log::Die("Could not write log to ROOT file");

      TObjString* macro = new TObjString( Log::GetMacro().c_str() );
      if(!macro->Write( "macro" ))
        Log::Die("Could not write macro to ROOT file");

      // Write the database trace (of access)
      TMap *dbtrace = Log::GetDBTraceMap();
      if(!dbtrace->Write("db", TObject::kSingleKey))
        Log::Die("Could not write db trace to ROOT file");

      // complete the meta information with the event counts etc.
      fMetaHelper.SetStoredEvents(fDSTree->GetEntries());
      MetaInformation::Get()->Finalise(fMetaHelper);

      // Write the MetaInformation (RAT and DB state)
      if(!MetaInformation::Get()->GetDSMeta().Write( "meta" ))
        Log::Die("Could not write meta to ROOT file");

      // Switch back to the previous ROOT global directory
      oldDirectory->cd();

      // Cleanup
      fFile->Close(); // also deletes TTree and TObjString?
      delete fFile;
      delete fDSBranch;
      delete fRunBranch;
    }

  // Don't delete branchDS because it looks like the TTree took ownership of it
}

void
OutROOTProc::SetS( const std::string& param,
                   const std::string& value )
{
  if( param == "file" )
    fOverrideFileName = value;
  else if( param == "updatefile" )
    {
      fOverrideFileName = value;
      fUpdate = true;
    }
  else if ( param == "suffix" ) {
    fFileNameSuffix = value;
  }
  else
    throw Processor::ParamUnknown( param );
}

void
OutROOTProc::SetI( const std::string& param,
                   const int value )
{
  if( param == "autosave" )
    {
      if( value < 0 )
        throw ParamInvalid( param, "autosave value must be >= 0" );
      fAutoSave = value;
      // SetAutoSave might be called after the file is already opened
      if( fFile != NULL )
        {
          fDSTree->SetAutoSave( fAutoSave );
          fRunTree->SetAutoSave( fAutoSave );
        }
    }
  else if( param == "autoflush" )
    {
      fAutoFlush = value;
      // SetAutoFlush might be called after the file is already opened
      if( fFile != NULL )
        {
          fDSTree->SetAutoFlush( fAutoFlush );
          fRunTree->SetAutoFlush( fAutoFlush );
        }
    }
}

void
OutROOTProc::BeginOfRun( DS::Run& )
{
  // Get the database file name, this is the default unless the user overrides it
  try
    {
      fDBFileName = DB::Get()->GetLink("IO")->GetS("default_output_filename");
      if( fDBFileName.find( "." ) == string::npos )
        fDBFileName += ".root";
    }
  catch( DBNotFoundError& e ) { Log::Die( "OutROOTProc::BeginOfRun: File name not found in IO[] ratdb." ); }

  if( fFile == NULL )
    OpenFile();

}

Processor::Result
OutROOTProc::DSEvent( DS::Run&,
                      DS::Entry& ds )
{
  if( fFile == NULL )
    OpenFile();

  // Don't let a crazy track with hundreds of thousands of steps flood
  // the output file
  if( ds.MCExists() )
    {
      RAT::DS::MC& mc = ds.GetMC();
      std::vector<unsigned int> trackIDs = mc.GetMCTrackIDs();
      for( size_t i = 0; i < trackIDs.size(); i++ )
      {
          RAT::DS::MCTrack& mctrack = mc.GetMCTrack(trackIDs.at(i));
          size_t nsteps = mctrack.GetMCTrackStepCount();
          if (nsteps > 10000) {
            warn << "outroot: trimming intermediate steps from event #"
                 << mc.GetMCID() << ", track #" << mctrack.GetTrackID()
                 << " which has " << nsteps << " steps" << newline;
            mctrack.PruneIntermediateMCTrackSteps();
          }
        }
    }

  TDirectory* const oldDirectory = gDirectory;
  fFile->cd();
  *fDSBranch = ds;
  if(-1 == fDSTree->Fill()) Log::Die("Could not fill DS tree!");

  oldDirectory->cd();

  return Processor::OK;
}

void
OutROOTProc::EndOfRun( DS::Run& run )
{
  if( fFile == NULL )
    OpenFile();

  TDirectory* const oldDirectory = gDirectory;
  fFile->cd();
  *fRunBranch = run;
  if(-1 == fRunTree->Fill()) Log::Die("Could not fill run tree!");
  oldDirectory->cd();

  fMetaHelper.AddRun(run);
}

void
OutROOTProc::OpenFile()
{
  string fileName = fDBFileName;
  if( !fOverrideFileName.empty() )
    fileName = fOverrideFileName;
  if (!fFileNameSuffix.empty()) {
    info << "OutROOTProc: Appending [" << fFileNameSuffix << "] to the output ROOT file name." << newline;
    string tmp_file = "";
    if( fileName.find( ".root" ) == string::npos ) {
      // -- there is no .root extension
      tmp_file = fileName;
    } else {
      tmp_file = fileName.substr(0,fileName.rfind(".root"));
    }
    fileName = tmp_file;
    fileName += "_";
    fileName += fFileNameSuffix;
    fileName += ".root";
  } else {
    debug << "OutROOTProc: There is **NO** appendix to be added [" << fFileNameSuffix << "] to the file name. " << newline;
  }

  if( fUpdate )
    {
      info << "OutROOTProc: Appending to " << fileName << newline;
      fFile = TFile::Open( fileName.c_str(), "UPDATE" );
    }
  else
    {
      info << "OutROOTProc: Writing to " << fileName << newline;
      fFile = TFile::Open( fileName.c_str(), "RECREATE" );
    }
  if( fFile == NULL )
    Log::Die( "OutROOTProc::OpenFile: Cannot open the file." );

  if( fUpdate )
    {
      fDSTree = reinterpret_cast<TTree*>( fFile->Get( "T" ) );
      fRunTree = reinterpret_cast<TTree*>( fFile->Get( "runT" ) );
      fDSTree->SetBranchAddress( "ds", &fDSBranch, 0 );
      fRunTree->SetBranchAddress( "run", &fRunBranch, 0 );

      // load up the meta information and start a new pass
      DS::Meta *meta = NULL;
      fFile -> GetObject("meta", meta);
      if(meta){
          MetaInformation::Get()->Initialise(meta);
          fMetaHelper.LoadLastPass(*meta);
      }
    }
  else
    {
      // Setup tree
      fDSTree = new TTree( "T", "RAT Tree" );
      fDSTree->Branch( "ds", fDSBranch->ClassName(), &fDSBranch, 32000, 99 );
      fRunTree = new TTree( "runT", "RAT Run Tree" );
      fRunTree->Branch( "run", fRunBranch->ClassName(), &fRunBranch, 32000, 99 );
    }

  fDSTree->SetAutoFlush( fAutoFlush );
  fDSTree->SetAutoSave( fAutoSave );
  fRunTree->SetAutoFlush( fAutoFlush );
  fRunTree->SetAutoSave( fAutoSave );
}