# Import the apd tools, versions customized for Snakemake
from apd.snakemake import get_analysis_data

# Get the APD dataset for my analysis
dataset = get_analysis_data("sl", "rds_hadronic")

# Parameters for the datasets to be analysed
OUTPUT = "output/"
CONFIG = "mc"
DATATYPE = "2012",
EVENTTYPE = ["13266069", "11266009"]
POLARITY = [ "magdown", "magup" ]

# to run this outside of lxplus and access input/output locations on EOS you
# need to wrap the path (str) with the remote method which can be imported
# from apd.snakemake and returns a Snakemake XRootD remote
rule all:
    input:
        f"{OUTPUT}bmass.root"

# templated rule to produce a ROOT file with the histogram for B_M in a
# specific sample, notice that:
#  - the input uses the dataset object and specifies the wildcards to use
#  - the output is local, and set to be deleted after used in a later rule with the temp() function
rule create_histo:
    input:
        data=lambda w: dataset(datatype=w.datatype, eventtype=w.eventtype, polarity=w.polarity)
    output: temp(f"{OUTPUT}bmass_{{config}}_{{datatype}}_{{eventtype}}_{{polarity}}.root")
    run:
        # Embedded script, for demo only, not a good idea in general
        import ROOT
        inputfiles =  [ f for f in input ]
        print(f"==== Reading {inputfiles}")
        f = ROOT.TFile.Open(output[0], "RECREATE")
        rdf = ROOT.RDataFrame("SignalTuple/DecayTree", input)
        hname = f"B_M_{wildcards.config}_{wildcards.datatype}_{wildcards.eventtype}_{wildcards.polarity}"
        h = rdf.Histo1D((hname, hname, 200, 0., 25e3), "B_M")
        h.Write()
        f.Close()
        print(f"==== Created {output[0]}")

# Rule to gather the files produced by create_histo
# - the Snakemake expand() method can be used to create all the combinations
#   of all parameters
rule gather:
        input:
            expand(OUTPUT+"bmass_{config}_{datatype}_{eventtype}_{polarity}.root",
                    config=CONFIG, datatype=DATATYPE,
                    eventtype=EVENTTYPE,
                    polarity=POLARITY)
        output:
            f"{OUTPUT}bmass.root"
        shell:
            "hadd {output} {input}"