import pandas as pd
import numpy as np
import uproot
from typing import Optional, Tuple
import logging
logger = logging.getLogger(__name__)

# Define detector component lists (copied from edm4hep_utils.py)
pixel_readouts = [
    'PixelBarrelReadout',
    'PixelEndcapReadout',
]
strip_readouts = [
    'ShortStripBarrelReadout',
    'ShortStripEndcapReadout',
    'LongStripBarrelReadout',
    'LongStripEndcapReadout'
]
ecal = [
    'ECalBarrelCollection',
    'ECalEndcapCollection',
]
hcal = [
    'HCalBarrelCollection',
    'HCalEndcapCollection',
]
all_trackers = pixel_readouts + strip_readouts
all_calos = ecal + hcal

digi_ecal = [
    'digiECalBarrelCollection',
    'digiECalEndcapCollection',
]
digi_hcal = [
    'digiHCalBarrelCollection', 
    'digiHCalEndcapCollection',
]

# digi_ecal = [
#     'ECALBarrel',  # Changed from 'digiECalBarrelCollection'
#     'ECALEndcap',  # Changed from 'digiECalEndcapCollection'  (if it exists)
# ]
# digi_hcal = [
#     'HCALBarrel',  # Changed from 'digiHCalBarrelCollection' (if it exists)
#     'HCALEndcap',  # Changed from 'digiHCalEndcapCollection' (if it exists)
# ]

all_digi_calos = digi_ecal + digi_hcal

# --- Helper Functions (copied/adapted from edm4hep_utils.py) ---

def _calculate_R(x, y, z=None):
    """Calculate R (radial distance) from x,y,z coordinates"""
    if z is None:
        return np.sqrt(x**2 + y**2)
    return np.sqrt(x**2 + y**2 + z**2)

def _calculate_theta(r, z):
    """Calculate theta (polar angle) from r,z coordinates"""
    # Add small epsilon to avoid division by zero or issues at z=0
    return np.arctan2(r, z + 1e-15)

def _calculate_eta(theta):
    """Calculate pseudorapidity from theta"""
    # Ensure tan(theta/2) is positive
    tan_theta_half = np.tan(theta / 2)
    # Handle cases close to zero or pi which might result in negative tan_theta_half due to precision
    # or theta values slightly outside [0, pi]
    valid_tan = np.where(tan_theta_half > 1e-15, tan_theta_half, 1e-15)
    eta = -np.log(valid_tan)
    # Handle potential infinities if theta is exactly 0 or pi
    eta = np.where(np.isinf(eta), np.sign(eta) * 1e6, eta) # Replace inf with large number
    return eta


def get_simulator_status_bits(int_val):
    """
    Decodes the simulator status integer into a dictionary of boolean flags.
    Based on edm4hep::MCRecoParticleStatus definition (check specific EDM4hep version).
    Bits are read from right to left (LSB is bit 0).
      static const int BITCreatedInSimulation = 30;
      static const int BITBackscatter = 29 ;
      static const int BITVertexIsNotEndpointOfParent = 28 ;
      static const int BITDecayedInTracker = 27 ;
      static const int BITDecayedInCalorimeter = 26 ;
      static const int BITLeftDetector = 25 ;
      static const int BITStopped = 24 ;
      static const int BITOverlay = 23 ;
    """
    # Ensure input is integer
    status = int(int_val)
    return {
        'created_in_simulation': bool(status & (1 << 30)),
        'backscatter': bool(status & (1 << 29)),
        'vertex_not_endpoint': bool(status & (1 << 28)),
        'decayed_in_tracker': bool(status & (1 << 27)),
        'decayed_in_calorimeter': bool(status & (1 << 26)),
        'has_left_detector': bool(status & (1 << 25)),
        'stopped': bool(status & (1 << 24)),
        'overlay': bool(status & (1 << 23))
    }

# --- DataFrame Building Functions (adapted from edm4hep_utils.py) ---

def _build_particle_df(events_tree, event_idx, full_load: bool = False):
    """Build particle dataframe and parent/daughter link dataframes for a single event"""
    particles = events_tree["MCParticles"].arrays(entry_start=event_idx, entry_stop=event_idx+1)
    
    # Check if event exists / has particles
    if len(particles["MCParticles.PDG"]) == 0 or len(particles["MCParticles.PDG"][0]) == 0:
        empty_cols_particles = ['PDG', 'generatorStatus', 'simulatorStatus', 'charge', 'time', 'mass', 'vx', 'vy', 'vz', 'px', 'py', 'pz', 'endpoint_x', 'endpoint_y', 'endpoint_z', 'parents_begin', 'parents_end', 'daughters_begin', 'daughters_end', 'pt', 'p', 'eta', 'phi']
        empty_cols_links = ['particle_id', 'collectionID']
        return pd.DataFrame(columns=empty_cols_particles), pd.DataFrame(columns=empty_cols_links), pd.DataFrame(columns=empty_cols_links)

    # Core particle data always loaded
    particle_dict = {
        'pdg_id': particles["MCParticles.PDG"][0],
        'simulator_status': particles["MCParticles.simulatorStatus"][0],
        'charge': particles["MCParticles.charge"][0],
        'time': particles["MCParticles.time"][0],
        'mass': particles["MCParticles.mass"][0],
        'vx': particles["MCParticles.vertex.x"][0],
        'vy': particles["MCParticles.vertex.y"][0],
        'vz': particles["MCParticles.vertex.z"][0],
        'px': particles["MCParticles.momentum.x"][0],
        'py': particles["MCParticles.momentum.y"][0],
        'pz': particles["MCParticles.momentum.z"][0],
        'parents_begin': particles["MCParticles.parents_begin"][0],
        'parents_end': particles["MCParticles.parents_end"][0],
    }
    
    # Add optional fields if full_load=True
    if full_load:
        particle_dict.update({
            'generator_status': particles["MCParticles.generatorStatus"][0],
            'endpoint_x': particles["MCParticles.endpoint.x"][0],
            'endpoint_y': particles["MCParticles.endpoint.y"][0],
            'endpoint_z': particles["MCParticles.endpoint.z"][0],
            'daughters_begin': particles["MCParticles.daughters_begin"][0],
            'daughters_end': particles["MCParticles.daughters_end"][0],
        })
    particles_df = pd.DataFrame(particle_dict)
    logger.debug(f"_build_particle_df: event_idx={event_idx}, particles rows={len(particles_df)}, cols={list(particles_df.columns)}")

    # Handle potentially missing parent/daughter link branches gracefully
    try:
        parents = events_tree["_MCParticles_parents"].arrays(entry_start=event_idx, entry_stop=event_idx+1)
        parent_dict = {
            'particle_id': parents["_MCParticles_parents.index"][0],
            'collectionID': parents["_MCParticles_parents.collectionID"][0]
        }
        parents_df = pd.DataFrame(parent_dict)
    except KeyError:
        logger.warning("Branch '_MCParticles_parents' not found.")
        parents_df = pd.DataFrame(columns=['particle_id', 'collectionID'])

    # Only load daughters if full_load=True
    if full_load:
        try:
            daughters = events_tree["_MCParticles_daughters"].arrays(entry_start=event_idx, entry_stop=event_idx+1)
            daughter_dict = {
                'particle_id': daughters["_MCParticles_daughters.index"][0],
                'collectionID': daughters["_MCParticles_daughters.collectionID"][0]
            }
            daughters_df = pd.DataFrame(daughter_dict)
        except KeyError:
            logger.warning("Branch '_MCParticles_daughters' not found.")
            daughters_df = pd.DataFrame(columns=['particle_id', 'collectionID'])
    else:
        daughters_df = pd.DataFrame(columns=['particle_id', 'collectionID'])


    # Calculate derived quantities only if dataframe is not empty
    if not particles_df.empty:
        particles_df['pt'] = np.sqrt(particles_df['px']**2 + particles_df['py']**2)
        particles_df['p'] = np.sqrt(particles_df['pt']**2 + particles_df['pz']**2)
        # Handle pt=0 case for eta calculation
        particles_df['eta'] = np.arcsinh(particles_df['pz'] / np.where(particles_df['pt'] == 0, 1e-15, particles_df['pt']))
        particles_df['phi'] = np.arctan2(particles_df['py'], particles_df['px'])

        # Ensure index columns have the correct data types
        for col in ['parents_begin', 'parents_end', 'daughters_begin', 'daughters_end']:
             if col in particles_df: particles_df[col] = particles_df[col].astype(int)
        for col in ['particle_id', 'collectionID']:
            if col in parents_df: parents_df[col] = parents_df[col].astype(int)
            if col in daughters_df: daughters_df[col] = daughters_df[col].astype(int)

    return particles_df, parents_df, daughters_df

def _process_single_tracker(events_tree, event_idx, detector_name, full_load: bool = False):
    """Process hits and links for a single tracker detector."""
    hits = events_tree[detector_name].arrays(entry_start=event_idx, entry_stop=event_idx+1)
    try:
        particle_links = events_tree[f"_{detector_name}_MCParticle"].arrays(entry_start=event_idx, entry_stop=event_idx+1)
    except:
        particle_links = events_tree[f"_{detector_name}_particle"].arrays(entry_start=event_idx, entry_stop=event_idx+1)

    # Check if event exists / has hits for this detector
    if len(hits[f"{detector_name}.cellID"]) == 0:
        return pd.DataFrame()

    # Core hit data always loaded
    hit_dict = {
        'cellID': hits[f"{detector_name}.cellID"][0],
        'time': hits[f"{detector_name}.time"][0],
        'x': hits[f"{detector_name}.position.x"][0],
        'y': hits[f"{detector_name}.position.y"][0],
        'z': hits[f"{detector_name}.position.z"][0],
    }

    # Add optional fields if full_load=True
    if full_load:
        hit_dict.update({
            'px': hits[f"{detector_name}.momentum.x"][0],
            'py': hits[f"{detector_name}.momentum.y"][0],
            'pz': hits[f"{detector_name}.momentum.z"][0],
            'pathLength': hits[f"{detector_name}.pathLength"][0],
            'quality': hits[f"{detector_name}.quality"][0],
        })
        
        # Handle different EDep naming conventions
        if f"{detector_name}.EDep" in events_tree[detector_name]:
            hit_dict['EDep'] = hits[f"{detector_name}.EDep"][0]
        elif f"{detector_name}.eDep" in events_tree[detector_name]:
             hit_dict['EDep'] = hits[f"{detector_name}.eDep"][0]
        else:
            hit_dict['EDep'] = np.nan # Or some default value

    # Try both particle link formats
    try:
        hit_dict['particle_id'] = particle_links[f"_{detector_name}_MCParticle.index"][0]
    except:
        hit_dict['particle_id'] = particle_links[f"_{detector_name}_particle.index"][0]

    hits_df = pd.DataFrame(hit_dict)

    return hits_df


def _build_tracker_df(events_tree, event_idx, full_load: bool = False):
    """Build combined dataframe for tracker hits for a single event"""
    all_hits_dfs = []

    for det in all_trackers:
        if det not in events_tree:
            continue

        hits_df = _process_single_tracker(events_tree, event_idx, det, full_load)
        hits_df['detector'] = det
        all_hits_dfs.append(hits_df)

    # Concatenate all hits and links
    combined_hits_df = pd.concat(all_hits_dfs, ignore_index=True) if all_hits_dfs else pd.DataFrame()

    # Calculate derived geometric quantities on the combined dataframe
    if not combined_hits_df.empty:
        combined_hits_df['r'] = _calculate_R(combined_hits_df['x'], combined_hits_df['y'])
        combined_hits_df['R'] = _calculate_R(combined_hits_df['x'], combined_hits_df['y'], combined_hits_df['z'])
        combined_hits_df['phi'] = np.arctan2(combined_hits_df['y'], combined_hits_df['x'])
        combined_hits_df['theta'] = _calculate_theta(combined_hits_df['r'], combined_hits_df['z'])
        combined_hits_df['eta'] = _calculate_eta(combined_hits_df['theta'])
        # Only calculate pt if momentum columns are present
        if full_load and 'px' in combined_hits_df.columns and 'py' in combined_hits_df.columns:
            combined_hits_df['pt'] = np.sqrt(combined_hits_df['px']**2 + combined_hits_df['py']**2)

    logger.debug(f"_build_tracker_df: event_idx={event_idx}, rows={len(combined_hits_df)}, cols={list(combined_hits_df.columns)}")
    return combined_hits_df


def _process_single_calo(events_tree, event_idx, detector_name, contrib_offset=None, full_load: bool = False):
    """Process hits, contributions, and links for a single calorimeter detector."""
    hits = events_tree[detector_name].arrays(entry_start=event_idx, entry_stop=event_idx+1)
    contributions = events_tree[f"{detector_name}Contributions"].arrays(entry_start=event_idx, entry_stop=event_idx+1)
    particle_links = events_tree[f"_{detector_name}Contributions_particle"].arrays(entry_start=event_idx, entry_stop=event_idx+1)

     # Check if event exists / has hits for this detector
    if len(hits[f"{detector_name}.cellID"]) == 0:
        return pd.DataFrame(), pd.DataFrame()

    hit_dict = {
        'cellID': hits[f"{detector_name}.cellID"][0],
        'energy': hits[f"{detector_name}.energy"][0],
        'x': hits[f"{detector_name}.position.x"][0],
        'y': hits[f"{detector_name}.position.y"][0],
        'z': hits[f"{detector_name}.position.z"][0],
        'contribution_begin': hits[f"{detector_name}.contributions_begin"][0],
        'contribution_end': hits[f"{detector_name}.contributions_end"][0]
    }
    hits_df = pd.DataFrame(hit_dict)

    # Core contribution data always loaded
    contrib_dict = {
        'energy': contributions[f"{detector_name}Contributions.energy"][0],
        'time': contributions[f"{detector_name}Contributions.time"][0],
        'particle_id': particle_links[f"_{detector_name}Contributions_particle.index"][0]
    }
    
    # Add optional fields if full_load=True
    if full_load:
        contrib_dict.update({
            'PDG': contributions[f"{detector_name}Contributions.PDG"][0],
            'step_x': contributions[f"{detector_name}Contributions.stepPosition.x"][0],
            'step_y': contributions[f"{detector_name}Contributions.stepPosition.y"][0],
            'step_z': contributions[f"{detector_name}Contributions.stepPosition.z"][0],
        })
    contrib_df = pd.DataFrame(contrib_dict)

    # Add hit positions to contributions
    contrib_df = _add_hit_positions_to_contributions(hits_df, contrib_df)

    if contrib_offset is not None:
        hits_df['contribution_begin'] += contrib_offset
        hits_df['contribution_end'] += contrib_offset
    
    return hits_df, contrib_df

def _process_single_digi_calo(events_tree, event_idx, detector_name, full_load: bool = False):
    """Process digitized hits for a single calorimeter detector.
    
    Simpler than truth calo - no contributions, just aggregated detector signals.
    """
    if detector_name not in events_tree:
        return pd.DataFrame()
    
    hits = events_tree[detector_name].arrays(entry_start=event_idx, entry_stop=event_idx+1)
    
    # Check if event has hits for this detector
    if len(hits[f"{detector_name}.cellID"]) == 0:
        return pd.DataFrame()
    
    hit_dict = {
        'cellID': hits[f"{detector_name}.cellID"][0],
        'energy': hits[f"{detector_name}.energy"][0],
        'energy_error': hits[f"{detector_name}.energyError"][0],
        'time': hits[f"{detector_name}.time"][0],
        'x': hits[f"{detector_name}.position.x"][0],
        'y': hits[f"{detector_name}.position.y"][0],
        'z': hits[f"{detector_name}.position.z"][0],
        'type': hits[f"{detector_name}.type"][0],
    }
    
    return pd.DataFrame(hit_dict)

def _build_calo_df(events_tree, event_idx, full_load: bool = False):
    """Build combined dataframes for calo hits, contributions, and particle links for a single event"""
    all_hits_dfs = []
    all_contrib_dfs = []
    global_contrib_offset = 0

    for det in all_calos:
        if det not in events_tree:
            continue

        hits_df, contrib_df = _process_single_calo(events_tree, event_idx, det, global_contrib_offset, full_load)

        # Add detector info
        hits_df['detector'] = det
        contrib_df['detector'] = det
        all_contrib_dfs.append(contrib_df)
        all_hits_dfs.append(hits_df)

        # Update offsets for the next detector
        global_contrib_offset += len(contrib_df) # Use length of contributions DF

    # Concatenate all dataframes
    combined_hits_df = pd.concat(all_hits_dfs, ignore_index=True) if all_hits_dfs else pd.DataFrame()
    combined_contrib_df = pd.concat(all_contrib_dfs, ignore_index=True) if all_contrib_dfs else pd.DataFrame()

    # Calculate derived geometric quantities on the combined hits dataframe
    combined_hits_df['r'] = _calculate_R(combined_hits_df['x'], combined_hits_df['y'])
    combined_hits_df['R'] = _calculate_R(combined_hits_df['x'], combined_hits_df['y'], combined_hits_df['z'])
    combined_hits_df['phi'] = np.arctan2(combined_hits_df['y'], combined_hits_df['x'])
    combined_hits_df['theta'] = _calculate_theta(combined_hits_df['r'], combined_hits_df['z'])
    combined_hits_df['eta'] = _calculate_eta(combined_hits_df['theta'])

    logger.debug(f"_build_calo_df: event_idx={event_idx}, hits rows={len(combined_hits_df)}, contrib rows={len(combined_contrib_df)}")
    return combined_hits_df, combined_contrib_df

def _build_digi_calo_df(events_tree, event_idx, full_load: bool = False):
    """Build combined dataframe for digitized calo hits for a single event."""
    all_hits_dfs = []
    
    for det in all_digi_calos:
        if det not in events_tree:
            logger.warning("Detector %s not found in events tree", det)
            continue
        
        hits_df = _process_single_digi_calo(events_tree, event_idx, det, full_load)
        if hits_df.empty:
            continue
            
        # Add detector info
        hits_df['detector'] = det
        all_hits_dfs.append(hits_df)
    
    # Concatenate all dataframes
    combined_hits_df = pd.concat(all_hits_dfs, ignore_index=True) if all_hits_dfs else pd.DataFrame()
    
    # Calculate derived geometric quantities (same as truth calo)
    if not combined_hits_df.empty:
        combined_hits_df['r'] = _calculate_R(combined_hits_df['x'], combined_hits_df['y'])
        combined_hits_df['R'] = _calculate_R(combined_hits_df['x'], combined_hits_df['y'], combined_hits_df['z'])
        combined_hits_df['phi'] = np.arctan2(combined_hits_df['y'], combined_hits_df['x'])
        combined_hits_df['theta'] = _calculate_theta(combined_hits_df['r'], combined_hits_df['z'])
        combined_hits_df['eta'] = _calculate_eta(combined_hits_df['theta'])
    
    logger.debug(f"_build_digi_calo_df: event_idx={event_idx}, hits rows={len(combined_hits_df)}")
    return combined_hits_df

def _add_hit_positions_to_contributions(hits_df, contrib_df):
    """Add mapping to hit index and bring over true hit cellID and positions.

    - Adds a `hit_index` column indicating the row index of the parent hit.
    - Merges the true detector `cellID` and hit positions (x, y, z) from hits_df.
    """
    import time
    import logging
    logger = logging.getLogger(__name__)
    
    t0_total = time.time()
    
    t0_repeat = time.time()
    num_contribs = hits_df.contribution_end - hits_df.contribution_begin
    hit_index = np.repeat(hits_df.index.values, num_contribs)
    t_repeat = time.time() - t0_repeat
    
    contrib_df["hit_index"] = hit_index

    # Merge true detector cellID and positions from hits_df using the hit row index
    t0_merge = time.time()
    contrib_df = contrib_df.merge(
        hits_df[["cellID", "x", "y", "z"]],
        left_on="hit_index",
        right_index=True,
        how="left",
    )
    t_merge = time.time() - t0_merge
    
    t_total = time.time() - t0_total
    logger.debug(f"  ⏱️  _add_hit_positions_to_contributions: {t_total:.4f}s ({len(contrib_df)} rows)")
    logger.debug(f"      • np.repeat: {t_repeat:.4f}s | merge: {t_merge:.4f}s")
    
    return contrib_df


def condense_calo_dataframes(
    calo_hits_df: pd.DataFrame,
    calo_contributions_df: pd.DataFrame,
) -> Tuple[pd.DataFrame, pd.DataFrame]:
    """Condense duplicated calorimeter contributions per (cellID, particle_id).

    Vectorized procedure that:
      1) Aggregates duplicate rows in contributions per (cellID, particle_id)
         with energy sum and energy-weighted mean time
      2) Restores geometry from hits and orders contributions to match hits order
      3) Recomputes contribution_begin/end ranges on hits

    For batch mode (with event_id), processes each event separately to maintain correct indices.

    Args:
        calo_hits_df: DataFrame of calorimeter hits with columns including
            'cellID', 'energy', 'x', 'y', 'z', 'contribution_begin', 'contribution_end', 'detector'.
        calo_contributions_df: DataFrame of calorimeter contributions with columns including
            'energy', 'time', 'particle_id', 'cellID', 'x', 'y', 'z', 'detector'.

    Returns:
        Tuple[pd.DataFrame, pd.DataFrame]:
            - new_calo_hits_df: same rows as input hits, with updated
              'contribution_begin'/'contribution_end'.
            - new_calo_contributions_df: condensed contributions in the
              exact order implied by hits.
    """
    import time
    import logging
    logger = logging.getLogger(__name__)
    
    t0_total = time.time()
    
    # Early exits for empty inputs
    if calo_hits_df is None or calo_contributions_df is None:
        return calo_hits_df, calo_contributions_df
    if calo_hits_df.empty or calo_contributions_df.empty:
        return calo_hits_df.copy(), calo_contributions_df.copy()

    # Check if this is batch mode (has event_id column)
    has_event_id = 'event_id' in calo_hits_df.columns and 'event_id' in calo_contributions_df.columns
    
    if has_event_id:
        # Batch mode: process each event separately
        logger.debug(f"⏱️  condense_calo_dataframes (batch mode): {len(calo_hits_df)} hits, {len(calo_contributions_df)} contributions")
        
        event_ids = calo_hits_df['event_id'].unique()
        condensed_hits_list = []
        condensed_contribs_list = []
        
        for event_id in event_ids:
            # Filter data for this event
            event_hits = calo_hits_df[calo_hits_df['event_id'] == event_id].copy()
            event_contribs = calo_contributions_df[calo_contributions_df['event_id'] == event_id].copy()
            
            if event_hits.empty or event_contribs.empty:
                condensed_hits_list.append(event_hits)
                if not event_contribs.empty:
                    condensed_contribs_list.append(event_contribs)
                continue
            
            # Reset indices for this event's processing
            event_hits_reset = event_hits.reset_index(drop=True)
            event_contribs_reset = event_contribs.reset_index(drop=True)
            
            # Condense this event (using per-event logic)
            new_hits, new_contribs = _condense_calo_event(event_hits_reset, event_contribs_reset)
            
            # Restore event_id column
            new_hits['event_id'] = event_id
            new_contribs['event_id'] = event_id
            
            condensed_hits_list.append(new_hits)
            condensed_contribs_list.append(new_contribs)
        
        # Concatenate all events
        new_calo_hits_df = pd.concat(condensed_hits_list, ignore_index=True) if condensed_hits_list else pd.DataFrame()
        new_calo_contributions_df = pd.concat(condensed_contribs_list, ignore_index=True) if condensed_contribs_list else pd.DataFrame()
    else:
        # Single event mode: use original logic
        logger.debug(f"⏱️  condense_calo_dataframes (single event mode): {len(calo_hits_df)} hits, {len(calo_contributions_df)} contributions")
        new_calo_hits_df, new_calo_contributions_df = _condense_calo_event(calo_hits_df, calo_contributions_df)
    
    t_total = time.time() - t0_total
    logger.debug(f"      ✓ condensation complete: {t_total:.4f}s ({len(new_calo_contributions_df)} condensed rows)")
    
    return new_calo_hits_df, new_calo_contributions_df


def _condense_calo_event(hits: pd.DataFrame, contribs: pd.DataFrame) -> Tuple[pd.DataFrame, pd.DataFrame]:
    """Condense a single event's calorimeter data.
    
    Helper function that performs the actual condensation logic for one event.
    """
    hits = hits.copy()
    contribs = contribs.copy()

    # Ensure matching unsigned dtype for cellID to avoid wraparound
    cellid_dtype = hits['cellID'].dtype
    if contribs['cellID'].dtype != cellid_dtype:
        try:
            contribs['cellID'] = contribs['cellID'].astype(cellid_dtype)
        except Exception:
            pass

    # Aggregate duplicate contributions per (cellID, particle_id)
    # Compute energy-weighted mean time
    contribs['energy_time'] = contribs['energy'] * contribs['time']
    aggregated = (
        contribs.groupby(['cellID', 'particle_id'], sort=False)
                .agg(energy=('energy', 'sum'), energy_time=('energy_time', 'sum'))
                .reset_index()
    )
    aggregated['time'] = aggregated['energy_time'] / aggregated['energy']
    aggregated = aggregated.drop(columns=['energy_time'])

    # Attach geometry from hits (many_to_one on cellID)
    aggregated = aggregated.merge(
        # hits[['cellID', 'x', 'y', 'z', 'detector']].drop_duplicates('cellID'),
        hits[['cellID', 'detector']].drop_duplicates('cellID'),
        on='cellID', how='left', validate='many_to_one'
    )

    # Order contributions by hits' cellID order, stable by particle_id
    cat = pd.Categorical(
        aggregated['cellID'],
        categories=hits['cellID'],
        ordered=True,
    )
    agg_sorted = (
        aggregated
            .assign(cellID=cat)
            .sort_values(['cellID', 'particle_id'], kind='mergesort')
            .reset_index(drop=True)
    )
    # Restore original dtype for cellID
    try:
        agg_sorted['cellID'] = agg_sorted['cellID'].astype(cellid_dtype)
    except Exception:
        pass

    # Build condensed contributions table
    # new_calo_contributions_df = agg_sorted[['energy', 'time', 'particle_id', 'cellID', 'x', 'y', 'z', 'detector']].copy()
    new_calo_contributions_df = agg_sorted[['energy', 'time', 'particle_id', 'cellID', 'detector']].copy()

    # Map cellID to hit_index (row index in hits). Assumes one hit per cellID.
    # If duplicates exist, the first occurrence is used, matching mergesort stability.
    cellid_to_hitidx = dict(zip(hits['cellID'].tolist(), hits.index.tolist()))
    
    # Debug: check for missing cellIDs (should not happen if contributions came from hits)
    mapped_indices = new_calo_contributions_df['cellID'].map(cellid_to_hitidx)
    if mapped_indices.isnull().any():
        missing_cellids = new_calo_contributions_df.loc[mapped_indices.isnull(), 'cellID'].unique()
        logger.warning(
            "%d cellIDs in contributions not found in hits: %s; contributions dtype=%s, hits dtype=%s",
            len(missing_cellids), missing_cellids[:10].tolist(),
            new_calo_contributions_df['cellID'].dtype, hits['cellID'].dtype
        )
        logger.debug("Sample hits cellIDs: %s", hits['cellID'].head(10).tolist())
    
    new_calo_contributions_df['hit_index'] = mapped_indices.astype(np.int64)

    # Vectorized begin/end recomputation
    counts_per_cell = agg_sorted.groupby('cellID', sort=False).size()
    counts = (
        hits['cellID']
            .map(counts_per_cell)
            .fillna(0)
            .astype(np.int64)
    )
    cum = counts.cumsum()
    new_end = cum.to_numpy()
    new_begin = (cum - counts).to_numpy()

    new_calo_hits_df = hits.copy()
    new_calo_hits_df['contribution_begin'] = new_begin
    new_calo_hits_df['contribution_end'] = new_end

    # Consistency assertions (safe prints could be added by caller)
    # Ensure length agreement
    try:
        if len(new_calo_hits_df) > 0:
            assert len(new_calo_contributions_df) == int(new_end[-1])
    except Exception:
        pass

    return new_calo_hits_df, new_calo_contributions_df


# --- Main Loading Function ---

def load_event_data(file_path, event_idx, full_load: bool = False):
    """
    Loads all relevant data for a single event from an EDM4hep ROOT file.

    Args:
        file_path (str): Path to the EDM4hep ROOT file.
        event_idx (int): The 0-based index of the event to load.

    Returns:
        dict: A dictionary containing the following Pandas DataFrames:
            - 'particles': MCParticles data.
            - 'parents': Parent links for MCParticles.
            - 'daughters': Daughter links for MCParticles.
            - 'tracker_hits': Combined tracker hits data.
            - 'tracker_links': Links between tracker hits and MCParticles.
            - 'calo_hits': Combined calorimeter hits data.
            - 'calo_contributions': Combined calorimeter contribution data.
            - 'calo_links': Links between calo contributions and MCParticles.
            - 'event_header': Basic event header info (as dict for now).
             Returns None if the event cannot be loaded or file not found.
    """
    try:
        root_file = uproot.open(file_path)
    except FileNotFoundError:
        logger.error("File not found at %s", file_path)
        return None
    except Exception as e:
        logger.error("Error opening file %s: %s", file_path, e)
        return None

    if "events" not in root_file:
        logger.error("'events' tree not found in %s", file_path)
        root_file.close()
        return None

    events_tree = root_file["events"]

    # Basic check for event index validity
    num_events = events_tree.num_entries
    if event_idx < 0 or event_idx >= num_events:
        logger.error("Event index %s is out of bounds (0-%s).", event_idx, num_events - 1)
        root_file.close()
        return None

    logger.debug(f"Loading event {event_idx} from {file_path}...")

    particles_df, parents_df, daughters_df = _build_particle_df(events_tree, event_idx, full_load)
    logger.debug(f"  Loaded {len(particles_df)} particles.")

    tracker_hits_df = _build_tracker_df(events_tree, event_idx, full_load)
    logger.debug(f"  Loaded {len(tracker_hits_df)} tracker hits.")

    calo_hits_df, calo_contrib_df = _build_calo_df(events_tree, event_idx, full_load)
    logger.debug(f"  Loaded {len(calo_hits_df)} calo hits and {len(calo_contrib_df)} contributions.")

    digi_calo_hits_df = _build_digi_calo_df(events_tree, event_idx, full_load)
    logger.debug(f"  Loaded {len(digi_calo_hits_df)} digi calo hits.")

    # Load basic event header info
    header_data = {}
    header_branch = "EventHeader"
    if header_branch in events_tree:
         header = events_tree[header_branch].arrays(entry_start=event_idx, entry_stop=event_idx+1)
         if len(header[f"{header_branch}.eventNumber"]) > 0: # Check if event exists
             header_data = {
                 'eventNumber': header[f"{header_branch}.eventNumber"][0],
                 'runNumber': header[f"{header_branch}.runNumber"][0],
                 'timeStamp': header[f"{header_branch}.timeStamp"][0],
                 'weight': header[f"{header_branch}.weight"][0]
            }
         else:
             logger.warning("EventHeader branch exists but seems empty for this event.")
    else:
        logger.warning("EventHeader branch not found.")


    root_file.close()

    return {
        "particles": particles_df,
        "parents": parents_df,
        "daughters": daughters_df,
        "tracker_hits": tracker_hits_df,
        "calo_hits": calo_hits_df,
        "calo_contributions": calo_contrib_df,
        "digi_calo_hits": digi_calo_hits_df,
        "event_header": header_data
    }


# --- Shared derivations for particles (used by both Event and EventBatch) ---

def _ensure_particle_identifier_columns(particles_df: pd.DataFrame, logger: Optional[logging.Logger] = None) -> pd.DataFrame:
    """Ensure particles_df has a stable per-event 'particle_id' and 'event_id'.

    - If 'particle_id' missing, use 'subentry' when available; otherwise use the index.
    - Returns a new DataFrame (copy) with ensured columns.
    """
    df = particles_df.copy()
    if 'event_id' not in df.columns and 'entry' in df.columns:
        df = df.rename(columns={'entry': 'event_id'})
    if 'particle_id' not in df.columns:
        if 'subentry' in df.columns:
            df['particle_id'] = df['subentry'].astype(int)
        else:
            # Fallback: use the current index position
            df = df.reset_index().rename(columns={'index': 'particle_id'})
    if logger:
        logger.debug(f"Particles identifier columns ensured: has event_id={ 'event_id' in df.columns }, has particle_id={ 'particle_id' in df.columns }")
    return df


def _augment_particle_kinematics(df: pd.DataFrame) -> pd.DataFrame:
    """Add vr, endpoint_r, energy, kinetic_energy to particles_df if available."""
    if df.empty:
        return df
    # Derived geometry and energy
    if {'vx', 'vy'}.issubset(df.columns):
        df['vr'] = np.sqrt(df['vx']**2 + df['vy']**2)
    if {'endpoint_x', 'endpoint_y'}.issubset(df.columns):
        df['endpoint_r'] = np.sqrt(df['endpoint_x']**2 + df['endpoint_y']**2)
    if {'px', 'py', 'pz', 'mass'}.issubset(df.columns):
        df['energy'] = np.sqrt(df['px']**2 + df['py']**2 + df['pz']**2 + df['mass']**2)
        df['kinetic_energy'] = df['energy'] - df['mass']
    return df


def _augment_particle_sim_flags(df: pd.DataFrame) -> pd.DataFrame:
    """Add created_in_simulation flag from simulator status if present."""
    status_col = None
    for cand in ('simulator_status', 'simulatorStatus'):
        if cand in df.columns:
            status_col = cand
            break
    if status_col is None:
        return df
    try:
        df['created_in_simulation'] = df[status_col].apply(get_simulator_status_bits).apply(lambda bits: bits.get('created_in_simulation', False))
    except Exception:
        pass
    return df


def _augment_particle_hit_counts(df: pd.DataFrame, tracker_hits_df: Optional[pd.DataFrame], calo_contrib_df: Optional[pd.DataFrame]) -> pd.DataFrame:
    """Add num_tracker_hits and num_calo_hits by grouping hits on (event_id, particle_id)."""
    import time
    import logging
    logger = logging.getLogger(__name__)
    
    if df.empty:
        return df
    
    # Ensure columns exist before grouping
    if tracker_hits_df is not None and not tracker_hits_df.empty and 'particle_id' in tracker_hits_df.columns:
        t0_tracker = time.time()
        logger.debug(f"  ⏱️  groupby tracker_hits ({len(tracker_hits_df)} rows)...")
        
        # If event_id is present, group by both event_id and particle_id, otherwise just particle_id
        if 'event_id' in tracker_hits_df.columns:
            groupby_cols = ['event_id', 'particle_id']
            merge_on = ['event_id', 'particle_id']
        else:
            groupby_cols = ['particle_id']
            merge_on = ['particle_id']
        
        t0_gb = time.time()
        tcounts = (
            tracker_hits_df.groupby(groupby_cols).size().rename('num_tracker_hits').reset_index()
        )
        t_gb = time.time() - t0_gb
        
        t0_merge = time.time()
        df = df.merge(tcounts, on=merge_on, how='left')
        t_merge = time.time() - t0_merge
        
        t_tracker = time.time() - t0_tracker
        logger.debug(f"      ✓ tracker: {t_tracker:.4f}s (groupby: {t_gb:.4f}s, merge: {t_merge:.4f}s)")
    else:
        df['num_tracker_hits'] = 0
    
    if calo_contrib_df is not None and not calo_contrib_df.empty and 'particle_id' in calo_contrib_df.columns:
        t0_calo = time.time()
        logger.debug(f"  ⏱️  groupby calo_contributions ({len(calo_contrib_df)} rows)...")
        
        # If event_id is present, group by both event_id and particle_id, otherwise just particle_id
        if 'event_id' in calo_contrib_df.columns:
            groupby_cols = ['event_id', 'particle_id']
            merge_on = ['event_id', 'particle_id']
        else:
            groupby_cols = ['particle_id']
            merge_on = ['particle_id']
        
        t0_gb = time.time()
        ccounts = (
            calo_contrib_df.groupby(groupby_cols).size().rename('num_calo_hits').reset_index()
        )
        t_gb = time.time() - t0_gb
        
        t0_merge = time.time()
        df = df.merge(ccounts, on=merge_on, how='left')
        t_merge = time.time() - t0_merge
        
        t_calo = time.time() - t0_calo
        logger.debug(f"      ✓ calo: {t_calo:.4f}s (groupby: {t_gb:.4f}s, merge: {t_merge:.4f}s)")
    else:
        df['num_calo_hits'] = 0
    
    # Fill NaNs after merges
    df['num_tracker_hits'] = df['num_tracker_hits'].fillna(0).astype(int)
    df['num_calo_hits'] = df['num_calo_hits'].fillna(0).astype(int)
    return df


def _augment_particle_geometry_flags(df: pd.DataFrame, detector_params: Optional[dict]) -> pd.DataFrame:
    """Add created_inside_tracker, ended_inside_tracker, and backscatter if geometry provided."""
    if df.empty or not detector_params:
        return df
    tracking_radius = detector_params.get('tracking_radius')
    tracking_z_max = detector_params.get('tracking_z_max')
    if tracking_radius is None or tracking_z_max is None:
        return df
    has_vertex = {'vr', 'vz'}.issubset(df.columns)
    has_end = {'endpoint_r', 'endpoint_z'}.issubset(df.columns)
    if has_vertex:
        df['created_inside_tracker'] = (df['vr'] < tracking_radius) & (df['vz'].abs() < tracking_z_max)
    if has_end:
        df['ended_inside_tracker'] = (df['endpoint_r'] < tracking_radius) & (df['endpoint_z'].abs() < tracking_z_max)
    if 'created_inside_tracker' in df.columns and 'ended_inside_tracker' in df.columns:
        df['backscatter'] = (~df['created_inside_tracker']) & (df['ended_inside_tracker'])
    return df


def augment_particles(
    particles_df: pd.DataFrame,
    tracker_hits_df: Optional[pd.DataFrame] = None,
    calo_contrib_df: Optional[pd.DataFrame] = None,
    detector_params: Optional[dict] = None,
    logger: Optional[logging.Logger] = None,
) -> pd.DataFrame:
    """Augment particles with derived kinematics, sim flags, hit counts, and geometry flags.

    Returns a new DataFrame. Safe to call multiple times.
    """
    import time
    t0_total = time.time()
    
    if logger:
        logger.debug(f"augment_particles: input particles rows={len(particles_df)}, cols={list(particles_df.columns)}")
        if tracker_hits_df is not None:
            logger.debug(f"augment_particles: tracker_hits rows={len(tracker_hits_df)}, cols={list(tracker_hits_df.columns)}")
        if calo_contrib_df is not None:
            logger.debug(f"augment_particles: calo_contrib rows={len(calo_contrib_df)}, cols={list(calo_contrib_df.columns)}")
    
    t0 = time.time()
    df = _ensure_particle_identifier_columns(particles_df, logger)
    t_identifiers = time.time() - t0
    
    t0 = time.time()
    df = _augment_particle_kinematics(df)
    t_kinematics = time.time() - t0
    
    t0 = time.time()
    df = _augment_particle_sim_flags(df)
    t_sim_flags = time.time() - t0
    
    t0 = time.time()
    df = _augment_particle_hit_counts(df, tracker_hits_df, calo_contrib_df)
    t_hit_counts = time.time() - t0
    
    t0 = time.time()
    df = _augment_particle_geometry_flags(df, detector_params)
    t_geometry = time.time() - t0
    
    if logger:
        logger.debug(f"augment_particles: output rows={len(df)}, cols={list(df.columns)}")
        t_total = time.time() - t0_total
        logger.debug(f"  ⏱️  augment_particles breakdown:")
        logger.debug(f"      • identifiers: {t_identifiers:.4f}s")
        logger.debug(f"      • kinematics: {t_kinematics:.4f}s")
        logger.debug(f"      • sim_flags: {t_sim_flags:.4f}s")
        logger.debug(f"      • hit_counts: {t_hit_counts:.4f}s (usually slowest - involves groupby)")
        logger.debug(f"      • geometry: {t_geometry:.4f}s")
        logger.debug(f"      TOTAL: {t_total:.4f}s")
    
    return df