import ROOT
import numpy as np

# df = ROOT.RDataFrame(10000)
# df = df.Define("angle", "gRandom->Uniform(0, 3.14)")
# df = df.Define("nTracks", "gRandom->Integer(3)")
# df.Snapshot("tree", "output.root")
df = ROOT.RDataFrame("tree", "output.root")
# verbosity = ROOT.Experimental.RLogScopedVerbosity(ROOT.Detail.RDF.RDFLogChannel(), ROOT.Experimental.ELogLevel.kInfo)

#  The idea is to bin the angles and find the number of events with 0, 1, 2, 3 tracks in each bin
bins = np.linspace(0, 3.14, 10000, dtype=np.double)
count_two_tracks = []
count_one_track = []
count_zero_track = []
for bin_idx in range(bins.size - 1):
    bin_low = bins[bin_idx]
    bin_high = bins[bin_idx+1]
    # print(f"Processing bin ({bin_idx}):\t{bin_low:.3f} - {bin_high:.3f}")
    rdf_bin = df.Filter(f"angle > {bin_low} && angle < {bin_high}", f"{bin_low} < angle < {bin_high}")
    count_two_tracks.append(rdf_bin.Filter("nTracks == 2").Count())
    count_one_track.append(rdf_bin.Filter("nTracks == 1").Count())
    count_zero_track.append(rdf_bin.Filter("nTracks == 0").Count())

# Make histograms of counts
def make_hist(name, bins, count):
    count = [x.GetValue() for x in count]
    hist = ROOT.TH1D(name, name, bins.size - 1, bins)
    for idx, bin_count in enumerate(count):
        hist.SetBinContent(idx+1, bin_count)
    return hist

histogram_two_tracks = make_hist("two_tracks", bins, count_two_tracks)
histogram_one_track = make_hist("one_track", bins, count_one_track)
histogram_zero_track = make_hist("zero_track", bins, count_zero_track)

hist_out = ROOT.TFile("histograms_rdf.root", "RECREATE")
histogram_two_tracks.Write()
histogram_one_track.Write()
histogram_zero_track.Write()
hist_out.Close()