import pandas as pd
import numpy as np
import time
from typing import Optional, Dict, List, Tuple, Iterator, Any


# Reuse detector naming and helper math from utils to stay consistent
from .utils import all_trackers, all_calos, all_digi_calos, _calculate_R, _calculate_theta, _calculate_eta, augment_particles, condense_calo_dataframes
import logging
logger = logging.getLogger(__name__)


class EDM4hepEventBatch:
    """Efficient multi-event loader with on-demand, per-collection IO.

    - Lazy per-collection: nothing is read until a getter is called.
    - Only branches necessary for the requested collection are loaded.
    - All returned DataFrames include an 'event_id' column for filtering.
    """

    def __init__(
        self,
        file_path: str,
        events: Optional[Tuple[int, int]] = None,
        detector_params: Optional[Dict[str, float]] = None,
        full_load: bool = False,
        condense_calo: bool = True,
    ):
        self.file_path = file_path
        self._events = events
        self.detector_params = detector_params
        self.full_load = full_load
        self.condense_calo = condense_calo
        logger.debug(f"EDM4hepEventBatch init: file={file_path}, events={events}, full_load={full_load}, condense_calo={condense_calo}")

        # Internal caches for loaded DataFrames
        self._loaded: set = set()
        self._event_ids: Optional[List[int]] = None
        self._event_indices: Optional[List[int]] = None
        self._particles_df: pd.DataFrame = pd.DataFrame()
        self._parents_df: pd.DataFrame = pd.DataFrame()
        self._daughters_df: pd.DataFrame = pd.DataFrame()
        self._tracker_hits_df: pd.DataFrame = pd.DataFrame()
        self._calo_hits_df: pd.DataFrame = pd.DataFrame()
        self._calo_contributions_df: pd.DataFrame = pd.DataFrame()
        self._digi_calo_hits_df: pd.DataFrame = pd.DataFrame()
        self._event_headers_df: pd.DataFrame = pd.DataFrame()

    # ------------------------- Public Accessors -------------------------

    def _ensure_particles_have_particle_id(self) -> None:
        """Ensure `self._particles_df` has a per-event `particle_id` column.

        Inputs:
            None

        Outputs:
            None. Mutates `self._particles_df` in-place (adds `particle_id` if missing).
        """
        if self._particles_df.empty:
            return
        if "particle_id" in self._particles_df.columns:
            return

        # Awkward->pandas commonly yields a `subentry` column after reset_index()
        if "subentry" in self._particles_df.columns:
            self._particles_df["particle_id"] = self._particles_df["subentry"].astype(int)
            return

        # Fallback: assume current row order is the particle order within each event
        if "event_id" in self._particles_df.columns:
            self._particles_df["particle_id"] = (
                self._particles_df.groupby("event_id", sort=False).cumcount().astype(int)
            )
            return

        # Last resort: single-event table without `event_id`
        self._particles_df["particle_id"] = np.arange(len(self._particles_df), dtype=int)

    def _add_daughter_ids_column(self, particles_df: pd.DataFrame) -> pd.DataFrame:
        """Attach a `daughter_ids` column to a particles DataFrame.

        Notes:
            - Requires `full_load=True` at batch construction time, because we need
              `daughters_begin/end` and the `_MCParticles_daughters` link table.
            - `daughter_ids` contains per-event particle indices (ints).

        Inputs:
            particles_df: Particle rows, must include `event_id`, `daughters_begin`, `daughters_end`.

        Outputs:
            Copy of `particles_df` with an additional `daughter_ids` object column (list[int]).
        """
        if not self.full_load:
            raise ValueError(
                "Daughters are not available unless EDM4hepEventBatch(full_load=True)."
            )

        if particles_df.empty:
            out = particles_df.copy()
            out["daughter_ids"] = pd.Series(dtype=object)
            return out

        required_cols = {"event_id", "daughters_begin", "daughters_end"}
        missing = required_cols.difference(set(particles_df.columns))
        if missing:
            raise KeyError(f"Missing required columns for daughters: {sorted(missing)}")

        daughters_df = self.get_daughters_df(event_id=None)
        if daughters_df.empty:
            out = particles_df.copy()
            out["daughter_ids"] = [[] for _ in range(len(out))]
            return out

        out = particles_df.copy()
        out["daughter_ids"] = [[] for _ in range(len(out))]

        # Compute per-event slices into the daughters link table.
        for eid, p_ev in out.groupby("event_id", sort=False):
            d_ev = daughters_df[daughters_df["event_id"] == eid]
            if d_ev.empty:
                continue

            daughter_ids_flat = d_ev["particle_id"].to_numpy()
            begins = p_ev["daughters_begin"].to_numpy(dtype=int)
            ends = p_ev["daughters_end"].to_numpy(dtype=int)

            lists = []
            n_links = len(daughter_ids_flat)
            for b, e in zip(begins, ends):
                # Guard against invalid or unset ranges
                if e <= b or b < 0 or e < 0 or b > n_links or e > n_links:
                    lists.append([])
                else:
                    lists.append(daughter_ids_flat[b:e].tolist())

            out.loc[p_ev.index, "daughter_ids"] = pd.Series(lists, index=p_ev.index, dtype=object)

        return out

    def get_event_ids(self) -> List[int]:
        self._ensure_loaded("event_ids")
        return self._event_ids
    
    def get_event_indices(self) -> List[int]:
        self._ensure_loaded("event_ids")
        return self._event_indices

    def get_particles_df(
        self,
        event_id: Optional[int] = None,
        include_daughters: bool = False,
    ) -> pd.DataFrame:
        """Return the particles table (optionally with daughter lists).

        Inputs:
            event_id: If provided, filter to a single event id.
            include_daughters: If True, add a `daughter_ids` list column.
                Requires constructing the batch with `full_load=True`.

        Outputs:
            Pandas DataFrame of particles.
        """
        self._ensure_loaded("particles")
        self._ensure_particles_have_particle_id()

        # Only augment if not already augmented (check for key columns)
        df = self._particles_df
        if not df.empty and not {'num_tracker_hits', 'num_calo_hits'}.issubset(df.columns):
            # Augment particles with derived fields and counts lazily
            t0_augment = time.time()
            df = augment_particles(
                self._particles_df,
                tracker_hits_df=(self._tracker_hits_df if not self._tracker_hits_df.empty else None),
                calo_contrib_df=(self._calo_contributions_df if not self._calo_contributions_df.empty else None),
                detector_params=self.detector_params,
                logger=logger,
            )
            t_augment = time.time() - t0_augment
            logger.debug(f"⏱️  augment_particles: {t_augment:.4f}s ({len(df)} particles)")
            self._particles_df = df  # cache augmented

        if event_id is not None:
            df = df[df["event_id"] == event_id]

        if include_daughters:
            df = self._add_daughter_ids_column(df)
        else:
            df = df.copy()

        return df

    def get_parents_df(self, event_id: Optional[int] = None) -> pd.DataFrame:
        self._ensure_loaded("particles")  # parents/daughters loaded together
        df = self._parents_df
        if event_id is None:
            return df.copy()
        return df[df["event_id"] == event_id].copy()

    def get_daughters_df(self, event_id: Optional[int] = None) -> pd.DataFrame:
        self._ensure_loaded("particles")
        df = self._daughters_df
        if event_id is None:
            return df.copy()
        return df[df["event_id"] == event_id].copy()

    def get_tracker_hits_df(self, event_id: Optional[int] = None) -> pd.DataFrame:
        self._ensure_loaded("tracker_hits")
        df = self._tracker_hits_df
        if event_id is None:
            return df.copy()
        return df[df["event_id"] == event_id].copy()

    def get_calo_hits_df(self, event_id: Optional[int] = None) -> pd.DataFrame:
        self._ensure_loaded("calo_hits")
        df = self._calo_hits_df
        if event_id is None:
            return df.copy()
        return df[df["event_id"] == event_id].copy()

    def get_calo_contributions_df(self, event_id: Optional[int] = None) -> pd.DataFrame:
        self._ensure_loaded("calo_contributions")
        df = self._calo_contributions_df
        if event_id is None:
            return df.copy()
        return df[df["event_id"] == event_id].copy()

    def get_digi_calo_hits_df(self, event_id: Optional[int] = None) -> pd.DataFrame:
        self._ensure_loaded("digi_calo_hits")
        df = self._digi_calo_hits_df
        if event_id is None:
            return df.copy()
        return df[df["event_id"] == event_id].copy()

    def condense_calo(self) -> None:
        """Condense calorimeter contributions by (cellID, particle_id).
        
        This method can be called after initialization to condense duplicate
        calorimeter contributions. It will:
        1. Ensure calo hits and contributions are loaded
        2. Apply condensation using the condense_calo_dataframes utility
        3. Update the internal cached dataframes
        
        Safe to call multiple times (will only condense once).
        """
        # Ensure data is loaded
        self._ensure_loaded("calo_hits")
        self._ensure_loaded("calo_contributions")
        
        if self._calo_hits_df.empty or self._calo_contributions_df.empty:
            logger.debug("No calo data to condense (empty dataframes)")
            return
        
        # Check if already condensed (simple heuristic: if we've done this before)
        # We could add a flag, but for now just check if the operation makes sense
        t0_condense = time.time()
        logger.debug("Condensing calorimeter contributions...")
        
        self._calo_hits_df, self._calo_contributions_df = condense_calo_dataframes(
            self._calo_hits_df, self._calo_contributions_df
        )
        
        t_condense = time.time() - t0_condense
        logger.debug(f"  ↳ condensation complete: {t_condense:.4f}s ({len(self._calo_contributions_df)} condensed rows)")

    def iter_event_ids(self) -> Iterator[int]:
        # Prefer particles for enumerating event ids, fall back to hits
        if not self._particles_df.empty:
            for eid in self._particles_df["event_id"].unique():
                yield int(eid)
            return
        if not self._tracker_hits_df.empty:
            for eid in self._tracker_hits_df["event_id"].unique():
                yield int(eid)
            return
        if not self._calo_hits_df.empty:
            for eid in self._calo_hits_df["event_id"].unique():
                yield int(eid)

    # ------------------------- Internal Loaders -------------------------

    def _ensure_loaded(self, which: str) -> None:
        if which in self._loaded:
            return
        if which == "particles":
            self._load_particles_and_links()
        elif which == "tracker_hits":
            self._load_tracker_hits()
        elif which == "calo_hits":
            self._load_calo_hits()
        elif which == "calo_contributions":
            self._load_calo_contributions()
        elif which == "digi_calo_hits":
            self._load_digi_calo_hits()
        elif which == "event_ids":
            self._load_event_ids()
        else:
            raise ValueError(f"Unknown collection '{which}'")
        self._loaded.add(which)

    def _open_tree(self):
        import uproot
        f = uproot.open(self.file_path)
        if "events" not in f:
            f.close()
            raise ValueError(f"'events' tree not found in {self.file_path}")
        return f, f["events"]

    def _load_event_ids(self) -> None:
        t0 = time.time()
        f, tree = self._open_tree()
        
        self._event_ids = tree["EventHeader"]["EventHeader.eventNumber"].arrays(library="np")['EventHeader.eventNumber']
        self._event_indices = list(range(len(self._event_ids)))
        f.close()
        
        elapsed = time.time() - t0
        logger.debug(f"⏱️  _load_event_ids: {elapsed:.4f}s ({len(self._event_ids)} events)")

    def _load_particles_and_links(self) -> None:
        import awkward as ak

        t0_total = time.time()
        f, tree = self._open_tree()
        try:
            start = stop = None
            if self._events is not None:
                start, stop = self._events
            logger.debug(f"_load_particles_and_links: start={start}, stop={stop}")
            
            t0_particles = time.time()

            # Core set of branches for particles
            particle_branches = [
                "MCParticles.PDG",
                "MCParticles.simulatorStatus",
                "MCParticles.charge",
                "MCParticles.time",
                "MCParticles.mass",
                "MCParticles.vertex.x",
                "MCParticles.vertex.y",
                "MCParticles.vertex.z",
                "MCParticles.momentum.x",
                "MCParticles.momentum.y",
                "MCParticles.momentum.z",
                "MCParticles.parents_begin",
                "MCParticles.parents_end",
            ]
            
            # Add optional branches if full_load=True
            if self.full_load:
                particle_branches.extend([
                    "MCParticles.generatorStatus",
                    "MCParticles.endpoint.x",
                    "MCParticles.endpoint.y",
                    "MCParticles.endpoint.z",
                    "MCParticles.daughters_begin",
                    "MCParticles.daughters_end",
                ])

            arrays = tree.arrays(
                particle_branches,
                entry_start=start,
                entry_stop=stop,
                library="ak",
            )
            p_df = ak.to_dataframe(arrays)

            # Bring event id into a flat column named 'event_id'
            if isinstance(p_df.index, pd.MultiIndex) and "entry" in p_df.index.names:
                p_df = p_df.reset_index()
                p_df = p_df.rename(columns={"entry": "event_id"})
            elif "entry" in p_df.columns:
                p_df = p_df.rename(columns={"entry": "event_id"})
            # Align event_id to absolute entries when using entry_start/stop (range)
            if start is not None and stop is not None and not p_df.empty and "event_id" in p_df:
                evmin = int(p_df["event_id"].min())
                evmax = int(p_df["event_id"].max())
                expected_max = (stop - start - 1)
                if evmin >= 0 and evmax <= expected_max:
                    p_df["event_id"] = p_df["event_id"] + int(start)

            # Rename to match internal conventions used elsewhere in this repo
            rename_cols = {
                "MCParticles.PDG": "PDG",
                "MCParticles.simulatorStatus": "simulatorStatus",
                "MCParticles.charge": "charge",
                "MCParticles.time": "time",
                "MCParticles.mass": "mass",
                "MCParticles.vertex.x": "vx",
                "MCParticles.vertex.y": "vy",
                "MCParticles.vertex.z": "vz",
                "MCParticles.momentum.x": "px",
                "MCParticles.momentum.y": "py",
                "MCParticles.momentum.z": "pz",
                "MCParticles.parents_begin": "parents_begin",
                "MCParticles.parents_end": "parents_end",
            }
            
            # Add optional column renames if they were loaded
            if self.full_load:
                rename_cols.update({
                    "MCParticles.generatorStatus": "generatorStatus",
                    "MCParticles.endpoint.x": "endpoint_x",
                    "MCParticles.endpoint.y": "endpoint_y",
                    "MCParticles.endpoint.z": "endpoint_z",
                    "MCParticles.daughters_begin": "daughters_begin",
                    "MCParticles.daughters_end": "daughters_end",
                })
            p_df = p_df.rename(columns=rename_cols)
            logger.debug(f"particles columns after rename: {list(p_df.columns)} rows={len(p_df)}")
            
            t_particles = time.time() - t0_particles
            logger.debug(f"  ↳ particles loaded: {t_particles:.4f}s ({len(p_df)} rows)")

            # Derived kinematics, vectorized
            if not p_df.empty:
                p_df["vr"] = np.sqrt(p_df.vx**2 + p_df.vy**2)
                if self.full_load and "endpoint_x" in p_df.columns:
                    p_df["endpoint_r"] = np.sqrt(p_df.endpoint_x**2 + p_df.endpoint_y**2)
                p_df["energy"] = np.sqrt(p_df.px**2 + p_df.py**2 + p_df.pz**2 + p_df.mass**2)
                p_df["kinetic_energy"] = p_df["energy"] - p_df["mass"]

            self._particles_df = p_df

            # Parents and daughters link tables (optional, if present)
            t0_links = time.time()
            link_tables: Dict[str, pd.DataFrame] = {}
            link_names = ["_MCParticles_parents"]
            if self.full_load:
                link_names.append("_MCParticles_daughters")
            
            for link_name in link_names:
                if link_name not in tree:
                    link_tables[link_name] = pd.DataFrame()
                    continue
                link_arrays = tree[link_name].arrays(
                    [f"{link_name}.index", f"{link_name}.collectionID"],
                    entry_start=start,
                    entry_stop=stop,
                    library="ak",
                )
                ldf = ak.to_dataframe(link_arrays)
                if isinstance(ldf.index, pd.MultiIndex) and "entry" in ldf.index.names:
                    ldf = ldf.reset_index().rename(columns={"entry": "event_id"})
                # Align link event ids to absolute entries when using range
                if start is not None and stop is not None and not ldf.empty and "event_id" in ldf:
                    evmin = int(ldf["event_id"].min())
                    evmax = int(ldf["event_id"].max())
                    expected_max = (stop - start - 1)
                    if evmin >= 0 and evmax <= expected_max:
                        ldf["event_id"] = ldf["event_id"] + int(start)
                ldf = ldf.rename(columns={
                    f"{link_name}.index": "particle_id",
                    f"{link_name}.collectionID": "collectionID",
                })
                link_tables[link_name] = ldf

            self._parents_df = link_tables.get("_MCParticles_parents", pd.DataFrame())
            self._daughters_df = link_tables.get("_MCParticles_daughters", pd.DataFrame())
            
            t_links = time.time() - t0_links
            logger.debug(f"  ↳ links loaded: {t_links:.4f}s (parents={len(self._parents_df)}, daughters={len(self._daughters_df)})")
            
            t_total = time.time() - t0_total
            logger.debug(f"⏱️  _load_particles_and_links: {t_total:.4f}s total")
        finally:
            f.close()

    def _load_tracker_hits(self) -> None:
        import awkward as ak

        t0_total = time.time()
        f, tree = self._open_tree()
        try:
            start = stop = None
            if self._events is not None:
                start, stop = self._events
            hit_frames: List[pd.DataFrame] = []
            logger.debug(f"_load_tracker_hits: start={start}, stop={stop}")
            
            detector_timings = {}

            for det in all_trackers:
                # Skip detectors absent in this file
                if det not in tree:
                    continue
                
                t0_det = time.time()

                sub = tree[det]
                edep_branch = f"{det}.EDep" if f"{det}.EDep" in sub else (f"{det}.eDep" if f"{det}.eDep" in sub else None)

                # Core branches always loaded
                branches = [
                    f"{det}.cellID",
                    f"{det}.time",
                    f"{det}.position.x",
                    f"{det}.position.y",
                    f"{det}.position.z",
                ]
                
                # Add optional branches if full_load=True
                if self.full_load:
                    branches.extend([
                        f"{det}.momentum.x",
                        f"{det}.momentum.y",
                        f"{det}.momentum.z",
                        f"{det}.pathLength",
                        f"{det}.quality",
                    ])
                    if edep_branch is not None:
                        branches.append(edep_branch)

                t0_read = time.time()
                arrays = sub.arrays(
                    branches,
                    entry_start=start,
                    entry_stop=stop,
                    library="ak",
                )
                t_read = time.time() - t0_read
                
                t0_to_df = time.time()
                df = ak.to_dataframe(arrays)
                t_to_df = time.time() - t0_to_df
                
                # Bring event id flat
                if isinstance(df.index, pd.MultiIndex) and "entry" in df.index.names:
                    df = df.reset_index().rename(columns={"entry": "event_id"})

                # Align event_id to absolute entries when using entry_start/stop (range)
                if start is not None and stop is not None and not df.empty and "event_id" in df:
                    evmin = int(df["event_id"].min())
                    evmax = int(df["event_id"].max())
                    expected_max = (stop - start - 1)
                    if evmin >= 0 and evmax <= expected_max:
                        df["event_id"] = df["event_id"] + int(start)

                # Rename common columns and add detector label
                rename_map = {
                    f"{det}.cellID": "cellID",
                    f"{det}.time": "time",
                    f"{det}.position.x": "x",
                    f"{det}.position.y": "y",
                    f"{det}.position.z": "z",
                }
                
                # Add optional column renames if they were loaded
                if self.full_load:
                    rename_map.update({
                        f"{det}.momentum.x": "px",
                        f"{det}.momentum.y": "py",
                        f"{det}.momentum.z": "pz",
                        f"{det}.pathLength": "pathLength",
                        f"{det}.quality": "quality",
                    })
                    if edep_branch is not None:
                        rename_map[edep_branch] = "EDep"
                df = df.rename(columns=rename_map)
                df["detector"] = det
                logger.debug(f"tracker {det}: rows={len(df)} cols={list(df.columns)}")

                # Particle link indices (optional, try both conventions)
                t0_links = time.time()
                particle_id_series = None
                try:
                    link = tree[f"_{det}_MCParticle"].arrays([f"_{det}_MCParticle.index"], entry_start=start, entry_stop=stop, library="ak")
                    link_df = ak.to_dataframe(link)
                    if isinstance(link_df.index, pd.MultiIndex) and "entry" in link_df.index.names:
                        link_df = link_df.reset_index().rename(columns={"entry": "event_id"})
                    particle_id_series = link_df[f"_{det}_MCParticle.index"]
                except Exception:
                    try:
                        link = tree[f"_{det}_particle"].arrays([f"_{det}_particle.index"], entry_start=start, entry_stop=stop, library="ak")
                        link_df = ak.to_dataframe(link)
                        if isinstance(link_df.index, pd.MultiIndex) and "entry" in link_df.index.names:
                            link_df = link_df.reset_index().rename(columns={"entry": "event_id"})
                        particle_id_series = link_df[f"_{det}_particle.index"]
                    except Exception:
                        particle_id_series = None
                t_links = time.time() - t0_links

                if particle_id_series is not None and len(particle_id_series) == len(df):
                    df["particle_id"] = particle_id_series.values

                # Derived quantities
                t0_derived = time.time()
                if not df.empty:
                    df["r"] = _calculate_R(df["x"], df["y"])
                    df["R"] = _calculate_R(df["x"], df["y"], df["z"])
                    df["phi"] = np.arctan2(df["y"], df["x"])
                    df["theta"] = _calculate_theta(df["r"], df["z"])
                    df["eta"] = _calculate_eta(df["theta"])
                    # Only calculate pt if momentum columns are present
                    if self.full_load and "px" in df.columns and "py" in df.columns:
                        df["pt"] = np.sqrt(df["px"]**2 + df["py"]**2)
                t_derived = time.time() - t0_derived

                # Post-filter events if necessary
                hit_frames.append(df)
                
                t_det = time.time() - t0_det
                detector_timings[det] = {
                    'total': t_det,
                    'read': t_read,
                    'to_df': t_to_df,
                    'links': t_links,
                    'derived': t_derived,
                    'rows': len(df)
                }

            # Log per-detector timings
            for det, timings in detector_timings.items():
                logger.debug(f"  ↳ {det}: {timings['total']:.4f}s ({timings['rows']} rows)")
                logger.debug(f"      • read: {timings['read']:.4f}s, to_df: {timings['to_df']:.4f}s, links: {timings['links']:.4f}s, derived: {timings['derived']:.4f}s")

            t0_concat = time.time()
            self._tracker_hits_df = pd.concat(hit_frames, ignore_index=True) if hit_frames else pd.DataFrame()
            t_concat = time.time() - t0_concat
            logger.debug(f"tracker_hits combined rows={len(self._tracker_hits_df)}")
            
            logger.debug(f"  ↳ concat: {t_concat:.4f}s")
            t_total = time.time() - t0_total
            logger.debug(f"⏱️  _load_tracker_hits: {t_total:.4f}s ({len(self._tracker_hits_df)} hits)")
        finally:
            f.close()

    def _load_calo_hits(self) -> None:
        import awkward as ak

        t0_total = time.time()
        f, tree = self._open_tree()
        try:
            start = stop = None
            if self._events is not None:
                start, stop = self._events
            hit_frames: List[pd.DataFrame] = []
            logger.debug(f"_load_calo_hits: start={start}, stop={stop}")
            
            detector_timings = {}

            for det in all_calos:
                if det not in tree:
                    continue
                
                t0_det = time.time()
                
                sub = tree[det]
                branches = [
                    f"{det}.cellID",
                    f"{det}.energy",
                    f"{det}.position.x",
                    f"{det}.position.y",
                    f"{det}.position.z",
                    f"{det}.contributions_begin",
                    f"{det}.contributions_end",
                ]

                t0_read = time.time()
                arrays = sub.arrays(
                    branches,
                    entry_start=start,
                    entry_stop=stop,
                    library="ak",
                )
                t_read = time.time() - t0_read
                
                t0_to_df = time.time()
                df = ak.to_dataframe(arrays)
                t_to_df = time.time() - t0_to_df
                
                if isinstance(df.index, pd.MultiIndex) and "entry" in df.index.names:
                    df = df.reset_index().rename(columns={"entry": "event_id"})
                # Align event_id to absolute entries when using entry_start/stop (range)
                if start is not None and stop is not None and not df.empty and "event_id" in df:
                    evmin = int(df["event_id"].min())
                    evmax = int(df["event_id"].max())
                    expected_max = (stop - start - 1)
                    if evmin >= 0 and evmax <= expected_max:
                        df["event_id"] = df["event_id"] + int(start)
                
                df = df.rename(columns={
                    f"{det}.cellID": "cellID",
                    f"{det}.energy": "energy",
                    f"{det}.position.x": "x",
                    f"{det}.position.y": "y",
                    f"{det}.position.z": "z",
                    f"{det}.contributions_begin": "contribution_begin",
                    f"{det}.contributions_end": "contribution_end",
                })
                df["detector"] = det
                logger.debug(f"calo {det}: rows={len(df)} cols={list(df.columns)}")

                # Derived geometry
                t0_derived = time.time()
                if not df.empty:
                    df["r"] = _calculate_R(df["x"], df["y"])
                    df["R"] = _calculate_R(df["x"], df["y"], df["z"])
                    df["phi"] = np.arctan2(df["y"], df["x"])
                    df["theta"] = _calculate_theta(df["r"], df["z"])
                    df["eta"] = _calculate_eta(df["theta"])
                t_derived = time.time() - t0_derived

                hit_frames.append(df)
                
                t_det = time.time() - t0_det
                detector_timings[det] = {
                    'total': t_det,
                    'read': t_read,
                    'to_df': t_to_df,
                    'derived': t_derived,
                    'rows': len(df)
                }

            # Log per-detector timings
            for det, timings in detector_timings.items():
                logger.debug(f"  ↳ {det}: {timings['total']:.4f}s ({timings['rows']} rows)")
                logger.debug(f"      • read: {timings['read']:.4f}s, to_df: {timings['to_df']:.4f}s, derived: {timings['derived']:.4f}s")

            t0_concat = time.time()
            self._calo_hits_df = pd.concat(hit_frames, ignore_index=True) if hit_frames else pd.DataFrame()
            t_concat = time.time() - t0_concat
            logger.debug(f"calo_hits combined rows={len(self._calo_hits_df)}")
            
            # Fix contribution_begin/end indices to be global within each event
            # ROOT gives us per-detector indices, but after concat we need cross-detector indices
            import os
            import psutil

            def _log_resource_usage(prefix=""):
                try:
                    process = psutil.Process(os.getpid())
                    mem_info = process.memory_info()
                    rss_mb = mem_info.rss / (1024 ** 2)
                    vms_mb = mem_info.vms / (1024 ** 2)
                    cpu_percent = process.cpu_percent(interval=0.01)
                    logger.debug(f"{prefix}Resource usage: RSS={rss_mb:.2f} MB, VMS={vms_mb:.2f} MB, CPU={cpu_percent:.2f}%")
                except Exception as res_e:
                    logger.warning(f"{prefix}Resource usage logging failed: {res_e}")

            if not self._calo_hits_df.empty and 'event_id' in self._calo_hits_df.columns:
                t0_fix_indices = time.time()
                logger.debug("  ↳ Fixing contribution indices to be global within each event...")
                _log_resource_usage("Before fixing indices: ")

                # Process each event separately
                fixed_hits = []
                for event_id in sorted(self._calo_hits_df['event_id'].unique()):
                    event_hits = self._calo_hits_df[self._calo_hits_df['event_id'] == event_id]
                    logger.debug(f"  ↳ event {event_id} has {len(event_hits)} hits")
                    _log_resource_usage(f"Event {event_id}: ")

                    # Track offset within this event
                    offset = 0
                    det_frames = []
                    
                    for det in all_calos:

                        det_hits = event_hits[event_hits['detector'] == det].copy()
                        logger.debug(f"  ↳ event {event_id} detector {det} has {len(det_hits)} hits")
                        if det_hits.empty:
                            continue
                        
                        # Shift this detector's indices by the offset
                        det_hits['contribution_begin'] = det_hits['contribution_begin'] + offset
                        det_hits['contribution_end'] = det_hits['contribution_end'] + offset
                        
                        # Update offset for next detector
                        max_idx = int(det_hits['contribution_end'].max())
                        offset = max_idx
                        
                        det_frames.append(det_hits)
                        _log_resource_usage(f"    Detector {det}: ")

                    if det_frames:
                        fixed_hits.append(pd.concat(det_frames, ignore_index=True))
                    logger.debug(f"  ↳ fixed hits: {len(fixed_hits)}")
                _log_resource_usage("After fixing indices: ")
                self._calo_hits_df = pd.concat(fixed_hits, ignore_index=True) if fixed_hits else pd.DataFrame()
                
                t_fix_indices = time.time() - t0_fix_indices
                logger.debug(f"  ↳ fix indices: {t_fix_indices:.4f}s")
            
            logger.debug(f"  ↳ concat: {t_concat:.4f}s")
            t_total = time.time() - t0_total
            logger.debug(f"⏱️  _load_calo_hits: {t_total:.4f}s ({len(self._calo_hits_df)} hits)")
        except Exception as e:
            logger.error(f"Error in _load_calo_hits: {e}")
            raise e
        finally:
            f.close()

    def _load_calo_contributions(self) -> None:
        import awkward as ak

        t0_total = time.time()
        f, tree = self._open_tree()
        try:
            start = stop = None
            if self._events is not None:
                start, stop = self._events
            contrib_frames: List[pd.DataFrame] = []
            logger.debug(f"_load_calo_contributions: start={start}, stop={stop}")
            
            detector_timings = {}

            for det in all_calos:
                # Require contributions and link
                contrib_coll = f"{det}Contributions"
                link_coll = f"_{det}Contributions_particle"
                if contrib_coll not in tree or link_coll not in tree:
                    continue
                
                t0_det = time.time()

                # Core calo contribution branches
                contrib_branches = [
                    f"{contrib_coll}.energy",
                    f"{contrib_coll}.time",
                ]
                
                # Add optional branches if full_load=True
                if self.full_load:
                    contrib_branches.extend([
                        f"{contrib_coll}.PDG",
                        f"{contrib_coll}.stepPosition.x",
                        f"{contrib_coll}.stepPosition.y",
                        f"{contrib_coll}.stepPosition.z",
                    ])

                # Load contribution arrays
                t0_contrib_read = time.time()
                arrays = tree[contrib_coll].arrays(
                    contrib_branches,
                    entry_start=start,
                    entry_stop=stop,
                    library="ak",
                )
                t_contrib_read = time.time() - t0_contrib_read
                
                # Convert to dataframe
                t0_to_df = time.time()
                df = ak.to_dataframe(arrays)
                t_to_df = time.time() - t0_to_df
                
                # Handle multiindex
                if isinstance(df.index, pd.MultiIndex) and "entry" in df.index.names:
                    df = df.reset_index().rename(columns={"entry": "event_id"})
                # Align event_id to absolute entries when using entry_start/stop (range)
                if start is not None and stop is not None and not df.empty and "event_id" in df:
                    evmin = int(df["event_id"].min())
                    evmax = int(df["event_id"].max())
                    expected_max = (stop - start - 1)
                    if evmin >= 0 and evmax <= expected_max:
                        df["event_id"] = df["event_id"] + int(start)
                
                # Core column renames
                rename_dict = {
                    f"{contrib_coll}.energy": "energy",
                    f"{contrib_coll}.time": "time",
                }
                
                # Add optional column renames if they were loaded
                if self.full_load:
                    rename_dict.update({
                        f"{contrib_coll}.PDG": "PDG",
                        f"{contrib_coll}.stepPosition.x": "step_x",
                        f"{contrib_coll}.stepPosition.y": "step_y",
                        f"{contrib_coll}.stepPosition.z": "step_z",
                    })
                
                df = df.rename(columns=rename_dict)

                # Load and attach particle ids
                t0_link_read = time.time()
                link_arrays = tree[link_coll].arrays([f"{link_coll}.index"], entry_start=start, entry_stop=stop, library="ak")
                t_link_read = time.time() - t0_link_read
                
                t0_link_to_df = time.time()
                link_df = ak.to_dataframe(link_arrays)
                t_link_to_df = time.time() - t0_link_to_df
                
                if isinstance(link_df.index, pd.MultiIndex) and "entry" in link_df.index.names:
                    link_df = link_df.reset_index().rename(columns={"entry": "event_id"})
                # Align link event ids to absolute entries when using range
                if start is not None and stop is not None and not link_df.empty and "event_id" in link_df:
                    evmin = int(link_df["event_id"].min())
                    evmax = int(link_df["event_id"].max())
                    expected_max = (stop - start - 1)
                    if evmin >= 0 and evmax <= expected_max:
                        link_df["event_id"] = link_df["event_id"] + int(start)
                df["particle_id"] = link_df[f"{link_coll}.index"].values if f"{link_coll}.index" in link_df else np.nan
                df["detector"] = det

                contrib_frames.append(df)
                
                t_det = time.time() - t0_det
                detector_timings[det] = {
                    'total': t_det,
                    'contrib_read': t_contrib_read,
                    'contrib_to_df': t_to_df,
                    'link_read': t_link_read,
                    'link_to_df': t_link_to_df,
                    'rows': len(df)
                }

            # Log per-detector timings
            for det, timings in detector_timings.items():
                logger.debug(f"  ↳ {det}: {timings['total']:.4f}s ({timings['rows']} rows)")
                logger.debug(f"      • read contrib: {timings['contrib_read']:.4f}s, to_df: {timings['contrib_to_df']:.4f}s")
                logger.debug(f"      • read links:   {timings['link_read']:.4f}s, to_df: {timings['link_to_df']:.4f}s")
            
            t0_concat = time.time()
            # Reorganize contributions strictly by event, then detector (all_calos order),
            # preserving original within-(event,detector) order from ROOT
            if contrib_frames:
                t0_concat_contribs = time.time()
                logger.debug("  ↳ Reorganizing contributions by event and detector...")

                all_contribs = pd.concat(contrib_frames, ignore_index=True)
                t_concat_contribs = time.time() - t0_concat_contribs
                logger.debug(f"  ↳ concat initial contributions: {t_concat_contribs:.4f}s")
                t0_reorg = time.time()
                if not all_contribs.empty and 'event_id' in all_contribs.columns:
                    self._calo_contributions_df = all_contribs.sort_values('event_id', kind='stable')
                    t_reorg = time.time() - t0_reorg
                    logger.debug(f"  ↳ reorganize: {t_reorg:.4f}s")
                else:
                    self._calo_contributions_df = all_contribs
            else:
                self._calo_contributions_df = pd.DataFrame()

            t_concat = time.time() - t0_concat
            logger.debug(f"calo_contrib combined rows={len(self._calo_contributions_df)}")
            logger.debug(f"  ↳ concat: {t_concat:.4f}s, concat initial contributions: {t_concat_contribs:.4f}s, reorganize: {t_reorg:.4f}s")
            
            self._ensure_loaded("calo_hits")
            if not self._calo_hits_df.empty:
                # First, we need to merge cellID from hits into contributions
                # In batch mode, contributions don't have cellID yet
                # Important: contribution_begin/end are per-event indices, so we must process per event
                t0_merge = time.time()
                
                # Process each event separately to respect contribution_begin/end indices
                # We need to assign hit_index per event to ensure correct mapping
                for event_id in self._calo_hits_df['event_id'].unique():
                    # Get hits and contributions for this event
                    event_hits = self._calo_hits_df[self._calo_hits_df['event_id'] == event_id]
                    event_contrib_mask = self._calo_contributions_df['event_id'] == event_id
                    num_event_contribs = event_contrib_mask.sum()
                    
                    # Build hit_index mapping for this event (contribution_begin/end are relative to this event)
                    num_contribs = event_hits.contribution_end - event_hits.contribution_begin
                    event_hit_index = np.repeat(event_hits.index.values, num_contribs)
                    
                    # Sanity check: lengths should match
                    if len(event_hit_index) != num_event_contribs:
                        logger.warning(f"Event {event_id}: hit_index length {len(event_hit_index)} != contributions {num_event_contribs}")
                    
                    # Assign hit_index to this event's contributions
                    self._calo_contributions_df.loc[event_contrib_mask, 'hit_index'] = event_hit_index
                
                # Merge only cellID from hits (not positions)
                self._calo_contributions_df = self._calo_contributions_df.merge(
                    self._calo_hits_df[['cellID', 'x', 'y', 'z']],
                    left_on='hit_index',
                    right_index=True,
                    how='left'
                )
                
                # Check for NaN cellIDs after merge
                nan_cellids = self._calo_contributions_df['cellID'].isna().sum()
                if nan_cellids > 0:
                    logger.error(f"After merge: {nan_cellids} contributions have NaN cellID!")
                    bad_indices = self._calo_contributions_df[self._calo_contributions_df['cellID'].isna()]['hit_index'].unique()
                    logger.error(f"  Bad hit_indices: {bad_indices[:20]}")
                    logger.error(f"  calo_hits_df index range: {self._calo_hits_df.index.min()} to {self._calo_hits_df.index.max()}")
                
                t_merge = time.time() - t0_merge
                logger.debug(f"  ↳ merge cellID from hits: {t_merge:.4f}s")

            # Optional condensation (if configured during initialization)
            if self.condense_calo and not self._calo_contributions_df.empty:
                t0_condense = time.time()
                logger.debug("Auto-condensing calorimeter contributions (condense_calo=True)...")
                
                self._calo_hits_df, self._calo_contributions_df = condense_calo_dataframes(
                    self._calo_hits_df, self._calo_contributions_df
                )
                
                t_condense = time.time() - t0_condense
                logger.debug(f"  ↳ condensation: {t_condense:.4f}s ({len(self._calo_contributions_df)} condensed rows)")
            
            t_total = time.time() - t0_total
            logger.debug(f"⏱️  _load_calo_contributions: {t_total:.4f}s ({len(self._calo_contributions_df)} contributions)")
        except Exception as e:
            logger.error(f"Error in _load_calo_contributions: {e}")
            raise e
        finally:
            f.close()

    def _load_digi_calo_hits(self) -> None:
        import awkward as ak
        
        t0_total = time.time()
        f, tree = self._open_tree()
        try:
            start = stop = None
            if self._events is not None:
                start, stop = self._events
            hit_frames: List[pd.DataFrame] = []
            logger.debug(f"_load_digi_calo_hits: start={start}, stop={stop}")
            
            for det in all_digi_calos:
                if det not in tree:
                    continue
                sub = tree[det]
                branches = [
                    f"{det}.cellID",
                    f"{det}.energy",
                    f"{det}.energyError",
                    f"{det}.time",
                    f"{det}.position.x",
                    f"{det}.position.y",
                    f"{det}.position.z",
                    f"{det}.type",
                ]
                
                arrays = sub.arrays(
                    branches,
                    entry_start=start,
                    entry_stop=stop,
                    library="ak",
                )
                df = ak.to_dataframe(arrays)
                if isinstance(df.index, pd.MultiIndex) and "entry" in df.index.names:
                    df = df.reset_index().rename(columns={"entry": "event_id"})
                
                # Align event_id to absolute entries
                if start is not None and stop is not None and not df.empty and "event_id" in df:
                    evmin = int(df["event_id"].min())
                    evmax = int(df["event_id"].max())
                    expected_max = (stop - start - 1)
                    if evmin >= 0 and evmax <= expected_max:
                        df["event_id"] = df["event_id"] + int(start)
                
                # Rename columns to match conventions
                df = df.rename(columns={
                    f"{det}.cellID": "cellID",
                    f"{det}.energy": "energy",
                    f"{det}.energyError": "energy_error",
                    f"{det}.time": "time",
                    f"{det}.position.x": "x",
                    f"{det}.position.y": "y",
                    f"{det}.position.z": "z",
                    f"{det}.type": "type",
                })
                df["detector"] = det
                logger.debug(f"digi calo {det}: rows={len(df)} cols={list(df.columns)}")
                
                # Derived geometry (same as truth calo)
                if not df.empty:
                    df["r"] = _calculate_R(df["x"], df["y"])
                    df["R"] = _calculate_R(df["x"], df["y"], df["z"])
                    df["phi"] = np.arctan2(df["y"], df["x"])
                    df["theta"] = _calculate_theta(df["r"], df["z"])
                    df["eta"] = _calculate_eta(df["theta"])
                
                hit_frames.append(df)
            
            self._digi_calo_hits_df = pd.concat(hit_frames, ignore_index=True) if hit_frames else pd.DataFrame()
            logger.debug(f"digi_calo_hits combined rows={len(self._digi_calo_hits_df)}")
            
            t_total = time.time() - t0_total
            logger.debug(f"⏱️  _load_digi_calo_hits: {t_total:.4f}s ({len(self._digi_calo_hits_df)} hits)")
        finally:
            f.close()


