#ifndef AUTOHIST2D_H
#define AUTOHIST2D_H

#include <TH2F.h>
#include <algorithm>
#include <iostream>
#include <limits>
#include <string>
#include <unordered_map>
#include <vector>

// 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();

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

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

    // 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