BmnParticleEqualizer1D.cxx

Go to the documentation of this file.

// -------------------------------------------------------------------------
// -----                BmnParticleEqualizer1D source file               -----
// -----          Created 01/05/2024 by K. Gertsenberger               -----
// -------------------------------------------------------------------------
 
#include "BmnParticleEqualizer1D.h"
 
#include "BmnGlobalTrack.h"
#include "BmnMCInfoDst.h"
#include "BmnTrackMatch.h"
#include "CbmMCTrack.h"
#include "CbmStack.h"
#include "FairLogger.h"
#include "FairMCEventHeader.h"
#include "FairRootManager.h"
#include "FairTrackParam.h"
#include "Math/MinimizerOptions.h"
#include "TCanvas.h"
#include "TChain.h"
#include "TClonesArray.h"
#include "TDatabasePDG.h"
#include "TFile.h"
#include "TFitResult.h"
#include "THashList.h"
#include "TLegend.h"
#include "TMath.h"
#include "TParticlePDG.h"
#include "TRandom.h"
#include "TSystem.h"
#include "TTreeReader.h"
#include "TTreeReaderValue.h"
#include "TVectorD.h"
#include "TVirtualMC.h"
 
// name and size of the object with the momentum grid (min P, max P, interval count)
// written to the file with the sample histograms to verify the generator configuration while reading them
static const char* kMomentumGridName = "MomentumGrid";
static const int kMomentumGridSize = 3;
 
// prefix of the name of the object with the relative reconstruction efficiency of a particle type
// (a TVectorD with one element per momentum interval), see ProduceEfficiencyMap
static const char* kEfficiencyNamePrefix = "Efficiency_pdg";
 
// binning of the Phi (azimuth angle) sample histograms of the momentum intervals. The whole azimuth range is covered
// by a small number of bins on purpose: an interval has much less statistics than the whole momentum range
static const int kIntervalPhiBinCount = 32;
 
///< Cut applied to the tracks of the sample (see SetSampleTrackCut)
Int_t BmnParticleEqualizer1D::fSampleTrackCut = 2;
Double_t BmnParticleEqualizer1D::fPMin = 0.4;   ///< Min value of momentum
Double_t BmnParticleEqualizer1D::fPMax = 5.4;   ///< Max value of momentum
///< Interval count for momentum to equalize particles at all the intervals
Int_t BmnParticleEqualizer1D::fIntervalCount = 25;
// particle PDG codes to be equalized:           p,   pi+,  pi-, K+,   K-      ,  e-,  e+
vector<Int_t> BmnParticleEqualizer1D::fParticles({2212, 211, -211, 321, -321});   //, 11, -11});
// their masses in MeV:                      938.272,139.6,139.6,493.7,493.7
 
// ranges of the sample histograms; they are internal to the translation unit ('static'), otherwise these common names
// would be exported by the library and could collide with the same ones of another translation unit
static Double_t fPtMin = 0, fPtMax = 3, fYMin = 0, fYMax = 3, fEtaMin = -4, fEtaMax = 8, fThetaMin = 0, fThetaMax = 3.2,
                fPhiMin = -3.2, fPhiMax = 3.2;
 
// Default 'stParticleInfo' constructor for storing particle distrubutons
BmnParticleEqualizer1D::stParticleInfo::stParticleInfo(int pdg_key,
                                                       Double_t min_p,
                                                       Double_t max_p,
                                                       Int_t interval_count)
    : fY0(0)
    , fSigma(0.72)
    , fT(0.223)
{
    // construct and initialize vector for particle numbers
    fParticleNumber.resize(interval_count);
    fill(fParticleNumber.begin(), fParticleNumber.end(), 0);
 
    // nothing has been added yet and the efficiency is not taken into account until its map is read
    fParticleDebt.resize(interval_count, 0.);
    fEfficiency.resize(interval_count, 1.);
 
    // construct vector with P histograms divided by momentum intervals
    fPHistVector.resize(interval_count);
    fEtaHistVector.resize(interval_count);
    fPhiHistVector.resize(interval_count);
 
    // the way to sample the values is defined after the sample histograms have been read
    fPSampling.resize(interval_count, EPSampling::kNone);
    fEtaHistIndex.resize(interval_count, -1);
    fPhiHistIndex.resize(interval_count, -1);
}
 
// 'BmnParticleEqualizer1D' constructor with input histograms
BmnParticleEqualizer1D::BmnParticleEqualizer1D(TString hist_file_name)
    : FairGenerator()
    , fIntervalStep((fPMax - fPMin) / fIntervalCount)
{
    fUseSampleStatisticsCut = kTRUE;
    fTargetPerInterval = 0.;
    fUseEfficiencyCorrection = kFALSE;
    fMinEfficiency = 0.05;
    fIsEfficiencyMapRead = kFALSE;
 
    fMaxParticles.resize(fIntervalCount);
    fill(fMaxParticles.begin(), fMaxParticles.end(), 0);
 
    // initialize map with particle information
    for (auto pdg_key : fParticles)
        fParticleInfos.emplace(pdg_key, make_unique<stParticleInfo>(pdg_key, fPMin, fPMax, fIntervalCount));
 
    isActive = ReadSampleHistograms(hist_file_name);
    if (!isActive)
        LOG(error) << "Errors occurred while initializing the BmnParticleEqualizer1D generator, "
                      "the simulation run will be stopped";
}
 
// Constructor with specified particle types
BmnParticleEqualizer1D::BmnParticleEqualizer1D(TString hist_file_name,
                                               vector<Int_t> pdg_codes,
                                               Double_t min_p,
                                               Double_t max_p,
                                               Int_t intervals)
    : FairGenerator()
{
    fPMin = min_p;
    fPMax = max_p;
    fIntervalCount = intervals;
    fIntervalStep = (max_p - min_p) / intervals;
 
    // set fParticles to given pdg_codes, dropping the duplicates: the map of the particle information keeps one
    // entry per PDG code, while the loops over 'fParticles' would process a duplicated code (and add its particles
    // to the event) as many times as it is listed
    fParticles.clear();
    for (auto pdg_key : pdg_codes) {
        if (find(fParticles.begin(), fParticles.end(), pdg_key) != fParticles.end()) {
            LOG(warning) << "The PDG code " << pdg_key
                         << " is listed more than once for the equalizing generator, "
                            "the extra entries are ignored";
            continue;
        }
        fParticles.push_back(pdg_key);
    }
 
    fUseSampleStatisticsCut = kTRUE;
    fTargetPerInterval = 0.;
    fUseEfficiencyCorrection = kFALSE;
    fMinEfficiency = 0.05;
    fIsEfficiencyMapRead = kFALSE;
 
    fMaxParticles.resize(fIntervalCount);
    fill(fMaxParticles.begin(), fMaxParticles.end(), 0);
 
    // initialize map with particle information
    for (auto pdg_key : fParticles)
        fParticleInfos.emplace(pdg_key, make_unique<stParticleInfo>(pdg_key, fPMin, fPMax, fIntervalCount));
 
    isActive = ReadSampleHistograms(hist_file_name);
    if (!isActive)
        LOG(error) << "Errors occurred while initializing the BmnParticleEqualizer1D generator, "
                      "the simulation run will be stopped";
}
 
// Initialize generator
Bool_t BmnParticleEqualizer1D::Init()
{
    return kTRUE;
}
 
// fit the full momentum histogram of a particle type and attach the function to it (as 'func_name'). The bulk of the
// distribution is described by two Gaussians and a Landau distribution, with the limits and the start values of the
// parameters defined by the histogram itself. The power law tail of the Landau term is much flatter than the steeply
// falling tails of the measured distributions, so above the momentum where the histogram falls below a small fraction
// of its peak ('p_tail') the bulk is replaced by an exponential fitted to the tail
// return kTRUE if the fits are valid
static Bool_t FitMomentumFunction(TH1D* hP, const TString& func_name, Double_t p_min, Double_t p_max)
{
    TF1* p_distribution =
        new TF1(func_name, "(x<[9])*(gaus(0)+gaus(3)+landau(6))+(x>=[9])*[10]*TMath::Exp(-[11]*(x-[9]))", p_min, p_max);
    Double_t p_avg = (p_min + p_max) / 2;
    Int_t max_bin = hP->GetMaximumBin();
    Double_t x_peak = hP->GetXaxis()->GetBinCenter(max_bin);
    Double_t x_max_value = hP->GetXaxis()->GetBinUpEdge(max_bin);
    Double_t y_peak = hP->GetBinContent(max_bin);
    // the spread of the distribution gives the scale of the start widths
    Double_t rms = TMath::Max(hP->GetRMS(), 0.2);
 
    // Every start value is put strictly inside the limits of its parameter. Minuit maps a limited parameter onto an
    // unlimited one, and the derivative of that mapping is zero exactly at a limit, so a parameter started there
    // cannot leave it and the fit ends as invalid
    auto set_parameter = [p_distribution](Int_t index, Double_t min_value, Double_t max_value, Double_t start_value) {
        Double_t margin = 0.01 * (max_value - min_value);
        p_distribution->SetParLimits(index, min_value, max_value);
        p_distribution->SetParameter(index, TMath::Range(min_value + margin, max_value - margin, start_value));
    };
 
    // The amplitudes are kept non-negative: the function is used to sample the momentum, so it must not be negative
    // anywhere, while terms of opposite signs compensating each other make it negative in the tails. The amplitude of
    // the Landau term is not its height: TMath::Landau() is about 0.18 at the most probable value
    set_parameter(0, 0., 10. * y_peak, 0.3 * y_peak);
    set_parameter(1, p_min, p_avg, x_peak - 0.5 * rms);
    set_parameter(2, 0.1, p_avg, 0.5 * rms);
    set_parameter(3, 0., 10. * y_peak, 0.5 * y_peak);
    set_parameter(4, x_max_value - 0.1, p_max, x_peak + 0.5 * rms);
    set_parameter(5, 0.1, p_avg, rms);
    set_parameter(6, 0., 100. * y_peak, 2. * y_peak);
    set_parameter(7, p_min, TMath::Min(x_max_value, p_avg), x_peak - 0.1);
    set_parameter(8, 0.1, p_avg, 0.3 * rms);
 
    // the start of the tail: the first bin to the right of the peak below the given fraction of the peak height. The
    // tail needs a few bins and some statistics to be fitted, otherwise the bulk covers the whole range (the step is
    // then put beyond the range, so that the function is the bulk one up to its upper limit inclusive)
    const Double_t tail_fraction = 0.05, min_tail_entries = 30;
    Double_t p_tail = p_max;
    for (Int_t bin = max_bin + 1; bin <= hP->GetNbinsX(); bin++) {
        if (hP->GetBinContent(bin) < tail_fraction * y_peak) {
            p_tail = hP->GetXaxis()->GetBinLowEdge(bin);
            break;
        }
    }
    Bool_t is_tail = (p_tail <= p_max - 3 * hP->GetBinWidth(max_bin))
                     && (hP->Integral(hP->FindBin(p_tail + 1.e-9), hP->GetNbinsX()) >= min_tail_entries);
    if (!is_tail)
        p_tail = p_max + 1.;
    p_distribution->FixParameter(9, p_tail);
    p_distribution->FixParameter(10, 0.);
    p_distribution->FixParameter(11, 0.);
 
    // The binned likelihood ('L') is used instead of the chi-square: the latter ignores the empty bins and is biased in
    // the bins with a few entries, i.e. exactly in the tails of the distribution this function is needed for
    // 1) the bulk below the tail
    p_distribution->SetRange(p_min, TMath::Min(p_tail, p_max));
    Bool_t is_valid = hP->Fit(p_distribution, "SQ0LR")->IsValid();
    // 2) the exponential tail with its own amplitude and slope. The amplitude is not tied to the value of the bulk at
    // 'p_tail': the Landau term already deviates from the histogram there, and the tail would inherit the deviation
    if (is_tail) {
        for (Int_t index = 0; index <= 8; index++)
            p_distribution->FixParameter(index, p_distribution->GetParameter(index));
        // the value of the bulk at the start of the tail (the function is still the bulk one just below the step)
        Double_t bulk_value = p_distribution->Eval(p_tail * (1. - 1.e-12));
        p_distribution->ReleaseParameter(10);
        set_parameter(10, 0., 100. * TMath::Max(bulk_value, 1.), bulk_value);
        p_distribution->ReleaseParameter(11);
        set_parameter(11, 0., 50., 2.);
        p_distribution->SetRange(p_tail, p_max);
        is_valid = hP->Fit(p_distribution, "SQ0LR")->IsValid() && is_valid;
    }
    // the fits are made in the subranges, while the function is used over the whole momentum range
    p_distribution->SetRange(p_min, p_max);
    TF1* fitted_function = hP->GetFunction(func_name);
    if (fitted_function != nullptr)
        fitted_function->SetRange(p_min, p_max);
 
    return is_valid;
}
 
// fit P histogram to generate the values according to the analytical functions
int BmnParticleEqualizer1D::FitPHistogram(unique_ptr<TH1D>& hP, int pdg_key, Double_t p_min, Double_t p_max)
{
    if (hP->GetEntries() == 0) {
        LOG(warning) << TString::Format("P (momentum) histogram is empty for PDG = %d", pdg_key);
        return -1;
    }
 
    // fit P distribution for the particle type by gaus + landau with an exponential tail
    if (!FitMomentumFunction(hP.get(), TString::Format("distP_pdg%d", pdg_key), p_min, p_max)) {
        LOG(error) << TString::Format("Fit of the P (momentum) histogram failed for PDG = %d", pdg_key);
        // return -2;
    }
 
    return 0;
}
 
// fit Pt histogram to generate the values according to the analytical functions
int BmnParticleEqualizer1D::FitPtHistogram(unique_ptr<TH1D>& hPt,
                                           int pdg_key,
                                           Double_t pdg_mass,
                                           Double_t pt_min,
                                           Double_t pt_max,
                                           Double_t& t_output)
{
    if (hPt->GetEntries() == 0) {
        LOG(warning) << TString::Format("Pt (transverse momentum) histogram is empty for PDG = %d", pdg_key);
        return -1;
    }
 
    // fit Pt distribution for the particle type by expo
    // the squared mass is formatted with a high precision on purpose: the default one (6 decimal digits) turns the
    // squared mass of the light particles into zero (2.6e-7 for e+/e-) and so removes the mass from the formula
    TString dist_func_name = TString::Format("distPt_pdg%d", pdg_key);
    TF1* pt_distribution = new TF1(
        dist_func_name, TString::Format("[0]*x*exp(-sqrt(x*x+%.10g)/[1])", pdg_mass * pdg_mass), pt_min, pt_max);
    pt_distribution->SetParameter(1, t_output);
    // the temperature is a positive value by definition, an unlimited parameter can be driven negative by the fit
    pt_distribution->SetParLimits(1, 1.e-3, 1.);
 
    TFitResultPtr r = hPt->Fit(pt_distribution, "SQ0");
    if (!r->IsValid()) {
        LOG(error) << TString::Format(
            "Fit of the Pt (transverse momentum) histogram failed for PDG = %d, the previous temperature value "
            "(%6.4f) is kept",
            pdg_key, t_output);
        return 0;
    }
    t_output = r->Parameter(1);
 
    return 0;
}
 
// fit Y histogram to generate the values according to the analytical functions
int BmnParticleEqualizer1D::FitYHistogram(unique_ptr<TH1D>& hY,
                                          int pdg_key,
                                          Double_t& y0_output,
                                          Double_t& sigma_putput)
{
    if (hY->GetEntries() == 0) {
        LOG(warning) << TString::Format("Rapidity histogram is empty for %d PDG code", pdg_key);
        return -1;
    }
 
    // fit rapidity distribution for the particle type by gauss
    if (pdg_key != 2212) {
        TFitResultPtr r = hY->Fit("gaus", "SQ0");
        y0_output = r->Parameter(1);
        sigma_putput = r->Parameter(2);
        // if proton (PDG == 2212)
    } else {
        Double_t max_bin_value = hY->GetMaximum();
        const string minimizer_type = ROOT::Math::MinimizerOptions::DefaultMinimizerType();
        ROOT::Math::MinimizerOptions::SetDefaultMinimizer("Genetic");
 
        // FITTING ONLY MEDIUM PEAK with background
        // TF1* fProtonYDistribution = new TF1("distProtonY", "gaus(0)+pol0(3)", 0.3, 2);
        // fProtonYDistribution->SetParLimits(0, 0, max_bin_value);
        // fProtonYDistribution->SetParLimits(1, 0.3, 2);
        // fProtonYDistribution->SetParName(1,"proton_y0");
        // fProtonYDistribution->SetParLimits(2, 0.05, 1);
        // fProtonYDistribution->SetParName(2,"proton_sigma");
        // fProtonYDistribution->SetParLimits(3, 0, max_bin_value);
        // TFitResultPtr r = hY->Fit(fProtonYDistribution,"SQ0","", 0.3, 2);
        // FITTING THREE PEAKS: TARGET FRAGMENTATION + PARTICIPANTS + SPECTATORS
        unique_ptr<TF1> fProtonYDistribution =
            make_unique<TF1>("distProtonY", "gaus(0)+gaus(3)+gaus(6)", 0.2, 5);   // 0, 5);
        fProtonYDistribution->SetParLimits(0, 0, max_bin_value);
        fProtonYDistribution->SetParLimits(1, 0, 0.05);
        fProtonYDistribution->SetParLimits(2, 0.01, 0.4);
        fProtonYDistribution->SetParameter(2, 0.2);
        fProtonYDistribution->SetParLimits(3, 0, max_bin_value);
        fProtonYDistribution->SetParName(4, "proton_y0");
        fProtonYDistribution->SetParLimits(4, 0.3, 2);
        fProtonYDistribution->SetParameter(4, 1);
        fProtonYDistribution->SetParName(5, "proton_sigma");
        fProtonYDistribution->SetParLimits(5, 0.05, 1);
        fProtonYDistribution->SetParameter(5, 0.5);
        fProtonYDistribution->SetParLimits(6, 0, max_bin_value);
        fProtonYDistribution->SetParLimits(7, 1.5, 3);
        fProtonYDistribution->SetParameter(7, 2.25);
        fProtonYDistribution->SetParLimits(8, 0.05, 1);
        fProtonYDistribution->SetParameter(8, 0.5);
 
        TFitResultPtr r = hY->Fit(fProtonYDistribution.get(), "SQ0", "", 0, 5);
        y0_output = r->Parameter(r->Index("proton_y0"));
        sigma_putput = r->Parameter(r->Index("proton_sigma"));
 
        ROOT::Math::MinimizerOptions::SetDefaultMinimizer(minimizer_type.c_str());
    }
 
    return 0;
}
 
// fit Eta (pseudorapidity) histogram to generate the values according to the analytical functions
int BmnParticleEqualizer1D::FitEtaHistogram(unique_ptr<TH1D>& hEta, int pdg_key, Double_t eta_min, Double_t eta_max)
{
    if (hEta->GetEntries() == 0) {
        LOG(warning) << TString::Format("Eta (pseudorapidity) histogram is empty for PDG = %d", pdg_key);
        return -1;
    }
 
    // fit P distribution for the particle type by gaus + landau
    TString dist_func_name = TString::Format("distEta_pdg%d", pdg_key);
    TF1* eta_distribution = new TF1(dist_func_name, "crystalball", eta_min, eta_max);
    eta_distribution->SetParLimits(1, 1, 4);
    eta_distribution->SetParameter(1, 2.5);
    eta_distribution->SetParLimits(2, 0.1, 1.5);
    eta_distribution->SetParameter(2, 0.5);
    eta_distribution->SetParLimits(3, 0.1, 2);
    eta_distribution->SetParameter(3, 1);
    eta_distribution->SetParLimits(4, 1, 10);
    eta_distribution->SetParameter(4, 3);
 
    TFitResultPtr r = hEta->Fit(eta_distribution, "SQ0");
    if (!r->IsValid()) {
        LOG(error) << TString::Format("Fit of the Eta (pseudorapidity) histogram failed for PDG = %d", pdg_key);
        // return -2;
    }
 
    return 0;
}
 
int BmnParticleEqualizer1D::ProduceSampleHistograms(TString input_list_file, TString output_histo_file)
{
    gSystem->ExpandPathName(input_list_file);
 
    // get file list from the input text file
    list<string> file_list;
    ifstream list_file(input_list_file.Data());
    string s;
    while (getline(list_file, s))
        file_list.push_back(s);
    if (file_list.empty()) {
        cout << "ERROR: There are no files in the list of the input text file: " << input_list_file << endl;
        return -14;
    }
 
    // form ROOT chain for the file list
    TChain dataChain("bmndata");
    for (auto const& cur_file_path : file_list) {
        TString strFilePath(cur_file_path);
        gSystem->ExpandPathName(strFilePath);
 
        dataChain.AddFile(strFilePath);
    }
 
    if (dataChain.GetEntries() < 1) {
        cout << "ERROR: There are no correct input files in the text list stored at: \"" << input_list_file << "\""
             << endl;
        return -15;
    }
 
    // flag whether the input files have SIM or DST format
    bool isDST = false;
    TClonesArray* mcTracks = nullptr;
    BmnMCInfoDst* mcInfo = nullptr;
    if (dataChain.GetBranch("MCTrack") == nullptr) {
        LOG(info) << "'MCTrack' branch was not found, searching for 'BmnMCInfo->MCTrack'...";
        if (dataChain.GetBranch("BmnMCInfo.") == nullptr) {
            LOG(error) << "'BmnMCInfo.' branch was not found too, please, check the structure of the files, exiting...";
            return -1;
        }
        dataChain.SetBranchAddress("BmnMCInfo.", &mcInfo);
        if (mcInfo == nullptr) {
            LOG(error)
                << "'BmnMCInfo.' branch was not read correctly, please, check the structure of the files, exiting...";
            return -2;
        }
        isDST = true;
        LOG(info) << "'BmnMCInfo->MCTrack' has been found";
        dataChain.SetBranchStatus("*", 0);
        dataChain.SetBranchStatus("BmnMCInfo*", 1);
    } else {
        dataChain.SetBranchAddress("MCTrack", &mcTracks);
        if (mcTracks == nullptr) {
            LOG(error)
                << "MCTrack' branch was not read correctly, please, check the structure of the files, exiting...";
            return -3;
        }
        dataChain.SetBranchStatus("*", 0);
        dataChain.SetBranchStatus("MCTrack*", 1);
    }
 
    int iEventNumber = dataChain.GetEntries();
    if (iEventNumber < 1) {
        cout << "ERROR: There are no correct entries in the file list stored at: \"" << input_list_file << "\"" << endl;
        return -16;
    }
 
    // the file is owned by the unique pointer so that it is closed and deleted on every exit from the function
    unique_ptr<TFile> hist_file(TFile::Open(output_histo_file, "RECREATE"));
    if ((!hist_file) || (!hist_file->IsOpen())) {
        cout << "ERROR: unable to create ROOT file for the sample histograms: " << output_histo_file << endl;
        return -1;
    }
 
    // create map with particle infos
    map<Int_t, unique_ptr<stParticleInfo>> particle_infos;
    for (auto pdg_key : fParticles)
        particle_infos.emplace(pdg_key, make_unique<stParticleInfo>(pdg_key, fPMin, fPMax, fIntervalCount));
 
    Double_t interval_width = (fPMax - fPMin) / fIntervalCount;
    // create all the histograms using make_unique
    for (auto pdg_key : fParticles) {
        TString histogram_name = TString::Format("hP_pdg%d", pdg_key);
        particle_infos[pdg_key]->fPHist = make_unique<TH1D>(histogram_name, histogram_name, 80, fPMin, fPMax);
        histogram_name = TString::Format("hPt_pdg%d", pdg_key);
        particle_infos[pdg_key]->fPtHist = make_unique<TH1D>(histogram_name, histogram_name, 60, fPtMin, fPtMax);
        histogram_name = TString::Format("hY_pdg%d", pdg_key);
        particle_infos[pdg_key]->fYHist = make_unique<TH1D>(histogram_name, histogram_name, 60, fYMin, fYMax);
        histogram_name = TString::Format("hEta_pdg%d", pdg_key);
        particle_infos[pdg_key]->fEtaHist = make_unique<TH1D>(histogram_name, histogram_name, 60, fEtaMin, fEtaMax);
        histogram_name = TString::Format("hTheta_pdg%d", pdg_key);
        particle_infos[pdg_key]->fThetaHist =
            make_unique<TH1D>(histogram_name, histogram_name, 61, fThetaMin, fThetaMax);
        histogram_name = TString::Format("hPhi_pdg%d", pdg_key);
        particle_infos[pdg_key]->fPhiHist = make_unique<TH1D>(histogram_name, histogram_name, 64, fPhiMin, fPhiMax);
        // detach the histograms from the output file (they are created while it is the current ROOT directory) so they
        // are owned only by the unique pointers - otherwise closing the file below would delete them and their
        // destructors would cause a double free
        particle_infos[pdg_key]->fPHist->SetDirectory(nullptr);
        particle_infos[pdg_key]->fPtHist->SetDirectory(nullptr);
        particle_infos[pdg_key]->fYHist->SetDirectory(nullptr);
        particle_infos[pdg_key]->fEtaHist->SetDirectory(nullptr);
        particle_infos[pdg_key]->fThetaHist->SetDirectory(nullptr);
        particle_infos[pdg_key]->fPhiHist->SetDirectory(nullptr);
        int hist_counter = 0;
        for (auto& p_hist : particle_infos[pdg_key]->fPHistVector) {
            histogram_name = TString::Format("hP%d_pdg%d", ++hist_counter, pdg_key);
            p_hist = make_unique<TH1D>(histogram_name, histogram_name, 40, fPMin + (hist_counter - 1) * interval_width,
                                       fPMin + hist_counter * interval_width);
            p_hist->SetDirectory(nullptr);
        }
        hist_counter = 0;
        for (auto& eta_hist : particle_infos[pdg_key]->fEtaHistVector) {
            histogram_name = TString::Format("hEta%d_pdg%d", ++hist_counter, pdg_key);
            eta_hist = make_unique<TH1D>(histogram_name, histogram_name, 60, fEtaMin, fEtaMax);
            eta_hist->SetDirectory(nullptr);
        }
        hist_counter = 0;
        for (auto& phi_hist : particle_infos[pdg_key]->fPhiHistVector) {
            histogram_name = TString::Format("hPhi%d_pdg%d", ++hist_counter, pdg_key);
            phi_hist = make_unique<TH1D>(histogram_name, histogram_name, kIntervalPhiBinCount, fPhiMin, fPhiMax);
            phi_hist->SetDirectory(nullptr);
        }
    }
 
    // fill all the histograms for specified particles by the PDG code
    int iCurrentEvent = 0;
    LOG(info) << "Processing events...";
    for (Int_t iEvent = 0; iEvent < iEventNumber; iEvent++) {
        // Next Event
        iCurrentEvent++;
        if (iCurrentEvent % 1000 == 0)
            LOG(info) << "Processing event #" << iCurrentEvent << " of " << iEventNumber;
        dataChain.GetEntry(iEvent);
 
        TClonesArray* mc_tracks = nullptr;
        if (isDST)
            mc_tracks = mcInfo->GetMCTracks();
        else
            mc_tracks = mcTracks;
 
        for (int i = 0; i < mc_tracks->GetEntries(); i++) {
            CbmMCTrack* pTrack = (CbmMCTrack*)mc_tracks->At(i);
            Int_t part_code = pTrack->GetPdgCode();
 
            // exclude secondary particles
            if (pTrack->GetMotherId() >= 0)
                continue;
            // filter particles
            if (fSampleTrackCut > 0) {
                // exclude tracks without points in Silicon and GEMs
                if ((pTrack->GetNPoints(kGEM) < 1) || (pTrack->GetNPoints(kSILICON) < 1))
                    continue;
                if (fSampleTrackCut > 1) {
                    // exclude tracks without points in ToF detectors. kTOF701 is not checked: CbmMCTrack keeps the
                    // points of TOF700 in a single bit field, returned by GetNPoints() both for kTOF and for kTOF701
                    if ((pTrack->GetNPoints(kTOF) < 1) && (pTrack->GetNPoints(kTOF1) < 1))
                        continue;
                }
            }
 
            if (auto part_iter = particle_infos.find(part_code); part_iter != particle_infos.end()) {
                int interval_number = (pTrack->GetP() - fPMin) / interval_width;
                Double_t track_p = pTrack->GetP();
                Double_t track_eta = 0.5 * TMath::Log((track_p + pTrack->GetPz()) / (track_p - pTrack->GetPz()));
                // the same definition of the azimuth angle has to be used while reading the histograms in ReadEvent()
                Double_t track_phi = TMath::ATan2(pTrack->GetPy(), pTrack->GetPx());
                part_iter->second->fPHist->Fill(track_p);
                part_iter->second->fPtHist->Fill(pTrack->GetPt());
                part_iter->second->fYHist->Fill(pTrack->GetRapidity());
                part_iter->second->fEtaHist->Fill(track_eta);
                part_iter->second->fThetaHist->Fill(TMath::ACos(pTrack->GetPz() / track_p));   // *TMath::RadToDeg()
                part_iter->second->fPhiHist->Fill(track_phi);                                  // *TMath::RadToDeg()
                // compare doubles: abs(min(a,b))*numeric_limits<double>::epsilon()*error_factor
                if ((track_p >= fPMin) && (isless(track_p, fPMax))) {
                    part_iter->second->fPHistVector[interval_number]->Fill(track_p);
                    part_iter->second->fEtaHistVector[interval_number]->Fill(track_eta);
                    part_iter->second->fPhiHistVector[interval_number]->Fill(track_phi);
                }
            }
        }   // for (int i = 0; i < mc_tracks->GetEntries(); i++) {
    }   // while (sim_reader.Next())
 
    // check correctness of all the histograms
    for (auto pdg_key : fParticles) {
        if (particle_infos[pdg_key]->fPHist->GetEntries() == 0) {
            LOG(error) << "Momentum histogram is empty for pdg=" << pdg_key;
            return -4;
        }
        if (particle_infos[pdg_key]->fPtHist->GetEntries() == 0) {
            LOG(error) << "Transverse momentum histogram is empty for pdg=" << pdg_key;
            return -5;
        }
        if (particle_infos[pdg_key]->fYHist->GetEntries() == 0) {
            LOG(error) << "Rapidity histogram is empty for pdg=" << pdg_key;
            return -6;
        }
        if (particle_infos[pdg_key]->fEtaHist->GetEntries() == 0) {
            LOG(error) << "Pseudorapidity histogram is empty for pdg=" << pdg_key;
            return -7;
        }
        if (particle_infos[pdg_key]->fThetaHist->GetEntries() == 0) {
            LOG(error) << "Theta angle histogram is empty for pdg=" << pdg_key;
            return -8;
        }
        if (particle_infos[pdg_key]->fPhiHist->GetEntries() == 0) {
            LOG(error) << "Phi angle histogram is empty for pdg=" << pdg_key;
            return -9;
        }
    }
 
    // fit histograms and calculate fY0, fSigma, fT for all the specified particle types
    for (auto pdg_key : fParticles) {
        // the start values are set inside the loop on purpose: being declared outside it, they would enter the fit of
        // every next particle type as the values fitted for the previous one, making the result depend on the order
        Double_t fY0 = 0, fSigma = 0.72, fT = 0.223;
        TDatabasePDG* pdgBase = TDatabasePDG::Instance();
        TParticlePDG* particle = pdgBase->GetParticle(pdg_key);
        if (!particle) {
            LOG(fatal) << "PDG code " << pdg_key << "is not defined. exiting...";
            return -10;
        }
        Double_t fPDGMass = particle->Mass();
 
        // fit full particle momentum histograms
        if (FitPHistogram(particle_infos[pdg_key]->fPHist, pdg_key, fPMin, fPMax) != 0)
            return -11;
 
        // fit Pt momentum histograms
        if (FitPtHistogram(particle_infos[pdg_key]->fPtHist, pdg_key, fPDGMass, fPtMin, fPtMax, fT) != 0)
            return -12;
 
        // fit Y (rapidity) histograms
        if (FitYHistogram(particle_infos[pdg_key]->fYHist, pdg_key, fY0, fSigma) != 0)
            return -13;
        TVector3 Y0SigmaT(fY0, fSigma, fT);
        TString vector_name = TString::Format("Y0SigmaT_pdg%d", pdg_key);
        hist_file->WriteObject(&Y0SigmaT, vector_name);
 
        // fit Eta (pseudorapidity) histograms
        // if (FitEtaHistogram(particle_infos[pdg_key]->fEtaHist, pdg_key, eta_min, eta_max) != 0)
        //    return -5;
 
        LOG(info) << TString::Format("pdg=%i y0=%4.2f sigma_y=%4.2f T_pt=%6.4f", pdg_key, fY0, fSigma, fT);
    }
 
    // write the momentum grid the histograms were produced with, so that a generator reading this file can verify
    // that its own grid is exactly the same: the interval histograms are read by name only and a shifted grid would
    // not be detected by itself (see ReadSampleHistograms)
    TVectorD momentum_grid(kMomentumGridSize);
    momentum_grid[0] = fPMin;
    momentum_grid[1] = fPMax;
    momentum_grid[2] = fIntervalCount;
    hist_file->WriteObject(&momentum_grid, kMomentumGridName);
 
    // write the histograms to the output ROOT file: they are detached from it, so the file is specified explicitly
    for (auto pdg_key : fParticles) {
        hist_file->WriteTObject(particle_infos[pdg_key]->fPHist.get());
        hist_file->WriteTObject(particle_infos[pdg_key]->fPtHist.get());
        hist_file->WriteTObject(particle_infos[pdg_key]->fYHist.get());
        hist_file->WriteTObject(particle_infos[pdg_key]->fEtaHist.get());
        hist_file->WriteTObject(particle_infos[pdg_key]->fThetaHist.get());
        hist_file->WriteTObject(particle_infos[pdg_key]->fPhiHist.get());
        for (auto& p_hist : particle_infos[pdg_key]->fPHistVector)
            hist_file->WriteTObject(p_hist.get());
        for (auto& eta_hist : particle_infos[pdg_key]->fEtaHistVector)
            hist_file->WriteTObject(eta_hist.get());
        for (auto& phi_hist : particle_infos[pdg_key]->fPhiHistVector)
            hist_file->WriteTObject(phi_hist.get());
        // hist_file->WriteObject(&particle_infos[pdg_key]->fPHistVector, "PHistVector");
        // hist_file->WriteObject(&particle_infos[pdg_key]->fEtaHistVector, "EtaHistVector");
    }
 
    hist_file->Close();
    return 0;
}
 
// switch on the correction of the number of the added particles by the reconstruction efficiency
Bool_t BmnParticleEqualizer1D::SetEfficiencyCorrection(Bool_t is_used, Double_t min_efficiency)
{
    if (is_used && (!fIsEfficiencyMapRead)) {
        LOG(error) << "The efficiency map is absent in the file with the sample histograms for at least one of the "
                      "particle types, so the correction by the efficiency cannot be switched on. Build the map with "
                      "BmnParticleEqualizer1D::ProduceEfficiencyMap() from the reconstructed (DST) files first";
        return kFALSE;
    }
    if (min_efficiency <= 0.) {
        LOG(error) << "The lower limit of the efficiency has to be positive, otherwise the correction factor is "
                      "unlimited; the correction is not switched on";
        return kFALSE;
    }
 
    fUseEfficiencyCorrection = is_used;
    fMinEfficiency = min_efficiency;
    return kTRUE;
}
 
// verify that the momentum grid of the sample histograms is the one this class is configured for
Bool_t BmnParticleEqualizer1D::CheckMomentumGrid(TFile* hist_file)
{
    if ((hist_file == nullptr) || (!hist_file->IsOpen())) {
        LOG(error) << "The file with the sample histograms is not open";
        return kFALSE;
    }
 
    unique_ptr<TVectorD> momentum_grid((TVectorD*)hist_file->Get(kMomentumGridName));
    if ((momentum_grid == nullptr) || (momentum_grid->GetNoElements() != kMomentumGridSize)) {
        LOG(error) << "The momentum grid ('" << kMomentumGridName
                   << "') is absent in the file with the sample histograms: " << hist_file->GetName()
                   << ". The file was produced by an older version of the equalizing generator, please, recreate it "
                      "with ProduceSampleHistograms()";
        return kFALSE;
    }
    // the momenta are compared with an absolute tolerance, because a zero value is a valid lower momentum limit
    if ((!TMath::AreEqualAbs((*momentum_grid)[0], fPMin, 1.e-9))
        || (!TMath::AreEqualAbs((*momentum_grid)[1], fPMax, 1.e-9))
        || (TMath::Nint((*momentum_grid)[2]) != fIntervalCount))
    {
        LOG(error) << TString::Format(
            "The sample histograms were produced for %d momentum intervals in [%g, %g] GeV, but the generator is "
            "configured for %d intervals in [%g, %g] GeV. Please, recreate the file '%s' with "
            "ProduceSampleHistograms() or configure the generator accordingly",
            TMath::Nint((*momentum_grid)[2]), (*momentum_grid)[0], (*momentum_grid)[1], fIntervalCount, fPMin, fPMax,
            hist_file->GetName());
        return kFALSE;
    }
 
    return kTRUE;
}
 
// build the map of the reconstruction efficiency from the reconstructed (DST) files
int BmnParticleEqualizer1D::ProduceEfficiencyMap(TString input_list_file, TString histo_file)
{
    gSystem->ExpandPathName(input_list_file);
 
    // get file list from the input text file
    list<string> file_list;
    ifstream list_file(input_list_file.Data());
    string s;
    while (getline(list_file, s))
        file_list.push_back(s);
    if (file_list.empty()) {
        LOG(error) << "There are no files in the list of the input text file: " << input_list_file;
        return -1;
    }
 
    // form ROOT chain for the file list
    TChain dataChain("bmndata");
    for (auto const& cur_file_path : file_list) {
        TString strFilePath(cur_file_path);
        gSystem->ExpandPathName(strFilePath);
        dataChain.AddFile(strFilePath);
    }
    if (dataChain.GetEntries() < 1) {
        LOG(error) << "There are no correct input files in the text list stored at: " << input_list_file;
        return -2;
    }
 
    // the sample histograms give the number of the generated tracks the efficiency is referred to, so their
    // momentum grid has to be verified before the (long) loop over the reconstructed files is started
    unique_ptr<TFile> hist_file(TFile::Open(histo_file, "UPDATE"));
    if ((!hist_file) || (!hist_file->IsOpen())) {
        LOG(error) << "Unable to open the file with the sample histograms for updating: " << histo_file;
        return -4;
    }
    if (!CheckMomentumGrid(hist_file.get()))
        return -7;
 
    // the reconstructed tracks are needed, so the input files have to be DST ones
    TClonesArray* globalTracks = nullptr;
    BmnMCInfoDst* mcInfo = nullptr;
    if (dataChain.GetBranch("BmnGlobalTrack") == nullptr) {
        LOG(error) << "'BmnGlobalTrack' branch was not found: the efficiency map is built from the reconstructed "
                      "(DST) files, please, check the structure of the given files";
        return -3;
    }
    dataChain.SetBranchAddress("BmnGlobalTrack", &globalTracks);
    // the MC information is present in the simulated files only: for the experimental ones the particle type is
    // defined by the mass measured by the ToF detectors. A reconstructed track is related to its MC track by the index
    // stored in the track itself (BmnGlobalTrack::GetRefIndex(), which always returns -1 for now) or by the separate
    // matching branch
    TClonesArray* globalMatches = nullptr;
    Bool_t isMC = (dataChain.GetBranch("BmnMCInfo.") != nullptr);
    if (isMC)
        dataChain.SetBranchAddress("BmnMCInfo.", &mcInfo);
    Bool_t isMatched = isMC && (dataChain.GetBranch("BmnGlobalTrackMatch") != nullptr);
    if (isMatched)
        dataChain.SetBranchAddress("BmnGlobalTrackMatch", &globalMatches);
    LOG(info) << "Building the efficiency map from the " << (isMC ? "simulated" : "experimental") << " DST files";
    if (isMC && (!isMatched))
        LOG(warning) << "'BmnGlobalTrackMatch' branch was not found, the particle types of the reconstructed tracks "
                        "without an MC index are defined by the mass measured by the ToF detectors";
 
    // count the reconstructed tracks of every particle type in every momentum interval
    map<Int_t, vector<Double_t>> reconstructed_count;
    for (auto pdg_key : fParticles)
        reconstructed_count[pdg_key].assign(fIntervalCount, 0.);
 
    Double_t interval_width = (fPMax - fPMin) / fIntervalCount;
    int iEventNumber = dataChain.GetEntries();
    for (int iEvent = 0; iEvent < iEventNumber; iEvent++) {
        if ((iEvent + 1) % 1000 == 0)
            LOG(info) << "Processing event #" << iEvent + 1 << " of " << iEventNumber;
        dataChain.GetEntry(iEvent);
        if (globalTracks == nullptr)
            continue;
        TClonesArray* mc_tracks = (isMC && (mcInfo != nullptr)) ? mcInfo->GetMCTracks() : nullptr;
 
        for (int iTrack = 0; iTrack < globalTracks->GetEntries(); iTrack++) {
            BmnGlobalTrack* global_track = (BmnGlobalTrack*)globalTracks->UncheckedAt(iTrack);
            if (global_track == nullptr)
                continue;
            // the reconstructed tracks are counted in the same acceptance as the sample ones (see SetSampleTrackCut):
            // with the default cut only the tracks matched to a ToF hit
            Bool_t has_tof = (global_track->GetTof1HitIndex() > -1) || (global_track->GetTof2HitIndex() > -1);
            if ((fSampleTrackCut > 1) && (!has_tof))
                continue;
 
            // define the particle type: by the related MC track for the simulated data, by the ToF mass and the charge
            // for the experimental ones
            BmnTrackMatch* global_match = ((globalMatches != nullptr) && (iTrack < globalMatches->GetEntriesFast()))
                                              ? (BmnTrackMatch*)globalMatches->UncheckedAt(iTrack)
                                              : nullptr;
            Double_t tof_mass = -1.;
            Int_t track_pdg = GetGlobalTrackPDG(global_track, global_match, mc_tracks, tof_mass);
            if (reconstructed_count.find(track_pdg) == reconstructed_count.end())
                continue;
 
            FairTrackParam* parFirst = global_track->GetParamFirst();
            if (parFirst == nullptr)
                continue;
            Double_t track_qp = parFirst->GetQp();
            if (TMath::Abs(track_qp) < 1.e-4)
                continue;
            Double_t track_p = 1. / TMath::Abs(track_qp);
            if ((track_p < fPMin) || (track_p >= fPMax))
                continue;
            int interval_number = (track_p - fPMin) / interval_width;
            if ((interval_number < 0) || (interval_number >= fIntervalCount))
                continue;
 
            reconstructed_count[track_pdg][interval_number] += 1.;
        }
    }
 
    for (auto pdg_key : fParticles) {
        TVectorD efficiency(fIntervalCount);
        Double_t max_efficiency = 0.;
        for (int cur_interval = 0; cur_interval < fIntervalCount; cur_interval++) {
            TString histogram_name = TString::Format("hP%d_pdg%d", cur_interval + 1, pdg_key);
            unique_ptr<TH1D> p_hist((TH1D*)hist_file->Get(histogram_name));
            if (p_hist == nullptr) {
                LOG(error) << "The momentum histogram '" << histogram_name
                           << "' does not exist in the file with the sample histograms: " << histo_file;
                return -5;
            }
            p_hist->SetDirectory(nullptr);
            Double_t sample_count = p_hist->GetEntries();
            efficiency[cur_interval] =
                (sample_count > 0) ? (reconstructed_count[pdg_key][cur_interval] / sample_count) : 0.;
            max_efficiency = TMath::Max(max_efficiency, efficiency[cur_interval]);
        }
 
        if (max_efficiency <= 0.) {
            LOG(error) << "No track of pdg=" << pdg_key << " was reconstructed in the given files, "
                       << "the efficiency map cannot be built";
            return -6;
        }
        // the numbers of the reconstructed and of the sample tracks come from different data sets, so only the shape
        // of the ratio is meaningful: it is normalized to its maximum, giving a RELATIVE efficiency
        efficiency *= 1. / max_efficiency;
        // 'overwrite' replaces the map of a previous run instead of adding one more cycle of the same key
        hist_file->WriteObject(&efficiency, TString::Format("%s%d", kEfficiencyNamePrefix, pdg_key), "overwrite");
 
        TString efficiency_values;
        for (int cur_interval = 0; cur_interval < fIntervalCount; cur_interval++)
            efficiency_values += TString::Format(" %4.2f", efficiency[cur_interval]);
        LOG(info) << "Relative efficiency for pdg=" << pdg_key << ":" << efficiency_values;
    }
 
    hist_file->Close();
    return 0;
}
 
// Calculate parameters for all the particle types
// according to the given histogram stored in the 'input_histo_file' file
Bool_t BmnParticleEqualizer1D::ReadSampleHistograms(TString input_histo_file)
{
    // open file with histograms
    TFile* hist_file = TFile::Open(input_histo_file);
    if ((hist_file == nullptr) || (hist_file->IsZombie())) {
        LOG(error) << "The file with the sample histograms was not opened: " << input_histo_file;
        return kFALSE;
    }
 
    // the interval histograms are read by name only, so a grid different from the one of this generator would
    // silently refer them to the wrong momentum ranges
    if (!CheckMomentumGrid(hist_file))
        return kFALSE;
 
    // get histograms from the file and estimate parameters for embedding new particles of this type
    // an empty list of the particle types would leave the flag set without a single map having been read
    Bool_t is_efficiency_map_read = (!fParticles.empty());
    for (auto pdg_key : fParticles) {
        TString histogram_name = TString::Format("hP_pdg%d", pdg_key);
        fParticleInfos[pdg_key]->fPHist = unique_ptr<TH1D>((TH1D*)hist_file->Get(histogram_name));
        if (fParticleInfos[pdg_key]->fPHist.get() == nullptr) {
            LOG(error) << "Momentum histogram does not exist for pdg=" << pdg_key;
            return kFALSE;
        }
        // detach the histogram from the file so it (and its fitted function) survives after the file is closed below,
        // and is owned only by the unique_ptr - otherwise closing the file would delete it and cause a double free
        fParticleInfos[pdg_key]->fPHist->SetDirectory(nullptr);
        if (fParticleInfos[pdg_key]->fPHist->GetEntries() == 0) {
            LOG(error) << "Momentum histogram is empty for pdg=" << pdg_key;
            return kFALSE;
        }
        // the fallback momentum sampling for intervals with too little statistics uses the fitted function of this full
        // histogram; if the histogram file was produced without the fit, create it now so the fallback always works
        if (fParticleInfos[pdg_key]->fPHist->GetFunction(TString::Format("distP_pdg%d", pdg_key)) == nullptr) {
            LOG(info) << "The fitted momentum function is absent for pdg=" << pdg_key << ", fitting it now";
            FitPHistogram(fParticleInfos[pdg_key]->fPHist, pdg_key, fPMin, fPMax);
        }
        histogram_name = TString::Format("hPt_pdg%d", pdg_key);
        fParticleInfos[pdg_key]->fPtHist = unique_ptr<TH1D>((TH1D*)hist_file->Get(histogram_name));
        if (fParticleInfos[pdg_key]->fPtHist.get() == nullptr) {
            LOG(error) << "Transverse momentum histogram does not exist for pdg=" << pdg_key;
            return kFALSE;
        }
        fParticleInfos[pdg_key]->fPtHist->SetDirectory(nullptr);
        if (fParticleInfos[pdg_key]->fPtHist->GetEntries() == 0) {
            LOG(error) << "Transverse momentum histogram is empty for pdg=" << pdg_key;
            return kFALSE;
        }
        histogram_name = TString::Format("hY_pdg%d", pdg_key);
        fParticleInfos[pdg_key]->fYHist = unique_ptr<TH1D>((TH1D*)hist_file->Get(histogram_name));
        if (fParticleInfos[pdg_key]->fYHist.get() == nullptr) {
            LOG(error) << "Rapidity histogram does not exist for pdg=" << pdg_key;
            return kFALSE;
        }
        fParticleInfos[pdg_key]->fYHist->SetDirectory(nullptr);
        if (fParticleInfos[pdg_key]->fYHist->GetEntries() == 0) {
            LOG(error) << "Rapidity histogram is empty for pdg=" << pdg_key;
            return kFALSE;
        }
        histogram_name = TString::Format("hEta_pdg%d", pdg_key);
        fParticleInfos[pdg_key]->fEtaHist = unique_ptr<TH1D>((TH1D*)hist_file->Get(histogram_name));
        if (fParticleInfos[pdg_key]->fEtaHist.get() == nullptr) {
            LOG(error) << "Pseudorapidity histogram does not exist for pdg=" << pdg_key;
            return kFALSE;
        }
        fParticleInfos[pdg_key]->fEtaHist->SetDirectory(nullptr);
        if (fParticleInfos[pdg_key]->fEtaHist->GetEntries() == 0) {
            LOG(error) << "Pseudorapidity histogram is empty for pdg=" << pdg_key;
            return kFALSE;
        }
        histogram_name = TString::Format("hTheta_pdg%d", pdg_key);
        fParticleInfos[pdg_key]->fThetaHist = unique_ptr<TH1D>((TH1D*)hist_file->Get(histogram_name));
        if (fParticleInfos[pdg_key]->fThetaHist.get() == nullptr) {
            LOG(error) << "Theta angle histogram does not exist for pdg=" << pdg_key;
            return kFALSE;
        }
        fParticleInfos[pdg_key]->fThetaHist->SetDirectory(nullptr);
        if (fParticleInfos[pdg_key]->fThetaHist->GetEntries() == 0) {
            LOG(error) << "Theta angle histogram is empty for pdg=" << pdg_key;
            return kFALSE;
        }
        histogram_name = TString::Format("hPhi_pdg%d", pdg_key);
        fParticleInfos[pdg_key]->fPhiHist = unique_ptr<TH1D>((TH1D*)hist_file->Get(histogram_name));
        if (fParticleInfos[pdg_key]->fPhiHist.get() == nullptr) {
            LOG(error) << "Phi angle histogram does not exist for pdg=" << pdg_key;
            return kFALSE;
        }
        fParticleInfos[pdg_key]->fPhiHist->SetDirectory(nullptr);
        if (fParticleInfos[pdg_key]->fPhiHist->GetEntries() == 0) {
            LOG(error) << "Phi angle histogram is empty for pdg=" << pdg_key;
            return kFALSE;
        }
        int hist_counter = 0;
        for (auto& p_hist : fParticleInfos[pdg_key]->fPHistVector) {
            histogram_name = TString::Format("hP%d_pdg%d", ++hist_counter, pdg_key);
            p_hist = unique_ptr<TH1D>((TH1D*)hist_file->Get(histogram_name));
            if (p_hist.get() == nullptr) {
                LOG(error) << "Momentum histogram does not exist for pdg=" << pdg_key
                           << " in interval=" << hist_counter;
                return kFALSE;
            }
            p_hist->SetDirectory(nullptr);
        }
        hist_counter = 0;
        for (auto& eta_hist : fParticleInfos[pdg_key]->fEtaHistVector) {
            histogram_name = TString::Format("hEta%d_pdg%d", ++hist_counter, pdg_key);
            eta_hist = unique_ptr<TH1D>((TH1D*)hist_file->Get(histogram_name));
            if (eta_hist.get() == nullptr) {
                LOG(error) << "Pseudorapidity histogram does not exist for pdg=" << pdg_key
                           << " in interval=" << hist_counter;
                return kFALSE;
            }
            eta_hist->SetDirectory(nullptr);
        }
        hist_counter = 0;
        for (auto& phi_hist : fParticleInfos[pdg_key]->fPhiHistVector) {
            histogram_name = TString::Format("hPhi%d_pdg%d", ++hist_counter, pdg_key);
            phi_hist = unique_ptr<TH1D>((TH1D*)hist_file->Get(histogram_name));
            if (phi_hist.get() == nullptr) {
                LOG(error) << "Phi angle histogram does not exist for pdg=" << pdg_key
                           << " in interval=" << hist_counter
                           << ". The file with the sample histograms was produced by an older version of the "
                              "equalizing generator, please, recreate it with ProduceSampleHistograms()";
                return kFALSE;
            }
            phi_hist->SetDirectory(nullptr);
        }
 
        // the pseudorapidity and the azimuth angle of a new particle are taken from the histograms of its momentum
        // interval, or from the nearest interval having statistics if the interval itself has none: this way the values
        // stay tied to a close momentum region instead of being taken from the whole momentum range
        DefineNearestFilledIntervals(fParticleInfos[pdg_key]->fEtaHistVector, fParticleInfos[pdg_key]->fEtaHistIndex);
        DefineNearestFilledIntervals(fParticleInfos[pdg_key]->fPhiHistVector, fParticleInfos[pdg_key]->fPhiHistIndex);
 
        // define how the momentum of a new particle will be sampled in every momentum interval: the P histogram of the
        // interval itself is used if the sample has enough statistics there, otherwise the fitted function of the full
        // P histogram is used as a fallback. It does not depend on the event, so it is defined once here
        TF1* p_dist_func = fParticleInfos[pdg_key]->fPHist->GetFunction(TString::Format("distP_pdg%d", pdg_key));
        if (p_dist_func != nullptr) {
            // TF1::GetRandom() draws a value from the cumulative distribution tabulated on a grid of GetNpx() points
            // over the whole function range and then rejects the values outside the requested momentum interval, so
            // the grid has to be fine enough for a single interval to keep the sampling unbiased and fast
            p_dist_func->SetNpx(TMath::Max(p_dist_func->GetNpx(), 40 * fIntervalCount));
        }
        // The lowest value of the fitted function an interval can be sampled at. A function negligible relative to its
        // maximum is as unusable as a non-positive one: the cumulative distribution TF1::GetRandom() tabulates cannot
        // resolve such an interval in double precision and undefined values are returned (it was observed below a
        // relative level of about 1e-13). It is the case for the tails of a function fitted to a small sample
        Double_t min_func_value =
            (p_dist_func != nullptr) ? TMath::Max(0., 1.e-9 * p_dist_func->GetMaximum(fPMin, fPMax)) : 0.;
        Double_t interval_p = fPMin;
        for (int cur_interval = 0; cur_interval < fIntervalCount; cur_interval++, interval_p += fIntervalStep) {
            if (fParticleInfos[pdg_key]->fPHistVector[cur_interval]->GetEntries() > 2)
                fParticleInfos[pdg_key]->fPSampling[cur_interval] = stParticleInfo::EPSampling::kIntervalHist;
            // the rejection sampling inside TF1::GetRandom() never returns if the fitted function has no positive
            // values in the requested interval, so such intervals have to be excluded here
            else if ((p_dist_func != nullptr)
                     && (p_dist_func->GetMaximum(interval_p, interval_p + fIntervalStep) > min_func_value))
                fParticleInfos[pdg_key]->fPSampling[cur_interval] = stParticleInfo::EPSampling::kFitFunction;
            else {
                fParticleInfos[pdg_key]->fPSampling[cur_interval] = stParticleInfo::EPSampling::kNone;
                LOG(warning) << TString::Format(
                    "There is neither sample statistics nor a significant fitted momentum function for pdg=%d in the "
                    "momentum interval [%4.2f, %4.2f] GeV: no particles will be added there",
                    pdg_key, interval_p, interval_p + fIntervalStep);
            }
        }
 
        if (TDatabasePDG::Instance()->GetParticle(pdg_key) == nullptr) {
            LOG(error) << "PDG code " << pdg_key << " is not defined in the PDG database";
            return kFALSE;
        }
 
        // the mean rapidity, its sigma and the temperature of the Pt distribution were obtained by fitting the sample
        // histograms while producing them, so they are simply read back here instead of repeating the fits at every
        // start of the simulation (the rapidity fit of the protons uses the slow "Genetic" minimizer)
        TString vector_name = TString::Format("Y0SigmaT_pdg%d", pdg_key);
        unique_ptr<TVector3> y0_sigma_t((TVector3*)hist_file->Get(vector_name));
        if (y0_sigma_t == nullptr) {
            LOG(error) << "The fitted distribution parameters ('" << vector_name
                       << "') do not exist in the file with the sample histograms. It was produced by an older "
                          "version of the equalizing generator, please, recreate it with ProduceSampleHistograms()";
            return kFALSE;
        }
        fParticleInfos[pdg_key]->fY0 = y0_sigma_t->X();
        fParticleInfos[pdg_key]->fSigma = y0_sigma_t->Y();
        fParticleInfos[pdg_key]->fT = y0_sigma_t->Z();
 
        LOG(info) << TString::Format("pdg=%i y0=%4.2f sigma_y=%4.2f T_pt=%6.4f", pdg_key, fParticleInfos[pdg_key]->fY0,
                                     fParticleInfos[pdg_key]->fSigma, fParticleInfos[pdg_key]->fT);
 
        // the map of the reconstruction efficiency is optional: it is only used if the correction by it is switched
        // on by SetEfficiencyCorrection(), which refuses to do it when the map has not been read here
        unique_ptr<TVectorD> efficiency(
            (TVectorD*)hist_file->Get(TString::Format("%s%d", kEfficiencyNamePrefix, pdg_key)));
        if ((efficiency != nullptr) && (efficiency->GetNoElements() == fIntervalCount)) {
            for (int cur_interval = 0; cur_interval < fIntervalCount; cur_interval++)
                fParticleInfos[pdg_key]->fEfficiency[cur_interval] = (*efficiency)[cur_interval];
        } else {
            if (efficiency != nullptr)
                LOG(error) << "The efficiency map of pdg=" << pdg_key << " has " << efficiency->GetNoElements()
                           << " elements instead of " << fIntervalCount << " momentum intervals, it is ignored";
            is_efficiency_map_read = kFALSE;
        }
    }
    fIsEfficiencyMapRead = is_efficiency_map_read;
 
    // all the histograms have been detached from the file above, so it can be safely closed and deleted
    hist_file->Close();
    delete hist_file;
 
    return kTRUE;
}
 
// Creates an event with equal number of particles at all momentum intervals
Bool_t BmnParticleEqualizer1D::ReadEvent(FairPrimaryGenerator* primGen)
{
    if (!isActive)
        return kFALSE;
 
    // the generator equalizes the particles put on the stack by the generators called before it, so it has to be the
    // last one in the list of FairPrimaryGenerator: being registered earlier it would silently miss the particles of
    // the generators registered after it and equalize an incomplete event
    TObjArray* generator_list = primGen->GetListOfGenerators();
    if ((generator_list == nullptr) || (generator_list->Last() != this)) {
        LOG(error) << "The BmnParticleEqualizer1D generator is not the last one registered in FairPrimaryGenerator: "
                      "it has to be added by AddGenerator() after all the other generators of the event";
        return kFALSE;
    }
 
    // the particles of the event are taken from the stack filled by the previous generators
    CbmStack* stack = (gMC == nullptr) ? nullptr : dynamic_cast<CbmStack*>(gMC->GetStack());
    if (stack == nullptr) {
        LOG(error) << "The BmnParticleEqualizer1D generator requires the particle stack of the simulation (CbmStack) "
                      "to count the particles of the event, but it is not available";
        return kFALSE;
    }
 
    /** CLEARING **/
    // reset particle count for all the particle types
    for (auto part_iter = fParticleInfos.begin(); part_iter != fParticleInfos.end(); part_iter++)
        fill(part_iter->second->fParticleNumber.begin(), part_iter->second->fParticleNumber.end(), 0);
    // clear maximum number of a particle type among all particles (divided by momentum intervals)
    fill(fMaxParticles.begin(), fMaxParticles.end(), 0);
 
    /** FORM EVENT DISTRIBUTION FOR GENERATION PARTICLES TO EQUALIZE THEIR NUMBERS **/
    // form particle distributions for the current event (obtained by the previous generators)
    TClonesArray* arrParticles = stack->GetListOfParticles();
    for (int i = 0; i < arrParticles->GetEntries(); i++) {
        // get particle information
        // exclude the secondary particles before anything is computed for them: TParticle::Eta() of a track with a
        // zero transverse momentum is infinite, and the values of a secondary particle are not needed at all
        TParticle* part = (TParticle*)arrParticles->At(i);
        if ((part == nullptr) || (!part->IsPrimary()))
            continue;
 
        Int_t part_code = part->GetPdgCode();
        Double_t part_p = part->P();
        Double_t part_eta = part->Eta();
        // TParticle::Phi() returns the angle in [0, 2pi), while the sample histograms are filled with ATan2(py, px),
        // i.e. in [-pi, pi], so the very same definition is used here
        Double_t part_phi = TMath::ATan2(part->Py(), part->Px());
 
        // check particle momentum is inside the momentum range: a momentum below fPMin would be truncated to interval 0
        // by the int conversion (not to a negative index), so the range is verified on the momentum value itself first
        if ((part_p < fPMin) || (part_p >= fPMax))
            continue;
        int interval_number = (part_p - fPMin) / fIntervalStep;
        if ((interval_number < 0) || (interval_number >= fIntervalCount))
            continue;
 
        // add found particle according to the momentum to the "PDG -> Particle Count" map (increase by 1)
        LOG(debug) << TString::Format("Found particle: PDG=%i P=%4.2f Eta=%4.2f", part_code, part_p, part_eta);
        if (auto part_iter = fParticleInfos.find(part_code); part_iter != fParticleInfos.end()) {
            // Skip the particles the generator cannot reproduce: those whose momentum, pseudorapidity or azimuth angle
            // falls into a bin where the sample of the interval has no statistics at all. The values are checked
            // separately, as they are also sampled separately, see BmnParticleEqualizer2D for the joint check
            if (fUseSampleStatisticsCut) {
                TH1D* p_hist = fParticleInfos[part_code]->fPHistVector[interval_number].get();
                if (p_hist->GetBinContent(p_hist->FindBin(part_p)) < 1)
                    continue;
                TH1D* eta_hist = fParticleInfos[part_code]->fEtaHistVector[interval_number].get();
                if (eta_hist->GetBinContent(eta_hist->FindBin(part_eta)) < 1)
                    continue;
                TH1D* phi_hist = fParticleInfos[part_code]->fPhiHistVector[interval_number].get();
                if (phi_hist->GetBinContent(phi_hist->FindBin(part_phi)) < 1)
                    continue;
            }
 
            part_iter->second->fParticleNumber[interval_number] += 1;
            // update maximum number among all the particle if greater (for the current momentum interval)
            if (part_iter->second->fParticleNumber[interval_number] > fMaxParticles[interval_number])
                fMaxParticles[interval_number] = part_iter->second->fParticleNumber[interval_number];
        }
    }   // for (int i = 0; i < arrParticles->GetEntries(); i++)
 
    // The sample histograms are filled with the momenta of the simulated particles, i.e. in the laboratory frame, so
    // the values sampled from them already include the beam angle. However FairPrimaryGenerator::AddTrack() rotates
    // every added track by the beam angle of the event once again, so the rotation is undone here to keep the added
    // particles in the very same frame as the sampled distributions. The angles are taken from the event header, where
    // they are written by FairPrimaryGenerator::MakeBeamAngle() and are left zero if the beam angle smearing is
    // switched off, so the compensation is an identity transformation in that case.
    // AddTrack() also rotates the track around the z-axis by its own event plane angle, but only if the rotation was
    // switched on by FairPrimaryGenerator::SetEventPlane(), which is not used in BmnRoot: the SetEventPlane() methods
    // called in the simulation macros belong to the generators themselves (MpdUnigenGenerator, MpdUrqmdGenerator, ...)
    // and those rotate their own particles, not the ones added here. That rotation cannot be compensated anyway: it is
    // applied before the beam angle one, and its angle is not available (the header field RotZ is not it, because the
    // generators listed above overwrite RotZ with their own reaction plane angle).
    // The beam angle rotation is TVector3::RotateUz(), the inverse of it is the transposed matrix, i.e. the projections
    // onto the images of the basis vectors
    FairMCEventHeader* event_header = primGen->GetEvent();
    TVector3 beam_x(1., 0., 0.), beam_y(0., 1., 0.), beam_z(0., 0., 1.);
    if ((event_header != nullptr) && ((event_header->GetRotX() != 0.) || (event_header->GetRotY() != 0.))) {
        TVector3 beam_direction(TMath::Tan(event_header->GetRotX()), TMath::Tan(event_header->GetRotY()), 1.);
        beam_direction = beam_direction.Unit();
        beam_x.RotateUz(beam_direction);
        beam_y.RotateUz(beam_direction);
        beam_z.RotateUz(beam_direction);
    }
 
    // generate new particles to equalize their numbers for different types
    for (auto pdg_key : fParticles) {
        LOG(debug) << TString::Format("Current particle type to be equalized: pdg=%i y0=%4.2f sigma_y=%4.2f T_pt=%6.4f",
                                      pdg_key, fParticleInfos[pdg_key]->fY0, fParticleInfos[pdg_key]->fSigma,
                                      fParticleInfos[pdg_key]->fT);
        // the fitted momentum function is the same for the whole particle type and is used as a fallback for the
        // intervals with too little sample statistics (the way to sample the momentum in every momentum interval was
        // defined in ReadSampleHistograms); fetch it once here to avoid searching for it for every added particle
        TF1* p_dist_func = fParticleInfos[pdg_key]->fPHist->GetFunction(TString::Format("distP_pdg%d", pdg_key));
 
        Double_t cur_p = fPMin;
        // add a number of particles for all momentum intervals to equalize their count
        LOG(debug) << "Adding paticle to equalize the numbers...";
        for (int cur_interval = 0; cur_interval < fIntervalCount; cur_interval++, cur_p += fIntervalStep) {
            // the momentum cannot be sampled in this interval at all (it was reported while reading the histograms)
            if (fParticleInfos[pdg_key]->fPSampling[cur_interval] == stParticleInfo::EPSampling::kNone)
                continue;
 
            // the number of the particles of this type to be reached in the interval: the target given by
            // SetTargetPerInterval() if it is set, the maximum number among the particle types found in this interval
            // of this event otherwise (the default behaviour, which equalizes the types but keeps the momentum
            // spectrum of the input events)
            Double_t target_number =
                (fTargetPerInterval > 0.) ? fTargetPerInterval : (Double_t)fMaxParticles[cur_interval];
            Double_t particle_number = target_number - fParticleInfos[pdg_key]->fParticleNumber[cur_interval];
            // the added particles are reconstructed with the efficiency of the interval, so more of them are needed
            // where it is low; the limit keeps the factor finite in the intervals with an almost zero efficiency.
            // A negative value means the event is already richer than the target and nothing is to be added, so it
            // must not be scaled: dividing a surplus by the efficiency would make it suppress the next events
            if ((fUseEfficiencyCorrection) && (particle_number > 0.))
                particle_number /= TMath::Max(fParticleInfos[pdg_key]->fEfficiency[cur_interval], fMinEfficiency);
 
            // The target may be fractional, so the part below one particle is accumulated between the events: without
            // it a target of, say, 0.3 particles per event would always be rounded down to nothing. With the default
            // target the increment is an integer and the accumulated value stays zero, i.e. nothing changes
            Double_t& particle_debt = fParticleInfos[pdg_key]->fParticleDebt[cur_interval];
            particle_debt += particle_number;
            int add_particle = (int)TMath::Floor(particle_debt);
            // the event already has more particles than the target: nothing is added and the surplus is kept, so that
            // fewer particles are added to this interval in the next events
            if (add_particle < 0)
                add_particle = 0;
            particle_debt -= add_particle;
            // Note: a particle skipped inside the loop below (its kinematics could not be sampled) is not returned to
            // the debt on purpose. All such skips are excluded by the checks made while reading the sample histograms,
            // so they should not happen at all; and were a skip systematic for some interval, returning the particles
            // would make the debt of that interval grow with every event and 'add_particle' grow together with it
 
            LOG(debug) << endl
                       << "   " << add_particle << " particle with PDG = " << pdg_key << " (" << cur_p << " GeV)";
            // the pseudorapidity and the azimuth angle are taken from the interval defined while reading the histograms
            // (the current one or the nearest one having statistics), so they are fetched once per interval here
            Int_t eta_interval = fParticleInfos[pdg_key]->fEtaHistIndex[cur_interval];
            Int_t phi_interval = fParticleInfos[pdg_key]->fPhiHistIndex[cur_interval];
            if ((eta_interval < 0) || (phi_interval < 0)) {
                LOG(warning) << "There is no sample statistics of the pseudorapidity or of the azimuth angle for pdg="
                             << pdg_key << " in any momentum interval, no particles are added";
                continue;
            }
            for (Int_t k = 0; k < add_particle; k++) {
                // sample the momentum either from the P histogram of the interval itself or, if the interval has too
                // little sample statistics, from the fitted momentum function of the full histogram (fetched above)
                Double_t p =
                    (fParticleInfos[pdg_key]->fPSampling[cur_interval] == stParticleInfo::EPSampling::kIntervalHist)
                        ? fParticleInfos[pdg_key]->fPHistVector[cur_interval]->GetRandom()
                        : p_dist_func->GetRandom(cur_p, cur_p + fIntervalStep);
                // Double_t pt = fParticleInfos[pdg_key]->fPtDistribution->GetRandom(cur_pt, cur_pt + fIntervalStep);
                // Double_t y = gRandom->Gaus(particle_info->fY0, particle_info->fSigma);
                // Double_t mt = TMath::Sqrt(fPDGMass * fPDGMass + pt * pt);
                // Double_t pz = mt * TMath::SinH(y);
                Double_t eta = fParticleInfos[pdg_key]->fEtaHistVector[eta_interval]->GetRandom();
                // the azimuth angle is sampled from the histogram of the sample as well, so the added particles are
                // distributed as the ones of the sample, i.e. inside the detector acceptance, and not uniformly
                Double_t phi = fParticleInfos[pdg_key]->fPhiHistVector[phi_interval]->GetRandom();
                // the fitted momentum function is not required to be positive over the whole interval and
                // TF1::GetRandom() takes the absolute values of its negative parts, so the sampled values are
                // verified before a particle with a zero or undefined momentum is put on the stack
                if ((!TMath::Finite(p)) || (p <= 0) || (!TMath::Finite(eta))) {
                    LOG(warning) << "An incorrect momentum (" << p << ") or pseudorapidity (" << eta
                                 << ") was sampled for pdg=" << pdg_key << ", the particle is skipped";
                    continue;
                }
                Double_t pt = p / TMath::CosH(eta);
                TVector3 mom(pt * TMath::Cos(phi), pt * TMath::Sin(phi), pt * TMath::SinH(eta));
 
                // undo the beam angle rotation that AddTrack() applies below (see the comment before the loop)
                mom.SetXYZ(mom * beam_x, mom * beam_y, mom * beam_z);
 
                LOG(debug) << "Particle generated: pdg=" << pdg_key << " p=" << p << " eta=" << eta;
                // AddTrack parameters: pdgid, px, py, pz, vx, vy, vz
                // the vertex (vx, vy, vz) is given relative to the event vertex: FairPrimaryGenerator::AddTrack adds
                // the event vertex itself, so (0, 0, 0) means "at the primary vertex"
                Int_t particle_count = arrParticles->GetEntriesFast();
                primGen->AddTrack(pdg_key, mom.X(), mom.Y(), mom.Z(), 0., 0., 0.);
                // mark the particle just put on the stack (the bit is carried over to CbmMCTrack)
                // the second possible way: use weight parameter of AddTrack(), then parse in CbmMCTrack
                // AddTrack() discards a particle whose PDG code is absent in the database,
                // hence the check that the stack has grown
                if (arrParticles->GetEntriesFast() > particle_count) {
                    TParticle* added_particle = (TParticle*)arrParticles->At(arrParticles->GetEntriesFast() - 1);
                    if (added_particle != nullptr)
                        added_particle->SetBit(kAddedParticleBit);
                }
            }
        }
    }
 
    return kTRUE;
}
 
// for every momentum interval define the index of the interval histogram to be used to sample a value
void BmnParticleEqualizer1D::DefineNearestFilledIntervals(const vector<unique_ptr<TH1D>>& hist_vector,
                                                          vector<Int_t>& index_vector)
{
    const Int_t interval_count = (Int_t)hist_vector.size();
    index_vector.assign(interval_count, -1);
 
    for (Int_t cur_interval = 0; cur_interval < interval_count; cur_interval++) {
        // look for the nearest interval having statistics, starting from the current one; at an equal distance the
        // lower momentum interval is preferred, as it usually has more statistics
        for (Int_t shift = 0; shift < interval_count; shift++) {
            const Int_t lower_interval = cur_interval - shift, upper_interval = cur_interval + shift;
            if ((lower_interval >= 0) && (hist_vector[lower_interval] != nullptr)
                && (hist_vector[lower_interval]->GetEntries() > 2))
            {
                index_vector[cur_interval] = lower_interval;
                break;
            }
            if ((upper_interval < interval_count) && (hist_vector[upper_interval] != nullptr)
                && (hist_vector[upper_interval]->GetEntries() > 2))
            {
                index_vector[cur_interval] = upper_interval;
                break;
            }
        }
    }
}
 
// static function to fit the obtained sample histograms (to show whether the results are appropriate)
void BmnParticleEqualizer1D::FitSampleHistograms(TString input_histo_file)
{
    TFile* f = TFile::Open(input_histo_file);
    if ((!f) || (!f->IsOpen())) {
        cout << "ERROR: unable to open ROOT file: " << input_histo_file << endl;
        return;
    }
    for (int particle_number : fParticles) {
        /*TCanvas* c = */ new TCanvas(TString::Format("c_%d", particle_number),
                                      TString::Format("Distributions for particle PDG = %d", particle_number), 1000,
                                      500);
        TH1D* h1 = (TH1D*)f->Get(TString::Format("hP_pdg%d", particle_number));
        if (h1 == nullptr) {
            cout << "ERROR: the momentum histogram is absent for PDG = " << particle_number << endl;
            continue;
        }
 
        Double_t x_max_value = h1->GetXaxis()->GetBinCenter(h1->GetMaximumBin());
        Double_t y_max_value = h1->GetBinContent(h1->GetMaximumBin());
        Double_t p_avg = (fPMin + fPMax) / 2;
        cout << "avg_value = " << p_avg << endl;
        cout << "x_max_value = " << x_max_value << endl;
        cout << "y_max_value = " << y_max_value << endl;
        // the very same function and fit as in FitPHistogram(), so what is shown is what the generator uses
        FitMomentumFunction(h1, "distFit", fPMin, fPMax);
        // the function is attached without drawing ('0' option of the fit), the flag is reset to show it here
        TF1* p_distribution = h1->GetFunction("distFit");
        if (p_distribution != nullptr)
            p_distribution->ResetBit(TF1::kNotDraw);
        h1->Draw();
    }
}
 
// static function to show the obtained sample histograms with obtained fit (to check whether the results are
// appropriate)
int BmnParticleEqualizer1D::ShowSampleHistograms(TString input_histo_file)
{
    TFile* f = TFile::Open(input_histo_file);
    if ((!f) || (!f->IsOpen())) {
        cout << "ERROR: unable to open ROOT file: " << input_histo_file << endl;
        return -1;
    }
 
    for (int particle_number : fParticles) {
        /*TCanvas* c = */ new TCanvas(TString::Format("c_%d", particle_number),
                                      TString::Format("Distributions for particle PDG = %d", particle_number), 1000,
                                      500);
        TH1D* h1 = (TH1D*)f->Get(TString::Format("hP_pdg%d", particle_number));
        if (h1 == nullptr) {
            cout << "ERROR: the momentum histogram is absent for PDG = " << particle_number << endl;
            continue;
        }
        TF1* f1 = h1->GetFunction(TString::Format("distP_pdg%d", particle_number));
        if (f1 == nullptr) {
            cout << "ERROR: the fitted momentum function is absent for PDG = " << particle_number << endl;
            continue;
        }
        f1->Draw();
        h1->Draw("SAME");
    }
 
    return 0;
}
 
// define PDG code of a particle by the given mass
// among particle types set by 'vecParticlePDG' vector or in the TDatabasePDG if the vector is empty
// tolerance: relative if tolerance > 0, absolute (in GeV) in case of tolerance <= 0
// return PDG code or 0 if no corresponding particle was found
Int_t BmnParticleEqualizer1D::GetPDGCodeByMass(Double_t mass, Double_t tolerance, const vector<int>& vecParticlePDG)
{
    TDatabasePDG* pdg_base = TDatabasePDG::Instance();
    if (!pdg_base) {
        cout << "ERROR: Could not get TDatabasePDG instance" << endl;
        return 0;
    }
 
    if (mass < 0)
        mass *= -1;
 
    if (vecParticlePDG.empty()) {
        TIter next(pdg_base->ParticleList());
        while (TParticlePDG* particle = (TParticlePDG*)next()) {
            if (TMath::Abs(particle->Mass() - mass) < (tolerance > 0 ? tolerance * mass : -1 * tolerance))
                return particle->PdgCode();
        }
    } else {
        for (int particle_number : vecParticlePDG) {
            TParticlePDG* particle = pdg_base->GetParticle(particle_number);
            if (particle == nullptr)
                continue;
            if (TMath::Abs(particle->Mass() - mass) < (tolerance > 0 ? tolerance * mass : -1 * tolerance))
                return particle->PdgCode();
        }
    }
 
    return 0;
}
 
// define the particle type of a reconstructed (DST) track: by the matched MC track or by the ToF mass and the charge
Int_t BmnParticleEqualizer1D::GetGlobalTrackPDG(BmnGlobalTrack* global_track,
                                                BmnTrackMatch* global_match,
                                                TClonesArray* mc_tracks,
                                                Double_t& tof_mass)
{
    tof_mass = -1.;
    if (global_track == nullptr)
        return 0;
 
    // simulated data: the type of the related MC track. The index stored in the global track itself is taken first:
    // BmnGlobalTrack::GetRefIndex() always returns -1 for now (it is kept for compatibility), but it is planned to be
    // filled again; until then the MC track is given by the matching
    if (mc_tracks != nullptr) {
        Int_t mc_index = global_track->GetRefIndex();
        if ((mc_index < 0) && (global_match != nullptr) && (global_match->GetNofLinks() > 0))
            mc_index = global_match->GetMatchedLink().GetIndex();
        if (mc_index > -1) {
            if (mc_index >= mc_tracks->GetEntriesFast())
                return 0;
            CbmMCTrack* mc_track = (CbmMCTrack*)mc_tracks->UncheckedAt(mc_index);
            return (mc_track == nullptr) ? 0 : mc_track->GetPdgCode();
        }
        // the track has the matching information, but no MC track is matched to it
        if (global_match != nullptr)
            return 0;
        // there is no information on the related MC track at all: the type is defined by the ToF mass below
    }
 
    // experimental data: the mass measured by ToF-400 or, if it is absent, by ToF-700
    Double_t mass_tof2 = global_track->GetMass2(1);
    if (mass_tof2 == -1000.0)
        mass_tof2 = global_track->GetMass2(2);
    // GetMass2 returns the squared mass, a negative value is caused by the ToF resolution
    if ((mass_tof2 == -1000.0) || (mass_tof2 < 0))
        return 0;
    FairTrackParam* par_first = global_track->GetParamFirst();
    if ((par_first == nullptr) || (par_first->GetQp() == 0))
        return 0;
 
    // the mass does not distinguish a particle from its antiparticle, so only the types with the charge of the track
    // (the sign of q/p) are the candidates. An empty list must not be passed on: GetPDGCodeByMass() would search the
    // whole PDG database then
    vector<Int_t> candidates;
    for (Int_t pdg_code : fParticles) {
        TParticlePDG* particle = TDatabasePDG::Instance()->GetParticle(pdg_code);
        if ((particle != nullptr) && (particle->Charge() * par_first->GetQp() > 0))
            candidates.push_back(pdg_code);
    }
    if (candidates.empty())
        return 0;
 
    tof_mass = TMath::Sqrt(mass_tof2);
    return GetPDGCodeByMass(tof_mass, -0.03, candidates);
}
 
// draw histogram with multiplicity of different particle types for all the momentum intervals
// as well as P, Pt, Y, Eta, Theta, Phi distributions after the equalizing
// it work both for SIM and DST files listed in the given text file ('input_list_file' var)
int BmnParticleEqualizer1D::ShowResultDistributions(TString input_list_file)
{
    gSystem->ExpandPathName(input_list_file);
 
    // get file list from the input text file
    list<string> file_list;
    ifstream list_file(input_list_file.Data());
    string s;
    while (getline(list_file, s))
        file_list.push_back(s);
    if (file_list.empty()) {
        cout << "ERROR: There are no files in the list of the input text file: " << input_list_file << endl;
        return -5;
    }
 
    // form ROOT chain for the file list
    TChain dataChain("bmndata");
    for (auto const& cur_file_path : file_list) {
        TString strFilePath(cur_file_path);
        gSystem->ExpandPathName(strFilePath);
 
        dataChain.Add(strFilePath);
    }
 
    // flag whether the input files have SIM or DST format
    bool isDST = false;
    TClonesArray *mcTracks = nullptr, *globalTracks = nullptr, *globalMatches = nullptr;
    BmnMCInfoDst* mcInfo = nullptr;
    if (dataChain.GetBranch("MCTrack") == nullptr) {
        if (dataChain.GetBranch("BmnMCInfo.") == nullptr) {
            LOG(error) << "'MCTrack' and 'BmnMCInfo.' branches were not found, please, check the structure of the "
                          "files, exiting...";
            return -1;
        }
        dataChain.SetBranchAddress("BmnMCInfo.", &mcInfo);
        if (mcInfo == nullptr) {
            LOG(error)
                << "'BmnMCInfo.' branch was not read correctly, please, check the structure of the files, exiting...";
            return -2;
        }
        dataChain.SetBranchAddress("BmnGlobalTrack", &globalTracks);
        if (globalTracks == nullptr) {
            LOG(error) << "'BmnGlobalTrack' branch was not read correctly, please, check the structure of the files, "
                          "exiting...";
            return -4;
        }
        // a reconstructed track is related to its MC track by the index stored in the track itself
        // (BmnGlobalTrack::GetRefIndex(), which always returns -1 for now) or by the separate matching branch
        if (dataChain.GetBranch("BmnGlobalTrackMatch") != nullptr)
            dataChain.SetBranchAddress("BmnGlobalTrackMatch", &globalMatches);
        else
            LOG(warning) << "'BmnGlobalTrackMatch' branch was not found, the particle types of the reconstructed "
                            "tracks without an MC index are defined by the mass measured by the ToF detectors";
        isDST = true;
        // dataChain.SetBranchStatus("*", false);
        // dataChain.SetBranchStatus("BmnMCInfo*", true);
        // dataChain.SetBranchStatus("BmnGlobalTrack", true);
    } else {
        dataChain.SetBranchAddress("MCTrack", &mcTracks);
        if (mcTracks == nullptr) {
            LOG(error)
                << "MCTrack' branch was not read correctly, please, check the structure of the files, exiting...";
            return -3;
        }
        // dataChain.SetBranchStatus("*", false);
        // dataChain.SetBranchStatus("MCTrack*", true);
    }
 
    // create histogram with bars
    TCanvas* c_number =
        new TCanvas("c_number", "Multiplicity of different particle types for all the momentum intervals", 1000, 500);
    gPad->SetGrid();
 
    int bin_count = fIntervalCount, cur_hist = 0;
    // create histograms for particle multiplicities
    map<int, TH1I*> mapParticleNumberHists;
    for (int particle_number : fParticles) {
        TH1I* hParticleNumber =
            new TH1I(TString::Format("hParticleNumber_%d", particle_number),
                     TString::Format("hParticleNumber_%d", particle_number), bin_count, fPMin, fPMax);
        // gStyle->SetTitleFontSize(0.03);    // only one approach to change histogram title size
        hParticleNumber->SetStats(kFALSE);
        hParticleNumber->LabelsDeflate();
        hParticleNumber->SetLabelFont(62);
        hParticleNumber->SetMarkerSize(0.6);
        // hParticleNumber->LabelsOption("v");
        hParticleNumber->SetTitle("");
        hParticleNumber->GetXaxis()->SetTitle("momentum interval");
        hParticleNumber->GetXaxis()->SetTitleSize(0.03);
        hParticleNumber->GetXaxis()->CenterTitle();
        hParticleNumber->GetXaxis()->SetLabelSize(0.02 - 0.00017 * bin_count);
        // hParticleNumber->GetXaxis()->CenterLabels();
        hParticleNumber->GetXaxis()->SetNdivisions(bin_count);
        hParticleNumber->GetYaxis()->SetTitle("particle multiplicity");
        // hParticleNumber->GetYaxis()->SetTitleSize(0.03);
        // hParticleNumber->GetYaxis()->SetNdivisions(1007);
        hParticleNumber->GetYaxis()->SetLabelSize(0.02);
        hParticleNumber->SetBarWidth(0.18);
        hParticleNumber->SetBarOffset(0.05 + cur_hist * 0.18);
        switch (cur_hist) {
            case 0:
                hParticleNumber->SetFillColor(49);
                break;
            case 1:
                hParticleNumber->SetFillColor(50);
                break;
            case 2:
                hParticleNumber->SetFillColor(42);
                break;
            case 3:
                hParticleNumber->SetFillColor(29);
                break;
            case 4:
                hParticleNumber->SetFillColor(36);
                break;
            default:
                break;
        }
 
        mapParticleNumberHists[particle_number] = hParticleNumber;
        cur_hist++;
    }
 
    // create histograms for P, Pt, Y, Eta, Theta, Phi
    map<int, TH1D*> mapPHists, mapPtHists, mapYHists, mapEtaHists, mapThetaHists, mapPhiHists;
    for (int particle_number : fParticles) {
        TH1D* hP = new TH1D(TString::Format("hP_%d", particle_number), TString::Format("hP_%d", particle_number), 80,
                            fPMin, fPMax);
        mapPHists[particle_number] = hP;
        TH1D* hPt = new TH1D(TString::Format("hPt_%d", particle_number), TString::Format("hPt_%d", particle_number), 60,
                             fPtMin, fPtMax);
        mapPtHists[particle_number] = hPt;
        TH1D* hY = new TH1D(TString::Format("hY_%d", particle_number), TString::Format("hY_%d", particle_number), 60,
                            fYMin, fYMax);
        mapYHists[particle_number] = hY;
        TH1D* hEta = new TH1D(TString::Format("hEta_%d", particle_number), TString::Format("hEta_%d", particle_number),
                              60, fEtaMin, fEtaMax);
        mapEtaHists[particle_number] = hEta;
        TH1D* hTheta = new TH1D(TString::Format("hTheta_%d", particle_number),
                                TString::Format("hTheta_%d", particle_number), 61, fThetaMin, fThetaMax);
        mapThetaHists[particle_number] = hTheta;
        TH1D* hPhi = new TH1D(TString::Format("hPhi_%d", particle_number), TString::Format("hPhi_%d", particle_number),
                              64, fPhiMin, fPhiMax);
        mapPhiHists[particle_number] = hPhi;
    }
 
    int iEventNumber = dataChain.GetEntries(), iCurrentEvent = 0;
    if (iEventNumber < 1) {
        cout << "ERROR: There are no correct entries in the file list stored at: \"" << input_list_file << "\"" << endl;
        return -6;
    }
 
    // fill all the histograms
    for (iCurrentEvent = 0; iCurrentEvent < iEventNumber; iCurrentEvent++) {
        // Next Event
        if ((iCurrentEvent + 1) % 1000 == 0)
            LOG(info) << "Processing event #" << iCurrentEvent + 1 << " of " << iEventNumber;
        dataChain.GetEntry(iCurrentEvent);
 
        Int_t track_count = 0;
        TClonesArray* mc_tracks = nullptr;
        if (isDST) {
            mc_tracks = mcInfo->GetMCTracks();
            track_count = globalTracks->GetEntries();
        } else {
            mc_tracks = mcTracks;
            track_count = mc_tracks->GetEntries();
        }
 
        // Getting tracks of the event and their corresponding components
        for (Int_t iTrack = 0; iTrack < track_count; iTrack++) {
            Int_t track_pdg = -1;
            Double_t track_mass = -1;   // real mass in GeV (an integer would truncate it to 0 and spoil the rapidity)
            Double_t track_p = 0, track_pt = 0, track_y = 0, track_eta = 0, track_theta = 0, track_phi = 0;
            CbmMCTrack* mc_track;
            BmnGlobalTrack* global_track;
 
            if (isDST) {
                // cout<<"iTrack = "<<iTrack<<", track_count = "<<track_count<<endl;
                global_track = (BmnGlobalTrack*)globalTracks->UncheckedAt(iTrack);
                // the reconstructed tracks are taken in the same acceptance as the sample ones (see SetSampleTrackCut):
                // with the default cut only the tracks matched to a ToF hit, i.e. the ones identified by ToF
                Bool_t has_tof = (global_track->GetTof1HitIndex() > -1) || (global_track->GetTof2HitIndex() > -1);
                if ((fSampleTrackCut > 1) && (!has_tof))
                    continue;
 
                // define the particle type: by the related MC track, or by the ToF mass and the charge if the track has
                // no information on it
                BmnTrackMatch* global_match = ((globalMatches != nullptr) && (iTrack < globalMatches->GetEntriesFast()))
                                                  ? (BmnTrackMatch*)globalMatches->UncheckedAt(iTrack)
                                                  : nullptr;
                Double_t tof_mass = -1.;
                track_pdg = GetGlobalTrackPDG(global_track, global_match, mc_tracks, tof_mass);
                if (find(fParticles.begin(), fParticles.end(), track_pdg) == fParticles.end())
                    continue;
                // the mass for the rapidity below: the measured one if the type is defined by it, the nominal otherwise
                if (tof_mass > 0)
                    track_mass = tof_mass;
                else {
                    TParticlePDG* particle = TDatabasePDG::Instance()->GetParticle(track_pdg);
                    if (particle == nullptr)
                        continue;
                    track_mass = particle->Mass();
                }
 
                // Getting track params predicted by Kalman filter ...
                FairTrackParam* parFirst = global_track->GetParamFirst();
                // FairTrackParam* parLast = track->GetParamLast();
                if (parFirst == nullptr)
                    continue;
                Double_t track_qp = parFirst->GetQp();
                track_p = (TMath::Abs(track_qp) > 1.e-4) ? 1. / TMath::Abs(track_qp) : 1.e4;
                TVector3 vecMom;
                parFirst->Momentum(vecMom);
                track_pt = TMath::Sqrt(vecMom.x() * vecMom.x() + vecMom.y() * vecMom.y());
                Double_t e = TMath::Sqrt(track_p * track_p + track_mass * track_mass);
                track_y = 0.5 * TMath::Log((e + vecMom.z()) / (e - vecMom.z()));
                track_eta = 0.5 * TMath::Log((track_p + vecMom.z()) / (track_p - vecMom.z()));
                Double_t track_tx = parFirst->GetTx(), track_ty = parFirst->GetTy();
                track_phi = TMath::ATan2(track_ty, track_tx);   // phi = arctan(tgy/tgx)
                // theta = arctan(sqrt(tgx^2 + tgy^2)): the equivalent 'arctan(tgx/cos(phi))' returns a negative angle
                // when cos(phi) is negative, and such values fall into the underflow bin of the Theta histogram
                track_theta = TMath::ATan(TMath::Sqrt(track_tx * track_tx + track_ty * track_ty));
            }   // if (isDST)
            else
            {
                mc_track = (CbmMCTrack*)mc_tracks->At(iTrack);
                if (mc_track == nullptr)
                    continue;
                if (mc_track->GetMotherId() >= 0)
                    continue;
 
                // filter particles
                if (fSampleTrackCut > 0) {
                    // exclude tracks without points in Silicon and GEMs
                    if ((mc_track->GetNPoints(kGEM) < 1) || (mc_track->GetNPoints(kSILICON) < 1))
                        continue;
                    if (fSampleTrackCut > 1) {
                        // exclude tracks without points in ToF detectors. kTOF701 is not checked: CbmMCTrack keeps the
                        // points of TOF700 in a single bit field, returned by GetNPoints() both for kTOF and for
                        // kTOF701
                        if ((mc_track->GetNPoints(kTOF) < 1) && (mc_track->GetNPoints(kTOF1) < 1))
                            continue;
                    }
                }
 
                track_pdg = mc_track->GetPdgCode();
                track_p = mc_track->GetP();
                track_pt = mc_track->GetPt();
                track_y = mc_track->GetRapidity();
                track_eta = 0.5 * TMath::Log((track_p + mc_track->GetPz()) / (track_p - mc_track->GetPz()));
                track_theta = TMath::ATan2(track_pt, mc_track->GetPz());
                track_phi = TMath::ATan2(mc_track->GetPy(), mc_track->GetPx());
            }   // if not (isDST)
 
            auto iter_pdg = find(fParticles.begin(), fParticles.end(), track_pdg);
            if (iter_pdg == fParticles.end())
                continue;
 
            if ((track_p >= fPMin) && (track_p < fPMax)) {
                mapParticleNumberHists[*iter_pdg]->Fill(track_p);
                mapPHists[*iter_pdg]->Fill(track_p);
                mapPtHists[*iter_pdg]->Fill(track_pt);
                mapYHists[*iter_pdg]->Fill(track_y);
                mapEtaHists[*iter_pdg]->Fill(track_eta);
                mapThetaHists[*iter_pdg]->Fill(track_theta);
                mapPhiHists[*iter_pdg]->Fill(track_phi);
            }
        }   // for (Int_t iTrack = 0; iTrack < track_count; iTrack++)
    }   // for (Int_t iEvent = 0; iEvent < iEventNumber; iEvent++) {
 
    // if no required branches in the tree or events then clear memory and exit
    if (iCurrentEvent == 0) {
        cout << "ERROR: There are no necessary tree branches or DST events in the files from the list: "
             << input_list_file << endl;
        for (const auto& pair : mapPHists)
            delete pair.second;
        for (const auto& pair : mapPtHists)
            delete pair.second;
        for (const auto& pair : mapYHists)
            delete pair.second;
        for (const auto& pair : mapEtaHists)
            delete pair.second;
        for (const auto& pair : mapThetaHists)
            delete pair.second;
        for (const auto& pair : mapPhiHists)
            delete pair.second;
        delete c_number;
        return -7;
    }
 
    // find max Y-axis value
    Double_t max_y_value = 0;
    for (const auto& pair_hist : mapParticleNumberHists) {
        if (max_y_value < pair_hist.second->GetMaximum())
            max_y_value = pair_hist.second->GetMaximum();
    }
 
    // draw bar histogram with particle multiplicity for all the particle types
    TLegend* legend = new TLegend(0.7, 0.74, 0.85, 0.88);
    TString hist_options = "bar2,hist,text";   // 'text' option for vertical captions with event count
    cur_hist = 0;
    for (const auto& pair_hist : mapParticleNumberHists) {
        if (cur_hist == 0) {
            if (pair_hist.second->GetMaximum() < max_y_value)
                pair_hist.second->SetMaximum(max_y_value + 10);
            pair_hist.second->DrawCopy(hist_options);
        } else {
            pair_hist.second->DrawCopy(hist_options + ",same");
        }
 
        legend->AddEntry(pair_hist.second, TString::Format("Multiplitcity for %d", pair_hist.first), "f");
        legend->Draw();
        cur_hist++;
    }
 
    // draw histograms (P, Pt, Y, Eta, Theta, Phi) for different particle types
    TCanvas* c_p = new TCanvas("c_p", "Particle Momentum", 1000, 500);
    c_p->Divide((fParticles.size() / 3) + 1, 3);
    cur_hist = 1;
    for (int particle_number : fParticles) {
        c_p->cd(cur_hist);
        mapPHists[particle_number]->Draw();
        cur_hist++;
    }
    TCanvas* c_pt = new TCanvas("c_pt", "Transverse Momentum", 1000, 500);
    c_pt->Divide((fParticles.size() / 3) + 1, 3);
    cur_hist = 1;
    for (int particle_number : fParticles) {
        c_pt->cd(cur_hist);
        mapPtHists[particle_number]->Draw();
        cur_hist++;
    }
    TCanvas* c_y = new TCanvas("c_y", "Rapidity", 1000, 500);
    c_y->Divide((fParticles.size() / 3) + 1, 3);
    cur_hist = 1;
    for (int particle_number : fParticles) {
        c_y->cd(cur_hist);
        mapYHists[particle_number]->Draw();
        cur_hist++;
    }
    TCanvas* c_eta = new TCanvas("c_eta", "Pseudorapidity", 1000, 500);
    c_eta->Divide((fParticles.size() / 3) + 1, 3);
    cur_hist = 1;
    for (int particle_number : fParticles) {
        c_eta->cd(cur_hist);
        mapEtaHists[particle_number]->Draw();
        cur_hist++;
    }
    TCanvas* c_theta = new TCanvas("c_theta", "Theta angle", 1000, 500);
    c_theta->Divide((fParticles.size() / 3) + 1, 3);
    cur_hist = 1;
    for (int particle_number : fParticles) {
        c_theta->cd(cur_hist);
        mapThetaHists[particle_number]->Draw();
        cur_hist++;
    }
    TCanvas* c_phi = new TCanvas("c_phi", "Phi Angle", 1000, 500);
    c_phi->Divide((fParticles.size() / 3) + 1, 3);
    cur_hist = 1;
    for (int particle_number : fParticles) {
        c_phi->cd(cur_hist);
        mapPhiHists[particle_number]->Draw();
        cur_hist++;
    }
 
    return 0;
}