import tensorflow as tf
from ..settings import ztypes as ztypes
from ..util import ztyping as ztyping
from ..util.exception import AnalyticIntegralNotImplemented as AnalyticIntegralNotImplemented, WorkInProgressError as WorkInProgressError
from .space import MultiSpace as MultiSpace, convert_to_space as convert_to_space, supports as supports
from _typeshed import Incomplete
from collections.abc import Callable as Callable, Mapping
from zfit import z as z
from zfit._interfaces import ZfitModel as ZfitModel, ZfitSpace as ZfitSpace

def auto_integrate(func, limits, n_axes=None, x=None, method: str = 'AUTO', dtype=..., mc_sampler=..., max_draws=None, tol=None, vectorizable=None, mc_options=None, simpsons_options=None): ...
def numeric_integrate() -> None: ...
def simpson(func, lower, upper, num_points: int = 1001, dtype=None): ...
def simpson_integrate(func, limits, num_points): ...
def mc_integrate(func: Callable, limits: ztyping.LimitsType, axes: ztyping.AxesTypeInput | None = None, x: ztyping.XType | None = None, n_axes: int | None = None, draws_per_dim: int = 40000, max_draws: int = 800000, tol: float = 1e-06, method: str | None = None, xfixed: Mapping | None = None, dtype: type = ..., mc_sampler: Callable = ..., importance_sampling: Callable | None = None, vectorizable=None) -> tf.Tensor: ...
def normalization_nograd(func, n_axes, batch_size, num_batches, dtype, space, x=None, shape_after=()): ...
def normalization_chunked(func, n_axes, batch_size, num_batches, dtype, space, x=None, shape_after=()): ...
def chunked_average(func, x, num_batches, batch_size, space, mc_sampler): ...

class AnalyticIntegral:
    def __init__(self, *args, **kwargs) -> None: ...
    def get_max_axes(self, limits: ztyping.LimitsType) -> tuple[int]: ...
    def get_max_integral(self, limits: ztyping.LimitsType, axes: ztyping.AxesTypeInput = None) -> None | Integral: ...
    def register(self, func: Callable, limits: ztyping.LimitsType, priority: int = 50, *, supports_norm: bool = False, supports_multiple_limits: bool = False) -> None: ...
    def integrate(self, x: ztyping.XType | None, limits: ztyping.LimitsType, axes: ztyping.AxesTypeInput = None, norm: ztyping.LimitsType = None, model: ZfitModel = None, params: dict | None = None) -> ztyping.XType: ...

class Integral:
    limits: Incomplete
    integrate: Incomplete
    axes: Incomplete
    priority: Incomplete
    def __init__(self, func: Callable, limits: ZfitSpace, priority: int | float) -> None: ...
    def __call__(self, *args, **kwargs): ...

class Integration:
    tol: Incomplete
    max_draws: Incomplete
    mc_sampler: Incomplete
    draws_per_dim: Incomplete
    draws_simpson: Incomplete
    def __init__(self, mc_sampler, draws_per_dim, tol, max_draws, draws_simpson) -> None: ...
