Indexing and Operations

class Locator:
    """Base class for histogram locators."""
    
    def __init__(self, offset: int = 0):
        """
        Parameters:
        - offset: Additional offset to apply
        """

    def __add__(self, offset: int):
        """Add offset to locator."""

    def __sub__(self, offset: int):
        """Subtract offset from locator."""

class loc(Locator):
    """Locate bin by value."""
    
    def __init__(self, value, offset: int = 0):
        """
        Parameters:
        - value: Value to locate
        - offset: Additional bin offset
        """

    def __call__(self, axis) -> int:
        """Get bin index for this value on the given axis."""

class at:
    """Locate bin by direct index."""
    
    def __init__(self, value: int):
        """
        Parameters:
        - value: Direct bin index
        """

    def __call__(self, axis) -> int:
        """Return the direct index."""

class Underflow(Locator):
    """Locate underflow bin."""
    
    def __call__(self, axis) -> int:
        """Return underflow bin index (-1)."""

class Overflow(Locator):
    """Locate overflow bin."""
    
    def __call__(self, axis) -> int:
        """Return overflow bin index (len(axis))."""

# Pre-created instances
underflow: Underflow
overflow: Overflow

class rebin:
    """Rebinning operation for histogram slicing."""
    
    def __init__(
        self,
        factor_or_axis=None,
        /,
        *,
        factor: int = None,
        groups = None,
        edges = None,
        axis = None
    ):
        """
        Create rebinning operation.

        Parameters:
        - factor_or_axis: Rebinning factor (int) or new axis definition
        - factor: Rebinning factor (combine this many bins into one)
        - groups: List of integers specifying custom groupings
        - edges: New edge positions for rebinning
        - axis: New axis definition to rebin to
        """

    def group_mapping(self, axis):
        """Get the grouping mapping for the given axis."""

    def axis_mapping(self, axis):
        """Get the axis mapping for rebinning."""

# Histogram indexing supports:
# - Integer indices: hist[5]
# - Slice objects: hist[10:20]
# - Locators: hist[bh.loc(2.5)]
# - Complex expressions: hist[{0: bh.rebin(2), 1: slice(10, 20)}]

def __getitem__(self, index):
    """
    Get histogram data or sub-histogram.

    Parameters:
    - index: Index expression (int, slice, locator, dict, etc.)

    Returns:
    Histogram, accumulator, or scalar depending on indexing
    """

def __setitem__(self, index, value):
    """
    Set histogram data.

    Parameters:
    - index: Index expression
    - value: Value(s) to set
    """

def Slicer():
    """
    Returns np.s_ for backward compatibility.
    Prefer using np.s_ directly.
    """
    return np.s_

# Dictionary-based indexing for multi-dimensional histograms
# {axis_index: operation} format
# Examples:
# hist[{0: slice(10, 20)}]  # Slice first axis
# hist[{0: bh.rebin(2), 1: bh.loc(5.0)}]  # Rebin axis 0, locate on axis 1
# hist[{1: bh.overflow}]  # Include overflow on axis 1

import boost_histogram as bh
import numpy as np

# Create and fill histogram
hist = bh.Histogram(bh.axis.Regular(100, 0, 10))
data = np.random.normal(5, 2, 1000)
hist.fill(data)

# Direct bin access
bin_5 = hist[5]  # Get contents of bin 5
print(f"Bin 5 content: {bin_5}")

# Slice access
subhist = hist[10:20]  # Get bins 10-19 as new histogram
values_slice = hist.values()[10:20]  # Get values directly

# Flow bin access
hist_with_flow = bh.Histogram(bh.axis.Regular(50, 0, 10))
hist_with_flow.fill([-1, 0, 5, 15])  # Some values in flow bins

underflow_count = hist_with_flow[bh.underflow]
overflow_count = hist_with_flow[bh.overflow]

# Create histogram with regular axis
hist = bh.Histogram(bh.axis.Regular(100, 0, 10))
hist.fill(np.random.uniform(0, 10, 1000))

# Locate by value
bin_at_25 = hist[bh.loc(2.5)]  # Find bin containing value 2.5
bin_near_75 = hist[bh.loc(7.5) + 1]  # Bin after the one containing 7.5

# Slice by values with offset
subset = hist[bh.loc(2.0):bh.loc(8.0)]  # Histogram from value 2.0 to 8.0

# Create fine-grained histogram
hist = bh.Histogram(bh.axis.Regular(100, 0, 10))
hist.fill(np.random.exponential(2, 5000))

# Rebin by factor of 5 (combine every 5 bins)
rebinned = hist[bh.rebin(5)]
print(f"Original bins: {hist.size}, Rebinned: {rebinned.size}")

# Custom rebinning with groups
# Combine first 10 bins into 1, next 20 into 2, rest into individual bins
groups = [10, 10, 10] + [1] * 70  # Total must equal original bin count
custom_rebin = hist[bh.rebin(groups=groups)]

# Rebin to specific edges
new_edges = [0, 1, 2, 5, 10]  # 4 bins with irregular spacing
edge_rebin = hist[bh.rebin(edges=new_edges)]

# Create 2D histogram
hist2d = bh.Histogram(
    bh.axis.Regular(50, 0, 5),
    bh.axis.Regular(50, 0, 5)
)

# Fill with correlated data
x = np.random.uniform(0, 5, 2000)
y = x + np.random.normal(0, 0.5, 2000)
hist2d.fill(x, y)

# Dictionary-based indexing
# Project onto y-axis by summing over all x
y_projection = hist2d[{0: sum}]

# Rebin first axis and slice second axis
modified = hist2d[{0: bh.rebin(2), 1: slice(10, 40)}]

# Select specific ranges on each axis
corner = hist2d[{0: slice(0, 10), 1: slice(40, 50)}]

# Combine locators and rebinning
complex_slice = hist2d[{
    0: bh.loc(2.5):bh.loc(3.5),  # Values between 2.5 and 3.5 on x
    1: bh.rebin(5)               # Rebin y-axis by factor 5
}]

# Create histogram with flow bins enabled
hist = bh.Histogram(bh.axis.Regular(20, 0, 10, underflow=True, overflow=True))

# Fill with data including out-of-range values
data = np.random.normal(5, 3, 1000)  # Some values outside [0, 10]
hist.fill(data)

# Access flow bins
underflow_content = hist[bh.underflow]
overflow_content = hist[bh.overflow]

# Include flow bins in operations
total_with_flow = hist.sum(flow=True)
values_with_flow = hist.values(flow=True)

# Create view including flow bins
full_projection = hist[{0: slice(None)}].values(flow=True)

print(f"Underflow: {underflow_content}")
print(f"Overflow: {overflow_content}")
print(f"Total (no flow): {hist.sum()}")
print(f"Total (with flow): {total_with_flow}")

# 3D histogram for complex slicing
hist3d = bh.Histogram(
    bh.axis.Regular(20, 0, 10),    # x
    bh.axis.Regular(20, 0, 10),    # y  
    bh.axis.Regular(20, 0, 10)     # z
)

# Fill with 3D data
x = np.random.uniform(0, 10, 5000)
y = np.random.uniform(0, 10, 5000)
z = np.random.uniform(0, 10, 5000)
hist3d.fill(x, y, z)

# Complex multi-axis operations
result = hist3d[{
    0: bh.rebin(2),           # Rebin x-axis
    1: bh.loc(5.0):bh.loc(7.0),  # Slice y between 5 and 7
    2: bh.overflow            # Include only overflow in z
}]

# Sequential slicing
temp = hist3d[{0: slice(5, 15)}]  # First operation
final = temp[{1: bh.rebin(2)}]    # Second operation

# Integration patterns
xy_projection = hist3d[{2: sum}]  # Integrate over z
x_profile = hist3d[{1: sum, 2: sum}]  # Integrate over y and z

# Create empty histogram
hist = bh.Histogram(bh.axis.Regular(10, 0, 10))

# Set individual bins
hist[0] = 5.0
hist[bh.loc(2.5)] = 10.0

# Set ranges
hist[3:7] = [1, 2, 3, 4]  # Set bins 3-6
hist[bh.loc(7.0):bh.loc(9.0)] = 15.0  # Set range by value

# Set with numpy arrays
import numpy as np
hist[:] = np.arange(10)  # Set all bins to 0, 1, 2, ..., 9

# Multi-dimensional setting
hist2d = bh.Histogram(bh.axis.Regular(5, 0, 5), bh.axis.Regular(5, 0, 5))
hist2d[{0: 2, 1: 3}] = 42.0  # Set specific 2D bin

# Set slices in 2D
hist2d[{0: slice(1, 4), 1: 2}] = [10, 20, 30]  # Set row of bins