#include <TFile.h>
#include <TTree.h>
#include <TObjArray.h>
#include <string>
#include <ICOMETLog.hxx>
#include "IRooTrackerFile.hxx"

using COMET::IHandle;
using COMET::IRooTrackerFile;

IRooTrackerFile::StringToRooTrackerMap IRooTrackerFile::fListOfFiles;

IHandle<IRooTrackerFile> IRooTrackerFile::Open(
        const std::string& filename,
        const std::string& mode,
        const std::string& treename){
    StringToRooTrackerMap::iterator i_file= fListOfFiles.find(filename);
    if (i_file!=fListOfFiles.end()) return i_file->second;
    IHandle<IRooTrackerFile> file(new IRooTrackerFile(filename,mode, treename));
    fListOfFiles[filename]=file;
    return file;
}

void IRooTrackerFile::CloseAll(){
    for(StringToRooTrackerMap::iterator i_file=fListOfFiles.begin();
            i_file!=fListOfFiles.end(); ++i_file){
        if(i_file->second) {
            i_file->second->Close();
        }
    }
    // IHandles should go out of scope when we clear them
    fListOfFiles.clear();
}

IRooTrackerFile::~IRooTrackerFile(){
    Close();
}

void IRooTrackerFile::Reset(){
  brEvtFlags = 0;
  brEvtCode  = 0;
  brEvtNum=0;
  brMaxTrajID=0;
  brEvtXSec=0;
  brEvtDXSec=0;
  brEvtWght=0;
  brEvtProb=0;
  brStdHepN=0;
  brEvtVtx[0]=brEvtVtx[1]=brEvtVtx[2]=brEvtVtx[3]=0;
  fShouldFill=false;
}

IRooTrackerFile::IRooTrackerFile(const std::string& filename,
        const std::string& mode, const std::string& treename){
    // Set all values to defaults
    Reset();

    // Get the current directory so we can cd back to it at the end
    TDirectory* oldDir=TDirectory::CurrentDirectory();
    fFile=TFile::Open(filename.c_str(), mode.c_str());
    if(!fFile){
        COMETError("Unable to find file '"<<filename<<"'!");
        throw ENoRooTrackerFile();
        return;
    }
    fIsOpen=true;
    fIsWriting=fFile->IsWritable();

    if(fIsWriting){
        fTree=new TTree(treename.c_str(),"RooTracker tree for primary particles");
        fTree->SetDirectory(fFile);
        SetupWritingBranches();
        fTree->SetAutoSave(100);
    }else{
        // Get the tree
        fTree = dynamic_cast<TTree*>(fFile->Get(treename.c_str()));
        if (!fTree) {
            COMETError("Unable to find tree with name '"<<treename<<"' in file '"<<filename<<"'");
            throw ENoRooTrackerTree();
            return;
        }
        // Get the list of monitor names and setup the index map
        TObjArray* monitorNames = dynamic_cast<TObjArray*>(fFile->Get("MonitorNames"));
        if(monitorNames){
            for(int i=0; i< monitorNames->GetEntries(); ++i){
                TObjString* name=(TObjString*) (*monitorNames)[i];
                AddMonitor(name->String().Data());
                COMETLog("Found monitor : " << name->String().Data() <<
                         " in input RT file");
            }
        }
        SetupReadingBranches();
    }

    // Change back to the original directory
    oldDir->cd();
}

void IRooTrackerFile::SetupWritingBranches(){
    fTree->Branch("EvtFlags",  "TBits", &brEvtFlags  );
    fTree->Branch("EvtCode",   "TObjString",   &brEvtCode      );
    fTree->Branch("EvtNum",       &brEvtNum       ,"EvtNum/I"                 );
    fTree->Branch("MaxTrajID",    &brMaxTrajID    ,"MaxTrajID/I"              );
    fTree->Branch("EvtXSec",      &brEvtXSec      ,"EvtXSec/D"                );
    fTree->Branch("EvtDXSec",     &brEvtDXSec     ,"EvtDXSec/D"               );
    fTree->Branch("EvtWght",      &brEvtWght      ,"EvtWght/D"                );
    fTree->Branch("EvtProb",      &brEvtProb      ,"EvtProb/D"                );
    fTree->Branch("EvtVtx",        brEvtVtx       ,"EvtVtx[4]/D"              );
    fTree->Branch("StdHepN",      &brStdHepN      ,"StdHepN/I"                );
    fTree->Branch("StdHepTrackId", brStdHepTrackId,"StdHepTrackId[StdHepN]/I" );
    fTree->Branch("StdHepPdg",     brStdHepPdg    ,"StdHepPdg[StdHepN]/I"     );
    fTree->Branch("StdHepStatus",  brStdHepStatus ,"StdHepStatus[StdHepN]/I"  );
    fTree->Branch("StdHepRescat",  brStdHepRescat ,"StdHepRescat[StdHepN]/I"  );
    fTree->Branch("StdHepX4",      brStdHepX4     ,"StdHepX4[StdHepN][4]/D"   );
    fTree->Branch("StdHepP4",      brStdHepP4     ,"StdHepP4[StdHepN][4]/D"   );
    fTree->Branch("StdHepPolz",    brStdHepPolz   ,"StdHepPolz[StdHepN][3]/D" );
    fTree->Branch("StdHepFd",      brStdHepFd     ,"StdHepFd[StdHepN]/I"      );
    fTree->Branch("StdHepLd",      brStdHepLd     ,"StdHepLd[StdHepN]/I"      );
    fTree->Branch("StdHepFm",      brStdHepFm     ,"StdHepFm[StdHepN]/I"      );
    fTree->Branch("StdHepLm",      brStdHepLm     ,"StdHepLm[StdHepN]/I"      );
    fTree->Branch("Weight",        brWeight       ,"Weight[StdHepN]/D"  );
    fTree->Branch("MonitorID",     brMonitorId    ,"MonitorID[StdHepN]/I"  );
}

void IRooTrackerFile::SetupReadingBranches(){
    fTree->SetBranchAddress("EvtFlags",     &brEvtFlags  );
    fTree->SetBranchAddress("EvtCode",      &brEvtCode      );
    fTree->SetBranchAddress("EvtNum",       &brEvtNum       );
    fTree->SetBranchAddress("MaxTrajID",    &brMaxTrajID    );
    fTree->SetBranchAddress("EvtXSec",      &brEvtXSec      );
    fTree->SetBranchAddress("EvtDXSec",     &brEvtDXSec     );
    fTree->SetBranchAddress("EvtWght",      &brEvtWght      );
    fTree->SetBranchAddress("EvtProb",      &brEvtProb      );
    fTree->SetBranchAddress("EvtVtx",        brEvtVtx       );
    fTree->SetBranchAddress("StdHepN",      &brStdHepN      );
    fTree->SetBranchAddress("StdHepTrackId", brStdHepTrackId);
    fTree->SetBranchAddress("StdHepPdg",     brStdHepPdg    );
    fTree->SetBranchAddress("StdHepStatus",  brStdHepStatus );
    fTree->SetBranchAddress("StdHepRescat",  brStdHepRescat );
    fTree->SetBranchAddress("StdHepX4",      brStdHepX4     );
    fTree->SetBranchAddress("StdHepP4",      brStdHepP4     );
    fTree->SetBranchAddress("StdHepPolz",    brStdHepPolz   );
    fTree->SetBranchAddress("StdHepFd",      brStdHepFd     );
    fTree->SetBranchAddress("StdHepLd",      brStdHepLd     );
    fTree->SetBranchAddress("StdHepFm",      brStdHepFm     );
    fTree->SetBranchAddress("StdHepLm",      brStdHepLm     );
    fTree->SetBranchAddress("Weight",        brWeight       );
    fTree->SetBranchAddress("MonitorID",     brMonitorId    );
}

void IRooTrackerFile::Close(){
    //guard against double Close
    if(!fIsOpen) return;
    COMETNamedLog("IRooTrackerFile","Closing file: '"<<fFile->GetName()<<"'");
    fIsOpen=false;
    TDirectory* oldDir=TDirectory::CurrentDirectory();
    if (fIsWriting) {
        fFile->cd();
        fTree->Write();
        TObjArray* monitorNames=new TObjArray();
        for(IntToStringMap::const_iterator i_name=fMonitorNames.begin();
                i_name!=fMonitorNames.end(); ++i_name){
            monitorNames->Add(new TObjString(i_name->c_str()));
        }
        monitorNames->Write("MonitorNames",TObject::kSingleKey);
    }
    fFile->Close();
    COMETNamedLog("IRooTrackerFile","Closed");
    oldDir->cd();
}

int IRooTrackerFile::Fill(){
    if(!fIsWriting){
        COMETError("Trying to fill a read-only tree");
        return 0;
    }
    int fill=0;
    if(fShouldFill){
        fill =fTree->Fill();
        brStdHepN=0;
        fShouldFill=false;
    } 
    return fill;
}

int IRooTrackerFile::GetEntry(int i){
    if(fIsWriting){
        COMETError("Trying to read a write-only tree");
        return 0;
    }
    return fTree->GetEntry(i);
}

int IRooTrackerFile::GetEntries(){
    int ret=0;
    if(fTree) ret= fTree->GetEntries();
    return ret;
}

int IRooTrackerFile::AddMonitor(const std::string& name){
    StringToIntMap::const_iterator it = fMonitorIds.find(name);
    if(it != fMonitorIds.end()){
        return it->second;
    }
    fMonitorNames.push_back(name);
    fMonitorIds[name]=fMonitorNames.size();
    return fMonitorNames.size();
}
void IRooTrackerFile::SetMonitorName(int index, const std::string& val){
    brMonitorName[index].Add(new TObjString(val.c_str()));
}

void IRooTrackerFile::cd(){
    fFile->cd();
}