import numpy as np
import numpy.typing as npt

from laddu.utils.variables import CosTheta, Mandelstam, Mass, Phi, PolAngle, PolMagnitude
from laddu.utils.vectors import Vector3, Vector4

class Event:
    p4s: list[Vector4]
    eps: list[Vector3]
    weight: float
    def __init__(self, p4s: list[Vector4], eps: list[Vector3], weight: float): ...
    def get_p4_sum(self, indices: list[int]) -> Vector4: ...

class Dataset:
    events: list[Event]
    weights: npt.NDArray[np.float64]
    def __init__(self, events: list[Event]): ...
    def __len__(self) -> int: ...
    def __getitem__(self, index: int) -> Event: ...
    def len(self) -> int: ...
    def weighted_len(self) -> float: ...
    def bin_by(
        self,
        variable: Mass | CosTheta | Phi | PolAngle | PolMagnitude | Mandelstam,
        bins: int,
        range: tuple[float, float],
    ) -> BinnedDataset: ...
    def bootstrap(self, seed: int) -> Dataset: ...

class BinnedDataset:
    bins: int
    range: tuple[float, float]
    edges: npt.NDArray[np.float64]
    def __len__(self) -> int: ...
    def len(self) -> int: ...
    def __getitem__(self, index: int) -> Dataset: ...

def open(path: str) -> Dataset: ...
