#include "CbmMCDataManager.h"

#include "CbmMCDataArray.h"
#include "CbmMCDataObject.h"

#include "FairRootManager.h"
#include "FairLogger.h"
#include "CbmFileSource.h"
#include "FairFileHeader.h"
#include "FairFileInfo.h"

#include <iostream>

using namespace std;

// --- Default constructor
CbmMCDataManager::CbmMCDataManager()
  : CbmMCDataManager("MCDataManager", 0)
//  : FairTask("MCDataManager", 0), fLegacy(0), fFileList(), fActive(), fActiveObj(), fFriends(), fbReadChainFromFirstFileHeader(kFALSE)
{
  fFileList.clear();
  fActive.clear();
  fActiveObj.clear();
  fFriends.clear();
}


// --- Standard constructor
CbmMCDataManager::CbmMCDataManager(const char* name, Int_t legacy)
  : FairTask(name,0), fLegacy(legacy), fFileList(), fActive(), fActiveObj(), fFriends(), fbReadChainFromFirstFileHeader(kFALSE)
{
  fFileList.clear();
  fActive.clear();
  fActiveObj.clear();
  fFriends.clear();
}


// --- Destructor
CbmMCDataManager::~CbmMCDataManager()
{
  map<TString, CbmMCDataArray*>::const_iterator p;
  for(p=fActive.begin();p!=fActive.end();++p) delete p->second;
  fActive.clear();

  map<TString, CbmMCDataObject*>::const_iterator o;
  for(o=fActiveObj.begin();o!=fActiveObj.end();++o) delete o->second;
  fActiveObj.clear();
}


// --- Add an additional input file
Int_t CbmMCDataManager::AddFile(const char* name)
{
  Int_t n=fFileList.size();
  fFileList.resize(n+1);
  fFileList[n].clear();
  fFileList[n].push_back(name);
  return n;
}


// --- Add a file to a given input chain
Int_t CbmMCDataManager::AddFileToChain(const char* name, Int_t number)
{
  Int_t i;
  Int_t n=fFileList.size();

  if (number<0) return -1111;
  if (number>=n)
  {
    fFileList.resize(number+1);
    for(i=n;i<=number;i++) fFileList[i].clear();
  }
  fFileList[number].push_back(name);

  return fFileList[number].size();
}


void CbmMCDataManager::ReadChainFromFirstFileHeader(Bool_t bReadFirstHeader)
{
  fbReadChainFromFirstFileHeader = bReadFirstHeader;
}


// --- End of event action
void CbmMCDataManager::FinishEvent()
{
  map<TString, CbmMCDataArray*>::const_iterator p;
  map<TString, CbmMCDataObject*>::const_iterator o;

  for(p=fActive.begin();p!=fActive.end();++p)
    p->second->FinishEvent();
  for(o=fActiveObj.begin();o!=fActiveObj.end();++o)
    o->second->FinishEvent();
}


// --- Initialisation
InitStatus CbmMCDataManager::Init()
{
  FairRootManager* fManager=FairRootManager::Instance();
  if (!fManager) return kFATAL;
  
  fManager->Register("MCDataManager", "Stack", (TNamed*)this, kFALSE);
  
  // Get access to the FairFileHeader object in the input file and construct
  // an input chain from all files referenced therein if requested
  if(fbReadChainFromFirstFileHeader)
  {
    if(0 < fFileList.size())
    {
      LOG(ERROR)<<"Chain 0 foreseen for FairFileHeader files already existing."<<FairLogger::endl;
      return kFATAL;
    }

    CbmFileSource* fileSource = dynamic_cast<CbmFileSource*>(fManager->GetSource());
    if(fileSource)
    {
      TFile* firstInputFile = fileSource->GetInFile();

      FairFileHeader* fileHeader(NULL);
      firstInputFile->GetObject("FileHeader", fileHeader);

      if(!fileHeader)
      {
        LOG(ERROR)<<"Could not retrieve FairFileHeader object from input file."<<FairLogger::endl;
        return kFATAL;
      }

      FairFileInfo* fileInfo(NULL);
      Int_t iFileIndex(0);
      Bool_t bExpansionSucceeded(kTRUE);

      while(fileHeader->GetFileInfo(0, iFileIndex))
      {
        fileInfo = fileHeader->GetFileInfo(0, iFileIndex);

        AddFileToChain((fileInfo->GetPath()).Data(), 0);

        if(0 == iFileIndex)
        {
          bExpansionSucceeded = ExpandFileHeaders(fileInfo->GetPath());
        }

        iFileIndex++;
      }

      if(!bExpansionSucceeded)
      {
        LOG(ERROR)<<"Could not read FairFileHeader objects recursively."<<FairLogger::endl;
        return kFATAL;
      }
    }
    else
    {
      LOG(ERROR)<<"Could not get CbmFileSource from FairRootManager."<<FairLogger::endl;
      return kFATAL;
    }
  }

  return kSUCCESS;
}


// --- Instantiate a data branch
CbmMCDataArray* CbmMCDataManager::InitBranch(const char* brname)
{
  CbmMCDataArray* arr = NULL;
  TString nm=brname;
  map<Int_t, Int_t>::const_iterator p;

  if (fActive.find(nm)!=fActive.end())
  {
    LOG(INFO) << "InitBranch: " << nm << " " << fActive[nm] << FairLogger::endl; 
    return fActive[nm];
  }  
  if (fLegacy==0)
  {
    arr=new CbmMCDataArray(brname, fFileList);
    for(p=fFriends.begin();p!=fFriends.end();++p)
    {
      arr->AddFriend(p->first, p->second);
    }
    fActive[nm]=arr;
  } //? Standard mode
  else
  {
  	if ( FairRootManager::Instance()->GetObject(brname) != NULL )
  	{
  	  arr=new CbmMCDataArray(brname);
 	  fActive[nm]=arr;
 	  LOG(INFO) << "InitBranch: " << nm << " " << arr << FairLogger::endl; 
  	} //? Branch found
  	else
  	{
  	  LOG(INFO) << "InitBranch: " << nm << " could not be initialised."
  	  << FairLogger::endl;
  	} //? Branch not found in FairRootManager
  } //? Legacy mode
  	
  return arr;
}

// --- Instantiate a data branch containing TObject
CbmMCDataObject* CbmMCDataManager::GetObject(const char* brname)
{
  CbmMCDataObject* arr;
  TString nm=brname;
  map<Int_t, Int_t>::const_iterator p;

  if (fActiveObj.find(nm)!=fActiveObj.end())
  {
    LOG(INFO) << "InitBranch: " << nm << " " << fActiveObj[nm] << FairLogger::endl; 
    return fActiveObj[nm];
  }  
  if (fLegacy==0)
  {
    arr=new CbmMCDataObject(brname, fFileList);
    for(p=fFriends.begin();p!=fFriends.end();++p)
    {
      arr->AddFriend(p->first, p->second);
    }
  }
  else
    arr=new CbmMCDataObject(brname);

  fActiveObj[nm]=arr;
  LOG(INFO) << "InitBranch: " << nm << " " << arr << FairLogger::endl; 
  return arr;
}


Bool_t CbmMCDataManager::ExpandFileHeaders(const TString& tFileName)
{
  TFile* tFile = new TFile(tFileName.Data(), "READ");

  FairFileHeader* fileHeader(NULL);
  tFile->GetObject("FileHeader", fileHeader);

  if(!fileHeader)
  {
    LOG(ERROR)<<Form("Could not retrieve FairFileHeader object from file %s.", tFileName.Data())<<FairLogger::endl;
    return kFALSE;
  }

  FairFileInfo* fileInfo(NULL);
  Int_t iFileIndex(0);
  Int_t iFileChainIndex(fFileList.size());
  Bool_t bExpansionSucceeded(kTRUE);

  while(fileHeader->GetFileInfo(0, iFileIndex))
  {
    fileInfo = fileHeader->GetFileInfo(0, iFileIndex);

    AddFileToChain((fileInfo->GetPath()).Data(), iFileChainIndex);

    if(0 == iFileIndex)
    {
      bExpansionSucceeded = ExpandFileHeaders(fileInfo->GetPath());
    }

    iFileIndex++;
  }

  tFile->Close();
  delete tFile;

  return kTRUE && bExpansionSucceeded;
}

ClassImp(CbmMCDataManager)