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

#include <G4UIdirectory.hh>
#include <G4UIcmdWithAString.hh>

#include <RAT/InROOTProducer.hh>
#include <RAT/ProcBlock.hh>
#include <RAT/SignalHandler.hh>
#include <RAT/Log.hh>
#include <RAT/DB.hh>
#include <RAT/DBCommandLoader.hh>
#include <RAT/MetaInformation.hh>
#include <RAT/DS/Entry.hh>
#include <RAT/DS/Run.hh>
#include <RAT/DS/BitMask.hh>
using namespace RAT;

#include <TChain.h>
#include <TFile.h>

#include <string>
#include <vector>

InROOTProducer::InROOTProducer()
  : fDSTree("T"), fRunTree("runT")
{
  fMainBlock = NULL;
  Init();
}

InROOTProducer::InROOTProducer( ProcBlock* block )
  : fDSTree("T"), fRunTree("runT")
{
  SetMainBlock(block);
  Init();
}

void InROOTProducer::Init()
{
  // Build commands
  G4UIdirectory* inrootDir = new G4UIdirectory( "/rat/inroot/" );
  inrootDir->SetGuidance( "Read Events from ROOT file" );

  fReadCmd = new G4UIcommand( "/rat/inroot/read", this );
  fReadCmd->SetGuidance( "Start reading in the root file" );
  fReadCmd->AvailableForStates( G4State_Idle );

  fLoadDefaultCmd = new G4UIcommand( "/rat/inroot/load_default", this );
  fLoadDefaultCmd->SetGuidance( "Load the default filename" );
  fLoadDefaultCmd->AvailableForStates( G4State_PreInit );

  fLoadCmd = new G4UIcmdWithAString( "/rat/inroot/load", this );
  fLoadCmd->SetParameterName( "filename", false );
  fLoadCmd->SetGuidance( "Load the filename" );
  fLoadCmd->AvailableForStates( G4State_PreInit );
}

G4String
InROOTProducer::GetCurrentValue( G4UIcommand* )
{
  Log::Die( "InROOTProducer::GetCurrentValue:invalid inroot \"get\" command" );
  return G4String( "You win! How did you manage to see this?." );
}


void InROOTProducer::SetNewValue( G4UIcommand* command, G4String newValue )
{
  if( command == fLoadCmd )
    LoadFile( newValue );
  else if( command == fLoadDefaultCmd ) {
    std::vector<std::string> default_files =DB::Get()->GetLink( "IO" )->GetSArray( "default_input_filename" );
    for (size_t i = 0; i < default_files.size(); ++i) {
      LoadFile( default_files.at(i));
    }
  } else if( command == fReadCmd )
    ReadEntries();
  else
    Log::Die( "InROOTProducer::SetNewValue: Invalid command" );
}

void
InROOTProducer::LoadFile( const std::string& fileName )
{
  // First of all check that the file actually exists
  TFile *testFile = TFile::Open( fileName.c_str() );
  Log::Assert( testFile != NULL, "InROOTProducer::LoadFile: File seems to be invalid!" );
  Log::Assert( testFile->IsZombie() != true, "InROOTProducer::LoadFile: File seems to be invalid!" );
  testFile->Close();
  delete testFile;

  info << "InROOTProducer::LoadFile : Loading file :: " << fileName << newline;
  fDSTree.Add( fileName.c_str() );
  fRunTree.Add( fileName.c_str() );
  // Load the Meta information
  DS::Meta *meta = NULL;
  fDSTree.GetFile()->GetObject("meta",meta);
  if(meta){
    MetaInformation::Get()->Initialise(meta);
  }
  else{
    warn<<"*** InROOTProducer::LoadEvents: WARNING ***"<<newline;
    warn<<"* No meta information found in "<<fileName<<": Creating new one."<<newline;
    warn<<"* This implies that the conditions under which this file was produced"<<newline;
    warn<<"* is unknown and the correct constants might NOT apply!"<<newline;
    warn<<"***********************************************************"<<newline;
    MetaInformation::Get()->Initialise();
  }
  DBCommandLoader::LoadCommands( MetaInformation::Get()->GetDSMeta() );
}

void
InROOTProducer::ReadEntries()
{
  Log::Assert( fMainBlock != NULL, "InROOTProducer::ReadEvents: No main block declared. (Place this command after processors)." );
  const Long64_t dsEntries = fDSTree.GetEntries();
  Log::Assert( dsEntries > 0 && fRunTree.GetEntries() > 0,"InROOTProducer::ReadEvents: Cannot load the T or runT in the file." );
  info << "InROOTProducer::ReadEvents:" << (int)dsEntries << " entries found. " << (int)fRunTree.GetEntries() << " runs found." << newline;

  // Load the branches
  DS::Entry* dsBranch = new DS::Entry();
  fDSTree.SetBranchAddress( "ds", &dsBranch );
  DS::Run* runBranch = new DS::Run();
  fRunTree.SetBranchAddress( "run", &runBranch );
  Long64_t iRun = 0;
  fRunTree.GetEntry( iRun );
  long long defaultNumEvents = DB::Get()->GetLink( "IO" )->GetI( "default_num_events" );
  long long nEntries; //number of entries to loop over
  //How many events to run over
  if (defaultNumEvents >= 0)
      nEntries = defaultNumEvents;
  else
      nEntries = fDSTree.GetEntries();

  for( Long64_t iEvent = 0; iEvent < nEntries && !SignalHandler::IsTermRequested(); iEvent++ )
    {
      fDSTree.GetEntry( iEvent );
      if( iEvent == 0 )
        BeginOfRun( *runBranch );
      else if( dsBranch->GetRunID() != runBranch->GetRunID() )
        {
          fMainBlock->EndOfRun( *runBranch );
          //Get the next run with same RunID than the DS entry
          while(dsBranch->GetRunID() != runBranch->GetRunID() && ++iRun < fRunTree.GetEntries()){
            fRunTree.GetEntry( iRun );
          }
          Log::Assert( dsBranch->GetRunID() == runBranch->GetRunID(), "InROOTProducer::ReadEvents: Missing run or run out of order in run tree, corrupt file." );
          BeginOfRun( *runBranch );
        }
      // Create new data cleaning flags for this pass
      const size_t pass = MetaInformation::Get()->GetCurrentPass();
      for( size_t iEV = 0; iEV < dsBranch->GetEVCount(); iEV++ )
        {
          dsBranch->GetEV( iEV ).GetDataCleaningFlags().SetFlags( pass, DS::BitMask() );
          dsBranch->GetEV( iEV ).GetDataCleaningFlags().SetApplied( pass, DS::BitMask() );
        }
      // Now run the processors
      fMainBlock->DSEvent( *runBranch, *dsBranch );
    }
  // Finished, end the run
  fMainBlock->EndOfRun( *runBranch );

  // Cleanup
  delete dsBranch;
  delete runBranch;
}

void
InROOTProducer::BeginOfRun( DS::Run& run )
{
  // Create new data quality flags for this pass
  const size_t pass = MetaInformation::Get()->GetCurrentPass();
  run.GetDataQualityFlags().SetFlags( pass, DS::BitMask() );
  run.GetDataQualityFlags().SetApplied( pass, DS::BitMask() );
  Producer::BeginOfRun( run );
  fMainBlock->BeginOfRun( run );
}