import numpy as np
import pandas as pd
import tensorflow as tf
import zfit
from .. import z as z
from ..exception import OutsideLimitsError as OutsideLimitsError
from ..serialization import SpaceRepr as SpaceRepr
from ..serialization.serializer import BaseRepr as BaseRepr, to_orm_init as to_orm_init
from ..settings import run as run, ztypes as ztypes
from ..util import ztyping as ztyping
from ..util.cache import GraphCachable as GraphCachable, invalidate_graph as invalidate_graph
from ..util.container import convert_to_container as convert_to_container
from ..util.exception import BreakingAPIChangeError as BreakingAPIChangeError, ObsIncompatibleError as ObsIncompatibleError, ShapeIncompatibleError as ShapeIncompatibleError, WorkInProgressError as WorkInProgressError
from ..util.temporary import TemporarilySet as TemporarilySet
from .baseobject import BaseObject as BaseObject, convert_param_values as convert_param_values
from .coordinates import convert_to_obs_str as convert_to_obs_str
from .dimension import BaseDimensional as BaseDimensional
from .serialmixin import SerializableMixin as SerializableMixin, ZfitSerializable as ZfitSerializable
from .space import Space as Space, convert_to_space as convert_to_space
from _typeshed import Incomplete
from collections.abc import Callable as Callable, Iterable, Mapping
from tensorflow.python.types.core import TensorLike as TensorLike
from typing import Literal
from zfit._interfaces import ZfitBinnedData as ZfitBinnedData, ZfitSpace as ZfitSpace, ZfitUnbinnedData as ZfitUnbinnedData

def convert_to_data(data, obs=None, *, check_limits: bool = False): ...

class DataMeta(type):
    def __call__(cls, data, obs=None, *args, **kwargs): ...

class Data(ZfitUnbinnedData, BaseDimensional, BaseObject, GraphCachable, SerializableMixin, ZfitSerializable, metaclass=DataMeta):
    USE_HASH: bool
    BATCH_SIZE: int
    dataset: Incomplete
    def __init__(self, data: LightDataset | pd.DataFrame | Mapping[str, np.ndarray] | tf.Tensor | np.ndarray | zfit.Data, *, obs: ztyping.ObsTypeInput = None, weights: TensorLike = None, name: str | None = None, label: str | None = None, dtype: tf.DType = None, use_hash: bool | None = None, guarantee_limits: bool = False) -> None: ...
    @property
    def label(self): ...
    @property
    def num_entries(self): ...
    @property
    def nevents(self) -> int: ...
    @property
    def samplesize(self) -> float: ...
    def enable_hashing(self) -> None: ...
    @property
    def hashint(self) -> int | None: ...
    @property
    def n_events(self): ...
    @property
    def has_weights(self): ...
    @property
    def dtype(self): ...
    @property
    def data_range(self): ...
    @invalidate_graph
    def set_data_range(self, data_range): ...
    @property
    def weights(self) -> tf.Variable: ...
    def with_weights(self, weights: ztyping.WeightsInputType) -> Data: ...
    @invalidate_graph
    def set_weights(self, weights: ztyping.WeightsInputType): ...
    @property
    def space(self) -> ZfitSpace: ...
    @classmethod
    def from_pandas(cls, df: pd.DataFrame, obs: ztyping.ObsTypeInput = None, *, weights: ztyping.WeightsInputType | str = None, name: str | None = None, label: str | None = None, dtype: tf.DType = None, use_hash: bool | None = None, guarantee_limits: bool = False) -> Data | ZfitBinnedData: ...
    @classmethod
    def from_mapping(cls, mapping: Mapping[str, ztyping.ArrayLike], obs: ztyping.ObsTypeInput = None, *, weights: TensorLike | None = None, label: str | None = None, name: str | None = None, dtype: tf.DType = None, use_hash: bool | None = None, guarantee_limits: bool | None = False) -> Data | ZfitBinnedData: ...
    @classmethod
    def from_root(cls, path: str, treepath: str, obs: ZfitSpace = None, *, weights: ztyping.WeightsStrInputType = None, obs_alias: Mapping[str, str] | None = None, name: str | None = None, label: str | None = None, dtype: tf.DType = None, root_dir_options=None, use_hash: bool | None = None, branches: list[str] | None = None, branches_alias: dict | None = None) -> Data: ...
    @classmethod
    def from_numpy(cls, obs: ztyping.ObsTypeInput, array: np.ndarray, *, weights: ztyping.WeightsInputType = None, name: str | None = None, label: str | None = None, dtype: tf.DType = None, use_hash=None, guarantee_limits: bool = False) -> Data | ZfitBinnedData: ...
    @classmethod
    def from_tensor(cls, obs: ztyping.ObsTypeInput, tensor: tf.Tensor, *, weights: ztyping.WeightsInputType = None, name: str | None = None, label: str | None = None, dtype: tf.DType = None, use_hash=None, guarantee_limits: bool = False) -> Data | ZfitBinnedData: ...
    def with_obs(self, obs: ztyping.ObsTypeInput, *, guarantee_limits: bool = False) -> Data: ...
    def to_pandas(self, obs: ztyping.ObsTypeInput = None, weightsname: str | None = None) -> pd.DataFrame: ...
    def unstack_x(self, obs: ztyping.ObsTypeInput = None, always_list=None): ...
    def value(self, obs: ztyping.ObsTypeInput = None, axis: int | None = None) -> tf.Tensor: ...
    def numpy(self) -> np.ndarray: ...
    @property
    def shape(self): ...
    def to_numpy(self) -> np.ndarray: ...
    def sort_by_axes(self, *_, **__) -> None: ...
    def sort_by_obs(self, *_, **__) -> None: ...
    def to_binned(self, space: ztyping.SpaceType, *, name: str | None = None, label: str | None = None, use_hash: bool | None = None) -> ZfitBinnedData: ...
    def __len__(self) -> int: ...
    def __getitem__(self, item): ...

class DataRepr(BaseRepr):
    hs3_type: Literal['Data']
    data: np.ndarray
    space: SpaceRepr | list[SpaceRepr]
    name: str | None
    weights: np.ndarray | None
    def extract_data(cls, values): ...
    def flatten_spaces(cls, v): ...
    def convert_data(cls, v): ...
    def convert_weights(cls, v): ...

def getitem_obs(self, item): ...
def check_cut_datamap_weights(limits, data, weights, guarantee_limits): ...
def check_cut_data_weights(limits: ZfitSpace, data: TensorLike | Mapping[str, TensorLike], weights: TensorLike | None = None, guarantee_limits: bool = False): ...

class SamplerData(Data):
    params: Incomplete
    n: Incomplete
    def __init__(self, data: LightDataset, *, sample_and_weights_func: Callable, sample_holder: tf.Variable, n: ztyping.NumericalScalarType | Callable, weights=None, weights_holder: tf.Variable | None = None, params: dict[zfit.Parameter, ztyping.NumericalScalarType] | None = None, obs: ztyping.ObsTypeInput = None, name: str | None = None, label: str | None = None, dtype: tf.DType = ..., use_hash: bool | None = None, guarantee_limits: bool = False) -> None: ...
    @property
    def fixed_params(self): ...
    @property
    def n_samples(self): ...
    @property
    def hashint(self) -> int | None: ...
    @classmethod
    def get_cache_counting(cls): ...
    @classmethod
    def from_sample(cls, sample_func: Callable, n: ztyping.NumericalScalarType, obs: ztyping.ObsTypeInput, fixed_params=None, name: str | None = None, weights=None, dtype=None, use_hash: bool | None = None): ...
    @classmethod
    def from_sampler(cls, *, sample_func: Callable | None = None, sample_and_weights_func: Callable | None = None, n: ztyping.NumericalScalarType, obs: ztyping.ObsTypeInput, params: ztyping.ParamValuesMap = None, fixed_params=None, name: str | None = None, label: str | None = None, dtype=None, use_hash: bool | None = None, guarantee_limits: bool = False): ...
    def update_data(self, sample: TensorLike | ZfitUnbinnedData, weights: TensorLike | None = None, guarantee_limits: bool = False): ...
    def resample(self, params: ztyping.ParamValuesMap = None, *, n: TensorLike = None, param_values: ztyping.ParamValuesMap = None): ...
    @classmethod
    def get_repr(cls): ...

def concat(datasets: Iterable[Data], *, obs: ztyping.ObsTypeInput = None, axis: int | str | None = None, name: str | None = None, label: str | None = None, use_hash: bool | None = None) -> Data: ...
def concat_data_obs(datasets, obs, name, label, use_hash): ...
def concat_data_index(datasets, obs, name, label, use_hash): ...

class LightDataset:
    def __init__(self, tensor=None, tensormap=None, ndims=None) -> None: ...
    def batch(self, _): ...
    @property
    def num_entries(self): ...
    @property
    def ndims(self): ...
    def __iter__(self): ...
    @classmethod
    def from_tensor(cls, tensor, ndims): ...
    def with_indices(self, indices: int | tuple[int] | list[int]): ...
    def value(self, index: int | tuple[int] | list[int] | None = None): ...
    def calc_hash(self): ...
    def __hash__(self): ...
    def __eq__(self, other): ...

def sum_samples(sample1: ZfitUnbinnedData, sample2: ZfitUnbinnedData, obs: ztyping.ObsTypeInput = None, weights: ztyping.WeightsInputType = None, shuffle: bool = False): ...

class Sampler(SamplerData):
    def __init__(self, *args, **kwargs) -> None: ...
