io.py

"""
I/O helpers for GENIE ROOT files.
"""

import cppyy
from ._genie_bindings import NtpMCEventRecord
from contextlib import contextmanager
import contextlib
import os, ctypes # to silence banners

# Load libc for low-level file descriptor dup2
libc = ctypes.CDLL(None)

@contextmanager
def silence_cpp():
    """
    Temporarily redirect both Python and C++ stdout/stderr to /dev/null.
    Useful to silence ROOT/GENIE banner spam.
    """
    devnull_fd = os.open(os.devnull, os.O_WRONLY)
    # Duplicate original FDs
    orig_stdout_fd = os.dup(1)
    orig_stderr_fd = os.dup(2)

    try:
        # Redirect to /dev/null
        os.dup2(devnull_fd, 1)
        os.dup2(devnull_fd, 2)

        # Also silence Python’s sys streams
        with open(os.devnull, "w") as dn, \
             contextlib.redirect_stdout(dn), \
             contextlib.redirect_stderr(dn):
            yield

    finally:
        # Restore original file descriptors
        os.dup2(orig_stdout_fd, 1)
        os.dup2(orig_stderr_fd, 2)
        os.close(orig_stdout_fd)
        os.close(orig_stderr_fd)
        os.close(devnull_fd)

# Using contextmanager for nice `with...` syntax
@contextmanager
def open_file(filename, treename):
    """Open a ROOT file and return (TFile, TTree). Closes upon end"""
    tf = cppyy.gbl.TFile.Open(filename)
    if not tf or tf.IsZombie():
        raise OSError(f"Failed to open ROOT file: {filename}")
    try:
        tt = tf.Get(treename)
        yield tf, tt
    finally:
        tf.Close()

def attach_record(tree, branch_name=None):
    """
    Attach a genie::NtpMCEventRecord branch.

    Parameters
    ----------
    tree : TTree
    branch_name : str or None
        If given, attach directly. If None, auto-discover the first branch
        whose class is 'genie::NtpMCEventRecord'.

    Returns
    -------
    (tree, record, branch_name)
    """
    import re
    from ._genie_bindings import NtpMCEventRecord

    if branch_name is None:
        # scan tree for NtpMCEventRecord branches
        blist = tree.GetListOfBranches()
        for b in blist:
            cname = b.GetClassName()
            if re.search("NtpMCEventRecord", cname):
                branch_name = b.GetName()
                break
        if branch_name is None:
            raise RuntimeError("No NtpMCEventRecord branch found.")

    # We need to define a tiny struct here to construct a pointer to the entry
    # because the CAFs stored it as a pointer
    if not hasattr(cppyy.gbl, "NtpMCEventRecordHolder"):
        cppyy.cppdef("""
        struct NtpMCEventRecordHolder {
        genie::NtpMCEventRecord * rec = nullptr;
        };
        """)

    with silence_cpp(): # No GENIE banners please - these are via cout
        #rec = cppyy.gbl.genie.NtpMCEventRecord()
        holder = cppyy.gbl.NtpMCEventRecordHolder()
    #res = tree.SetBranchAddress(branch_name, rec)
    res = tree.SetBranchAddress(branch_name, holder)
    tree._ntp_ptr = holder
    return tree, holder, branch_name

def attach_record_with_flux(tree, branch_name=None, flux_name=None):
    """
    Attach a genie::NtpMCEventRecord branch alongside the bsim::Dk2Nu flux record

    Parameters
    ----------
    tree : TTree
    branch_name : str or None
    flux_name   : str or None
        If given, attach directly. If None, auto-discover the first branch
        whose class is 'genie::NtpMCEventRecord' (branch) of 'bsim::Dk2Nu' (flux).

    Returns
    -------
    (tree, event record, flux record)
    """
    import re
    from ._genie_bindings import NtpMCEventRecord, Dk2Nu

    if branch_name is None:
        # scan tree for NtpMCEventRecord branches
        blist = tree.GetListOfBranches()
        for b in blist:
            cname = b.GetClassName()
            if re.search("NtpMCEventRecord", cname):
                branch_name = b.GetName()
                break
        if branch_name is None:
            raise RuntimeError("No NtpMCEventRecord branch found.")
    if flux_name is None:
        # scan tree for Dk2Nu branches
        blist = tree.GetListOfBranches()
        for b in blist:
            cname = b.GetClassName()
            if re.search("Dk2Nu", cname):
                flux_name = b.GetName()
                break
        if flux_name is None:
            raise RuntimeError("No NtpMCEventRecord branch found.")

    # We need to define a tiny struct here to construct a pointer to the entry
    if not hasattr(cppyy.gbl, "NtpMCEventRecordHolder"):
        cppyy.cppdef("""
        struct NtpMCEventRecordHolder {
        genie::NtpMCEventRecord * rec = nullptr;
        };
        """)
    if not hasattr(cppyy.gbl, "Dk2NuHolder"):
        cppyy.cppdef("""
        struct Dk2NuHolder {
        bsim::Dk2Nu * dk2nu = nullptr;
        };
        """)

    with silence_cpp(): # No GENIE banners please - these are via cout
        #rec = cppyy.gbl.genie.NtpMCEventRecord()
        holder  = cppyy.gbl.NtpMCEventRecordHolder()
        fholder = cppyy.gbl.Dk2NuHolder()
    #res = tree.SetBranchAddress(branch_name, rec)
    res = tree.SetBranchAddress(branch_name, holder)
    flx = tree.SetBranchAddress(flux_name, fholder)
    tree._ntp_ptr = holder
    tree._flx_ptr = fholder
    return tree, holder, fholder