#include "IRooTrackerFile.hxx"
#include "IRooTrackerHistogram.hxx"
#include <ICOMETLog.hxx>
#include <EoaCore.hxx>

#include <TFile.h>
#include <TDirectory.h>
#include <THnSparse.h>
#include <TH1.h>
#include <TH2.h>
#include <TRandom3.h>
#include <TMath.h>

#include <sstream>
#include <cstdio>
#include <algorithm>
#include <string>
#include <sstream>

bool COMET::IOutputRooTrackerHistogram::InitialiseForWriting(const char* mode){
    // Open the output file
    fFile=TFile::Open(fFilename.c_str(),mode);
    if(!fFile || fFile->IsZombie()){
        COMETError("Problem opening output file: "<<fFilename);
        return false;
    }

    return true;
}

void COMET::IOutputRooTrackerHistogram::BuildDefaultAxes(){
    fAxisProperties[kPosX]=AxisProperties(  1e3,-1e3,1e4,"X Position (mm)"       );
    fAxisProperties[kPosY]=AxisProperties(  1e3,-1e3,1e3,"Y Position (mm)"       );
    fAxisProperties[kPosZ]=AxisProperties(  1e3,-1e3,1e4,"Z Position (mm)"       );
    fAxisProperties[kMomX]=AxisProperties(  500,0,1e3,"Momentum in X (MeV/c)" );
    fAxisProperties[kMomY]=AxisProperties(  500,0,1e3,"Momentum in Y (MeV/c)" );
    fAxisProperties[kMomZ]=AxisProperties(  500,0,1e3,"Momentum in Z (MeV/c)" );
    fAxisProperties[kTime]=AxisProperties(  1e4,   0,1e5,"Time (ns)"             );
    fAxisProperties[kEnergy]=AxisProperties(1e4,   0,1e4,"Energy (MeV)"           );
}

bool COMET::IOutputRooTrackerHistogram::SetupAxis(const std::string& axis_specification){
    char axis[10];
    int bins=0;
    double low=0, high=0;
    int read=sscanf(axis_specification.c_str(),"%s =(%d,%lg,%lg)",axis,&bins,&low,&high);
    if(read!=4){
        COMETError("Bad axis specification: '"<<axis_specification<<"' should be of the form: axis =(bins,low,high)" );
        return false;
    }

    Dimensions_t axis_id=DimensionStr(axis);
    if(axis_id==kNumDimensions){
        COMETError("Axis name is not valid: "<<axis);
        return false;
    }

    fAxisProperties[axis_id].bins=bins;
    fAxisProperties[axis_id].low=low;
    fAxisProperties[axis_id].high=high;

    return true;
}

void COMET::IOutputRooTrackerHistogram::AddParticle(int pid,
           const TLorentzVector& position,
           const TLorentzVector& momentum,
           double weight){
    ++fParticlesInThisEvent[pid];

    // Obtain or build the necessary histograms for this pid
    THnBase* histogram=GetElseMakeDistribution(pid);

    // Fill with distribution of particle properties
    double vals[kNumDimensions];
    vals[kPosX]=position.X();
    vals[kPosY]=position.Y();
    vals[kPosZ]=position.Z();
    vals[kMomX]=momentum.Px();
    vals[kMomY]=momentum.Py();
    vals[kMomZ]=momentum.Pz();
    vals[kTime]=position.T();
    vals[kEnergy]=momentum.E();

    histogram->Fill(vals,weight);
}

THnBase* COMET::IOutputRooTrackerHistogram::GetElseMakeDistribution(int pid){
    // Do we already have a histogram for this PID?
    ParticleDistributionHistograms::iterator i_hist=fParticleDistributions.find(pid);
    if(i_hist!=fParticleDistributions.end()){
        return i_hist->second;
    }

    // Need to make a new histogram
    Int_t bins[kNumDimensions];
    Double_t low[kNumDimensions], high[kNumDimensions];
    std::stringstream title,name;
    name<<"hDistribution_"<<pid;
    title<<"Properties of "<<pid;
    for(int i=0; i<kNumDimensions;++i){
        const AxisProperties& axis=fAxisProperties.at(static_cast<Dimensions_t>(i));
        bins[i]=axis.bins;
        low[i]=axis.low;
        high[i]=axis.high;
        title<<";"<<axis.title;
    }
    THnBase* hist=new THnSparseF(name.str().c_str(),title.str().c_str(), 
                                 kNumDimensions, bins, low, high);
    fParticleDistributions[pid]=hist;
    return hist;
}

void COMET::IOutputRooTrackerHistogram::EndEvent(){
    if(fParticlesInThisEvent.empty()) return;
    ++fEventsIncluded;

    // transfer the particles found in this event to the particle frequency
    // histogram
    for(ParticleTally::const_iterator i_particle=fParticlesInThisEvent.begin();
            i_particle!=fParticlesInThisEvent.end(); ++i_particle){
        TH1* histogram=GetElseMakeFrequency(i_particle->first);
        // Ensure this frequency is not outside the range of maximum frequencies
        int tally = i_particle->second;
        bool isOverflow = histogram->IsBinOverflow(histogram->FindFixBin(tally));
        if(isOverflow) throw COMET::ETooManyParticles();
        // Otherwise, fill the histogram
        histogram->Fill(i_particle->second);
    }
}

TH1* COMET::IOutputRooTrackerHistogram::GetElseMakeFrequency(int pid){
    // Do we already have a histogram for this PID?
    ParticleFrequencyHistograms::iterator i_hist=fParticleFrequencies.find(pid);
    if(i_hist!=fParticleFrequencies.end()){
        return i_hist->second;
    }

    // Need to make a new histogram
    std::stringstream name,title;
    name<<"hFrequency_"<<pid;
    title<<"Number of pdg="<<pid<<" particles in a single event";
    TH1* hist=new TH1F(name.str().c_str(),title.str().c_str(),
                       IRooTrackerFile::kNPmax+1, 0,
                       IRooTrackerFile::kNPmax+1);
    fParticleFrequencies[pid]=hist;
    return hist;
}

void COMET::IOutputRooTrackerHistogram::SetFreqTotalEvents(const int totalEvents){
    // Work out for each particle type how many events had none of this particle
    for(ParticleFrequencyHistograms::iterator 
            i_hist=fParticleFrequencies.begin();
            i_hist!=fParticleFrequencies.end(); 
            ++i_hist){
        TH1* hist=i_hist->second;
        int pid=i_hist->first;
        // Integrate from 1 particle per event upwards
        int low=hist->FindFixBin(1);
        int nFilled=hist->Integral(low,-1);
        // Fill with the number of events that didn't contain this particle type
        int binContent = totalEvents - nFilled;
        if (binContent < 0){
            // Print an error
            COMETNamedError("IRooTrackerHistogram",
                            "Negative bin content attempted for zero bin of"
                            " frequency histogram for" <<
                            "\n  Particle ID " << pid <<
                            "\n  Filled Events " << nFilled <<
                            "\n  Total Events " << totalEvents <<
                            "\n  Attempted Empty Events " << binContent);
            // Throw an error
            throw ENegativeBinContent();
        }
        // Set the zero bin content
        int zero=hist->FindFixBin(0);
        hist->SetBinContent(zero, binContent);
        // Compute the integral for safety in using SetBinContent
        double integral = hist->ComputeIntegral();
    }
    // Set the total number of events
    fTotalEvents = totalEvents;
    COMETLog("Total Events: "<<GetNumTotalEvents());
}

void COMET::IOutputRooTrackerHistogram::Finalize(const int eventsRead){
    // Print the contents
    for(ParticleFrequencyHistograms::iterator 
            i_hist=fParticleFrequencies.begin();
            i_hist!=fParticleFrequencies.end(); 
            ++i_hist){
        // Integrate from 1 particle per event upwards
        int low=i_hist->second->FindFixBin(1);
        int nFilled=i_hist->second->Integral(low,-1);
        COMETLog("Events with "<<i_hist->first<<" = "<<nFilled);
    }
    // Add in enough empty events
    SetFreqTotalEvents(eventsRead);
}

void COMET::IOutputRooTrackerHistogram::Close(){
    TDirectory* oldDir=TDirectory::CurrentDirectory();
    fFile->cd();
    fFile->Close();
    oldDir->cd();
}

void COMET::IOutputRooTrackerHistogram::WriteHistogramsToFile(TFile* aFile){
    TDirectory* oldDir=TDirectory::CurrentDirectory();
    aFile->cd();
    for(ParticleFrequencyHistograms::iterator i_hist=fParticleFrequencies.begin();
            i_hist!=fParticleFrequencies.end(); ++i_hist){
        i_hist->second->Write();
    }
    for(ParticleDistributionHistograms::iterator i_hist=fParticleDistributions.begin();
            i_hist!=fParticleDistributions.end(); ++i_hist){
        i_hist->second->Write();
    }
    oldDir->cd();
}


bool COMET::IOutputRooTrackerHistogram::Add(const COMET::IRooTrackerHistogram& hist,
                                            double scale,
                                            bool skipCheck){
    // If we are checking before adding, ensure the two RooTracker histograms
    // can be added
    bool cannotAdd = (!skipCheck) &&
                     (!COMET::IOutputRooTrackerHistogram::CanAdd(*this, hist));
    if(cannotAdd){
        throw ECannotAddHistograms();
        return false;
    }
    // Get all the PIDs mentioned in the histogram to add
    std::set<int> pidToAdd = hist.GetDescribedParticles();
    // Get all the PIDs that this RooTracker histogram has
    std::set<int> pidHave = GetDescribedParticles();
    // Iteate through the PIDs to add and add tem for both the frequency
    // histograms and distribution histograms
    for(std::set<int>::iterator aPID = pidToAdd.begin();
                                aPID != pidToAdd.end();
                                aPID++){
        // Check if we have this one or not
        bool havePID = pidHave.find(*aPID) != pidHave.end();
        // If it is not present, then just clone the one that we want to add
        if(!havePID) {
            fParticleFrequencies[*aPID] =
                (TH1*)hist.GetFrequencyHist(*aPID)->Clone();
            fParticleFrequencies[*aPID]->Scale(scale);
            fParticleDistributions[*aPID] =
                (THnBase*)hist.GetDistributionHist(*aPID)->Clone();
            fParticleDistributions[*aPID]->Scale(scale);
        } else {
            // If it is present, then just add it
            fParticleFrequencies[*aPID]->Add(hist.GetFrequencyHist(*aPID), scale);
            fParticleDistributions[*aPID]->Add(hist.GetDistributionHist(*aPID), scale);
        }
    }
    return true;
}

bool COMET::IOutputRooTrackerHistogram::CanAdd(const IRooTrackerHistogram& hA,
                                               const IRooTrackerHistogram& hB){
    // Check that both histograms have particles
    std::set<int> pidA = hA.GetDescribedParticles();
    std::set<int> pidB = hB.GetDescribedParticles();
    if(pidB.size() == 0){
        COMETNamedError("IRooTrackerHistogram", "Cannot add an empty RooTracker "
                        "histogram to an empty RooTracker histogram");
        return false;
    }
    // Return true if hA is empty
    if (pidA.size() == 0) {
        COMETNamedWarn("IRooTrackerHistogram", "Adding histogram to entirely"
                       " empty histogram");
        return true;
    }
    // Check that the frequency histogram binning is compatable
    // TODO get ride potential memory leak problem
    if (!CanAdd(hA.GetFrequencyHist(*pidA.begin()),
                hB.GetFrequencyHist(*pidB.begin()))){
        COMETNamedError("IRooTrackerHistogram", "Cannot add a RooTracker "
                        "histograms: Frequency histogram has different binning");
        return false;
    }
    // Check that the sparse array histogram projections yeild the same binning
    // as well
    const THnBase* distA = hA.GetDistributionHist(*pidA.begin());
    const THnBase* distB = hB.GetDistributionHist(*pidB.begin());
    for (int dim = 0; dim < kNumDimensions; dim++){
        Dimensions_t thisDim = static_cast<Dimensions_t>(dim);
        TH1* aProj = distA->Projection(thisDim);
        aProj->SetName("distA");
        TH1* bProj = distB->Projection(thisDim);
        bProj->SetName("distB");
        if (!CanAdd(aProj, bProj)){
            COMETNamedError("IRooTrackerHistogram", "Cannot add a RooTracker "
                            "histograms: Binning for dimension " <<
                            DimensionStr(thisDim) << " is different");
            return false;
        }
    }
    // Otherwise, return true
    return true;
}

bool COMET::IOutputRooTrackerHistogram::CanAdd(const TH1* hA,
                                               const TH1* hB){
    // Ensure they have the same dimensions
    if(hA->GetDimension() != hB->GetDimension()) {
        COMETNamedError("IRooTrackerHistogram", "Cannot add histograms "
                        "of different dimensions");
        return false;
    }
    // Ensure they are both only one dimension
    Int_t dim = hA->GetDimension();
    if(dim != 1){
        COMETNamedError("IRooTrackerHistogram", "CanAdd(TH1*, TH1*) only valid "
                        "for 1D histograms");
        return false;
    }

    // Check whether the histograms have the same number of bins.
    if(hA->GetNbinsX() != hB->GetNbinsX()){
        COMETNamedError("IRooTrackerHistogram", "Cannot add histograms "
                        "of different number of bins");
        return false;
    }
    // Check axis limits
    if(!TMath::AreEqualRel(hA->GetXaxis()->GetXmin(),
                            hB->GetXaxis()->GetXmin(),
                            1.E-12) ||
        !TMath::AreEqualRel(hA->GetXaxis()->GetXmax(),
                            hB->GetXaxis()->GetXmax(),
                            1.E-12)){
        COMETNamedError("IRooTrackerHistogram", "Cannot add histograms "
                        "with different axis limits");
        return false;
    }

    // check bin limits
    const TArrayD* h1Array = hA->GetXaxis()->GetXbins();
    const TArrayD* h2Array = hB->GetXaxis()->GetXbins();
    Int_t fN = h1Array->fN;
    if(fN != 0) {
        bool returnVal = true;
        if(h2Array->fN != fN) returnVal = false;
        else{
            for(int i = 0; i < fN; ++i) {
                if(! TMath::AreEqualRel( h1Array->GetAt(i), h2Array->GetAt(i), 1E-10 ) ) {
                    returnVal = false;
                }
            }
        }
        if(returnVal == false){
            COMETNamedError("IRooTrackerHistogram", "Cannot add histograms "
                            "with different axis binning");
            return false;
        }
    }

    // Otherwise, return true!
    return true;
}



std::string COMET::IRooTrackerHistogram::DimensionStr(Dimensions_t dim){
    switch(dim){
        case kPosX:  return "X";
        case kPosY:  return "Y";
        case kPosZ:  return "Z";
        case kMomX:  return "Px";
        case kMomY:  return "Py";
        case kMomZ:  return "Pz";
        case kTime:  return "T";
        case kEnergy:return "E";
        case kNumDimensions:return "unknown";
    }
    return "unknown";
}

COMET::IRooTrackerHistogram::Dimensions_t COMET::IRooTrackerHistogram::DimensionStr(const std::string& dim){
    // Make the string lower case before comparison
    std::string dimLow = dim;
    std::transform(dimLow.begin(), dimLow.end(), dimLow.begin(), ::tolower);
         if(dimLow=="x"  )return kPosX  ;
    else if(dimLow=="y"  )return kPosY  ;
    else if(dimLow=="z"  )return kPosZ  ;
    else if(dimLow=="px" )return kMomX  ;
    else if(dimLow=="py" )return kMomY  ;
    else if(dimLow=="pz" )return kMomZ  ;
    else if(dimLow=="t"  )return kTime  ;
    else if(dimLow=="e"  )return kEnergy;
    return kNumDimensions;
}

const COMET::IRooTrackerHistogram::PrimaryList& COMET::IRooTrackerHistogram::GeneratePrimaries()
{
    // Clear out the existing primaries
    fGeneratedPrimaries.clear();

    double properties[kNumDimensions];

    for(ParticleFrequencyHistograms::iterator i_hist=fParticleFrequencies.begin();
            i_hist!=fParticleFrequencies.end(); ++i_hist){
        // Distributions for this particle type
        const int pid=i_hist->first;
        const TH1* frequency_hist=i_hist->second;
        THnBase* properties_hist=fParticleDistributions.at(pid);
        // Work out how many of each particle type to generate
        const int n_particles=frequency_hist->GetRandom();
        // Assign each particle its 4-position and 4-momentum
        for(int i=0; i<n_particles;++i){
            // Sample randomly within a random bin
            properties_hist->GetRandom(properties, kTRUE);
            // Add the  particle to the returned primaries
            fGeneratedPrimaries.push_back(PrimaryParticle_t(
                pid,
                properties[kTime],   properties[kPosX], properties[kPosY], properties[kPosZ],
                properties[kEnergy], properties[kMomX], properties[kMomY], properties[kMomZ]));
        }
    }
    return fGeneratedPrimaries;
}

const COMET::IRooTrackerHistogram::PrimaryList& COMET::IRooTrackerHistogram::GenerateSpecificPrimaries(){
    if(fSpecificTracks.size() < 1){
        COMETError("No requirement is set.");
        throw ECannotGeneratePrimary();
    }

    // Clear out the existing primaries
    fGeneratedPrimaries.clear();
    double properties[kNumDimensions];
    ParticleFrequencyHistograms::iterator i_hist;
    THnBase* properties_hist;
    int pid = 0;
    int n_particles;

    for(SpecificTracks::const_iterator spec = fSpecificTracks.begin();
                                       spec != fSpecificTracks.end();
                                       spec++){

        if(spec->IsSpecificPID()){
            // ifpid is specified, get frequency histo in advance.
            i_hist = fParticleFrequencies.find(spec->pid);
            if(i_hist == fParticleFrequencies.end()){
                COMETError("Cannot find the specific particle: " << spec->pid << ".");
                throw ECannotGeneratePrimary();
            }
            pid = spec->pid;
        }

        // (1) Get N-partcles
        do{
            if(not spec->IsSpecificPID()){
                // ifpid is not specified, randomly pick up a particle kind
                static TRandom3 random(0);
                int rndm = random.Uniform(fSpecificTracks.size());
                i_hist = fParticleFrequencies.begin();
                for(int j=0; j < rndm; j++) i_hist++;
                pid = i_hist->first;
            }

            // Generate 1 or multi track(s) and check its number
            n_particles = i_hist->second->GetRandom();
        }while(not spec->CheckNTracks(n_particles));

        // (2) Get track properties
        properties_hist=fParticleDistributions.at(pid);
        // Look for a set of the track properties satisfies the requirement
        for(int i=0; i<n_particles; ++i){
            int nFails = 0;
            do{
                // Too much retries; user must check the requirement
                if(nFails++ > 1000000){
                    COMETError("Cannot generate any track satisfying the requirement.");
                    throw ECannotGeneratePrimary();
                }
                // Generate a set oftrack properties and check it
                properties_hist->GetRandom(properties);
            }while(not spec->CheckTrackProperties(properties));

            // Add to primaries
            fGeneratedPrimaries.push_back(PrimaryParticle_t(
                        pid,
                        properties[kTime]   , properties[kPosX] , properties[kPosY] , properties[kPosZ],
                        properties[kEnergy] , properties[kMomX] , properties[kMomY] , properties[kMomZ]
                        ));
        }
    }
    return fGeneratedPrimaries;
}

COMET::IRooTrackerHistogram::PrimaryParticle_t COMET::IRooTrackerHistogram::GenerateOnePrimary(){
    if(fParticleWeights.empty()) InitializeParticleWeights();

    // Select a random particle to generate
    const double select=gRandom->Uniform();
    THnBase* distribution=NULL;
    int pid=0;
    for(ParticleWeights::const_iterator i_pid=fParticleWeights.begin();
            i_pid!=fParticleWeights.end() ; ++i_pid){
        if(i_pid->first>select){
            pid=i_pid->second;
            distribution=fParticleDistributions.at(pid);
            break;
        }
    }
    if(!distribution){
        COMETError("No distribution.  Weird behaviour...");
        throw ECannotGeneratePrimary();
    }
    double properties[kNumDimensions];
    distribution->GetRandom(properties);
    PrimaryParticle_t particle(
                pid,
                properties[kTime]   , properties[kPosX] , properties[kPosY] , properties[kPosZ],
                properties[kEnergy] , properties[kMomX] , properties[kMomY] , properties[kMomZ]
                );
    return particle;
}

void COMET::IRooTrackerHistogram::InitializeParticleWeights(){
    double cumulative=0;
    ParticleWeights tmp_weights;
    for(ParticleFrequencyHistograms::const_iterator i_hist=fParticleFrequencies.begin(); i_hist!=fParticleFrequencies.end(); ++i_hist){
        // Integrate from 1 particle per event upwards
        int low=i_hist->second->FindFixBin(1);
        const double integral=i_hist->second->Integral(low,-1);
        cumulative+=integral;
        tmp_weights[cumulative]=i_hist->first;
    }
    // Now scale everything so that the total is 1
    for(ParticleWeights::const_iterator i_weight=tmp_weights.begin(); i_weight!=tmp_weights.end(); ++i_weight){
        fParticleWeights[i_weight->first/cumulative]=i_weight->second;
    }
}

void COMET::IRooTrackerHistogram::PrimaryParticle_t::Print(const std::string& prefix)const{
    COMETLog(prefix<<pid<<", pos: ("<<posX<<", "<<posY<<", "<<posZ<<"), t="<<time<<", 4-mom: ("<<momX<<", "<<momY<<", "<<momZ<<", E="<<energy);
}

COMET::IRooTrackerHistogram* COMET::IRooTrackerHistogram::Open(const std::string& filename){
    TDirectory* old_dir=TDirectory::CurrentDirectory();

    TFile* input_file = TFile::Open(filename.c_str(),"READ");
    if(!input_file || !input_file->IsOpen()) {
        old_dir->cd();
        COMETError("Cannot find input file for RooTrackHist: '"<<filename<<"'");
        return NULL;
    }

    COMETLog("Obtain RooTrackHist from "<<filename);

    // Get the IRooTrackerHistogram object in the file:
    TObject* obj=input_file->Get(IRooTrackerHistogram::GetNameOnFile().c_str());
    if(!obj){
        old_dir->cd();
        COMETError("Cannot find IRooTrackerHistogram object (called '"<<IRooTrackerHistogram::GetNameOnFile()<<"') in file: "<<filename);
        return NULL;
    }
    old_dir->cd();

    IRooTrackerHistogram* returnVal = dynamic_cast<IRooTrackerHistogram*>(obj);
    // Close the file that was used
    input_file->Close();
    return returnVal;
}

void COMET::IRooTrackerHistogram::PrintParticles()const{
    for(ParticleFrequencyHistograms::const_iterator i_hist=fParticleFrequencies.begin(); i_hist!=fParticleFrequencies.end(); ++i_hist){
        // Integrate from 1 particle per event upwards
        int low=i_hist->second->FindFixBin(1);
        const double integral=i_hist->second->Integral(low,-1);
        COMETLog("particle: "<<i_hist->first<<" rel. rate: "<<integral);
    }
}

TDirectory* COMET::IRooTrackerHistogram::GetElseMakeDirectory(const std::string& name, 
                                                              TDirectory* parentDir){
    if(!parentDir){
        parentDir=TDirectory::CurrentDirectory();
    }
    TDirectory* dir=parentDir->GetDirectory(name.c_str());
    if(!dir) dir=parentDir->mkdir(name.c_str());
    return dir;
}

void COMET::IRooTrackerHistogram::Draw(TFile* aFile, bool only1DHistos){
    // Check if the file is okay
    if(!aFile || aFile->IsZombie()){
        COMETError("Problem opening output file for drawing");
        throw ECannotDrawHistogram();
    }
    // Cache the old directory we are in
    TDirectory* oldDir=TDirectory::CurrentDirectory();
    // Change directory into the file to save the output in
    aFile->cd();
    // Get the top level plots directory
    TDirectory* plotDir = GetElseMakeDirectory("Plots_All", aFile);
    plotDir->cd();
    // Get all the PIDs in this file
    std::set<int> partsToDraw = GetDescribedParticles();
    // Plot all information for each particle
    for (std::set<int>::const_iterator part = partsToDraw.begin();
                                       part != partsToDraw.end();
                                       part++){
        // Make the directory name for this PID
        std::stringstream pidName;
        pidName << "Plots_"<< *part;
        std::string pidStr = pidName.str();
        TDirectory* pidDir = GetElseMakeDirectory(pidStr.c_str(), plotDir);
        pidDir->cd();
        // Draw all plots for this PID
        DrawByPID(*part, only1DHistos);
    }
    oldDir->cd();
}

void COMET::IRooTrackerHistogram::DrawByPID(int pid, bool only1DHistos){
    // Write all the particle frequency histograms
    const TH1* thisFreq = GetFrequencyHist(pid);
    thisFreq->Write(thisFreq->GetName(), TObject::kOverwrite);
    // Get the projections of each dimension
    const THnBase* thisDist = GetDistributionHist(pid);
    for(int i=0; i<kNumDimensions;++i){
        // Get the projection
        TH1* aProj = thisDist->Projection(i);
        // Name the projection
        std::string projStr = thisDist->GetName();
        projStr += "_"+DimensionStr(static_cast<Dimensions_t>(i));
        // Write the projection
        aProj->Write(projStr.c_str(), TObject::kOverwrite);
    }
    // Return now if we do not want to draw 2D histograms as well
    if (only1DHistos) return;
    // Define all needed 2D projections
    typedef std::vector<std::pair<Dimensions_t, Dimensions_t> > ProjectionAxes;
    ProjectionAxes projAxes2D;
    // Get 2D projections in XY, XZ, YZ in space
    projAxes2D.push_back(std::make_pair(kPosX, kPosY));
    projAxes2D.push_back(std::make_pair(kPosX, kPosZ));
    projAxes2D.push_back(std::make_pair(kPosZ, kPosY));
    // Get 2D projections in XY, XZ, YZ in momentum
    projAxes2D.push_back(std::make_pair(kMomX, kMomY));
    projAxes2D.push_back(std::make_pair(kMomX, kMomZ));
    projAxes2D.push_back(std::make_pair(kMomZ, kMomY));
    // Get a 2D correlation plot of energy vs. time
    projAxes2D.push_back(std::make_pair(kTime, kEnergy));
    // Iterate through all defined pairs of 2D projections
    for(ProjectionAxes::iterator axes = projAxes2D.begin();
                                 axes != projAxes2D.end(); 
                                 axes++){
        // Get the 2D projection
        TH2* aProj2D = thisDist->Projection(axes->second, axes->first);
        // Name the 2D projection
        std::string projStr = thisDist->GetName();
        projStr += "_"+DimensionStr(static_cast<Dimensions_t>(axes->first));
        projStr += DimensionStr(static_cast<Dimensions_t>(axes->second));
        // Write the 2D projection
        aProj2D->Write(projStr.c_str(), TObject::kOverwrite);
    }
}

std::set<int> COMET::IRooTrackerHistogram::GetDescribedParticles()const{
    std::set<int> particles;
    for(ParticleFrequencyHistograms::const_iterator i_particle=fParticleFrequencies.begin();
            i_particle!=fParticleFrequencies.end(); ++i_particle){
        particles.insert(i_particle->first);
    }
    return particles;
}

const TH1* COMET::IRooTrackerHistogram::GetFrequencyHist(int pid)const{
    ParticleFrequencyHistograms::const_iterator i_hist=fParticleFrequencies.find(pid);
    if(i_hist==fParticleFrequencies.end()) return NULL;
    return i_hist->second;
}

const THnBase* COMET::IRooTrackerHistogram::GetDistributionHist(int pid)const{
    ParticleDistributionHistograms::const_iterator i_hist=fParticleDistributions.find(pid);
    if(i_hist==fParticleDistributions.end()) return NULL;
    return i_hist->second;
}