ANASEN_analysis/Armory/AutoHist2D.h

#ifndef AUTOHIST2D_H
#define AUTOHIST2D_H

#include <TCanvas.h>
#include <TH1.h>
#include <TH2F.h>
#include <TDirectory.h>
#include <TTree.h>
#include <TObject.h>
#include <algorithm>
#include <iostream>
#include <limits>
#include <string>
#include <unordered_map>
#include <vector>
#include <filesystem>

// Buffers (name -> x,y) pairs during the event loop and, on WriteAll(),
// books a TH2F per name with a fixed number of bins per axis and a range
// derived from the buffered data, then Write()s it to the current
// TDirectory/TFile so it shows up alongside the tree in a TBrowser.
class AutoHist2D {
public:
    static const int kDefaultBins = 200;

    // Buffers one (xValue, yValue) pair under the histogram 'name'. Axis
    // titles are optional and only need to be given once (first call wins).
    static void Fill(const std::string& name, double xValue, double yValue,
                      const std::string& xTitle = "", const std::string& yTitle = "") {
        Buffer& buf = Registry()[name];
        buf.x.push_back(xValue);
        buf.y.push_back(yValue);
        if (buf.xTitle.empty() && !xTitle.empty()) buf.xTitle = xTitle;
        if (buf.yTitle.empty() && !yTitle.empty()) buf.yTitle = yTitle;
    }

    // Books, fills and Write()s every registered histogram (nbins x nbins,
    // range taken from min/max of the buffered data) to gDirectory, then
    // clears the buffers.
    static void WriteAll(int nbins = kDefaultBins) {
        for (auto& kv : Registry()) {
            const std::string& name = kv.first;
            Buffer& buf = kv.second;
            if (buf.x.empty()) continue;

            double xlo = *std::min_element(buf.x.begin(), buf.x.end());
            double xhi = *std::max_element(buf.x.begin(), buf.x.end());
            double ylo = *std::min_element(buf.y.begin(), buf.y.end());
            double yhi = *std::max_element(buf.y.begin(), buf.y.end());
            PadRange(xlo, xhi);
            PadRange(ylo, yhi);

            std::string title = name + ";" + (buf.xTitle.empty() ? "x" : buf.xTitle)
                                      + ";" + (buf.yTitle.empty() ? "y" : buf.yTitle);
            TH2F h(name.c_str(), title.c_str(), nbins, xlo, xhi, nbins, ylo, yhi);
            for (size_t i = 0; i < buf.x.size(); ++i) h.Fill(buf.x[i], buf.y[i]);
            h.Write();
            SavePNG(h, name);

            std::cout << "AutoHist2D: wrote \"" << name << "\" (" << buf.x.size()
                       << " entries, " << nbins << "x" << nbins << " bins, range ["
                       << xlo << "," << xhi << "] x [" << ylo << "," << yhi << "])" << std::endl;
        }
        Registry().clear();
    }

    // Writes the standard tree-derived plots from testMacro.C into the
    // current ROOT directory, along with the requested vertex-Z residual.
    static void WriteTreePlotHistograms(TTree* tree) {
        if (!tree) return;

        const auto drawAndWrite = [tree](const char* drawExpression,
                                         const char* selection,
                                         const char* histogramName) {
            tree->Draw(drawExpression, selection, "goff");
            TH1* histogram = dynamic_cast<TH1*>(gDirectory->Get(histogramName));
            if (histogram) {
                histogram->Write("", TObject::kOverwrite);
            } else {
                std::cerr << "Could not create ROOT histogram: " << histogramName << std::endl;
            }
        };

        drawAndWrite("Ex_kin:vZ_track>>hExKinVsVZTrack(800,-500,300,110,-1,10)",
                     "sx3ID >= 0 && !TMath::IsNaN(vZ_track) && EPC > 0.02",
                     "hExKinVsVZTrack");
        drawAndWrite("Ex:vZ>>hExVsVZ",
                     "sx3ID >= 0 && !TMath::IsNaN(vZ) && EPC > 0.02",
                     "hExVsVZ");
        drawAndWrite("theta_recon:thetab>>hThetaReconVsThetaB",
                     "sx3ID >= 0 && Esx3 > 0",
                     "hThetaReconVsThetaB");
        drawAndWrite("vZ_track:vZ>>hVZTrackVsVZ",
                     "sx3ID >= 0 && !TMath::IsNaN(vZ_track)",
                     "hVZTrackVsVZ");
        drawAndWrite("vZ_kin:vZ>>hVZKinVsVZ",
                     "sx3ID >= 0 && !TMath::IsNaN(vZ_kin)",
                     "hVZKinVsVZ");
        drawAndWrite("EBeam_track:beamEnergy>>hEBeamTrackVsBeamEnergy",
                     "sx3ID >= 0 && !TMath::IsNaN(EBeam_track)",
                     "hEBeamTrackVsBeamEnergy");
        drawAndWrite("EBeam_kin:beamEnergy>>hEBeamKinVsBeamEnergy",
                     "sx3ID >= 0 && !TMath::IsNaN(EBeam_kin)",
                     "hEBeamKinVsBeamEnergy");
        drawAndWrite("Ex_kin>>hExKin",
                     "sx3ID >= 0 && Ex_kin > 0",
                     "hExKin");
        drawAndWrite("Ex>>hEx",
                     "sx3ID >= 0 && Ex > 0",
                     "hEx");
        drawAndWrite("vZ_track-vZ>>hVZTrackMinusVZ",
                     "sx3ID >= 0 && EPC > 0.03",
                     "hVZTrackMinusVZ");
    }

private:
    static constexpr const char* kPlotDir = "Plots";

    struct Buffer {
        std::vector<double> x, y;
        std::string xTitle, yTitle;
    };

    // Draws 'h' on a throwaway canvas and saves it under kPlotDir/<name>.png,
    // creating the directory first if it doesn't already exist; overwrites
    // any existing file of that name.
    static void SavePNG(TH2F& h, const std::string& name) {
        std::filesystem::create_directories(kPlotDir);
        TCanvas c(("c_" + name).c_str(), "", 800, 600);
        h.Draw("COLZ");
        c.SaveAs((std::string(kPlotDir) + "/" + name + ".png").c_str());
    }

    // Registered histograms, keyed by name. Function-local static avoids
    // needing a separate translation unit for a header-only class.
    static std::unordered_map<std::string, Buffer>& Registry() {
        static std::unordered_map<std::string, Buffer> registry;
        return registry;
    }

    // Pads a [lo, hi] range by 2% on each side so points at the extremes
    // aren't binned into the edge/overflow bin; handles the lo == hi case.
    static void PadRange(double& lo, double& hi) {
        if (lo == hi) {
            lo -= 1.0;
            hi += 1.0;
            return;
        }
        double pad = 0.02 * (hi - lo);
        lo -= pad;
        hi += pad;
    }
};

#endif