import abc
import numpy as np
import tensorflow as tf
import zfit
from abc import ABCMeta, abstractmethod
from collections.abc import Callable as Callable
from uhi.typing.plottable import PlottableHistogram
from zfit.util import ztyping as ztyping

class ZfitObject: ...

class ZfitDimensional(ZfitObject, metaclass=abc.ABCMeta):
    @property
    @abstractmethod
    def obs(self) -> ztyping.ObsTypeReturn: ...
    @property
    @abstractmethod
    def axes(self) -> ztyping.AxesTypeReturn: ...
    @property
    @abstractmethod
    def n_obs(self) -> int: ...

class ZfitOrderableDimensional(ZfitDimensional, metaclass=ABCMeta):
    @abstractmethod
    def with_obs(self, obs: ztyping.ObsTypeInput | None, allow_superset: bool = True, allow_subset: bool = True) -> ZfitOrderableDimensional: ...
    @abstractmethod
    def with_axes(self, axes: ztyping.AxesTypeInput | None, allow_superset: bool = True, allow_subset: bool = True) -> ZfitOrderableDimensional: ...
    @abstractmethod
    def with_autofill_axes(self, overwrite: bool = False) -> ZfitOrderableDimensional: ...
    @abstractmethod
    def reorder_x(self, x: tf.Tensor | np.ndarray, *, x_obs: ztyping.ObsTypeInput = None, x_axes: ztyping.AxesTypeInput = None, func_obs: ztyping.ObsTypeInput = None, func_axes: ztyping.AxesTypeInput = None) -> ztyping.XTypeReturnNoData: ...
    @abstractmethod
    def get_reorder_indices(self, obs: ztyping.ObsTypeInput = None, axes: ztyping.AxesTypeInput = None) -> tuple[int]: ...

class ZfitData(ZfitDimensional, metaclass=abc.ABCMeta):
    @abstractmethod
    def value(self, obs: list[str] | None = None) -> ztyping.XType: ...
    @property
    @abstractmethod
    def weights(self): ...
    @property
    def shape(self) -> tuple: ...
    @property
    @abstractmethod
    def samplesize(self) -> float: ...
    @property
    @abstractmethod
    def num_entries(self) -> int: ...
    @abstractmethod
    def to_binned(self, binning: ztyping.BinningTypeInput): ...

class ZfitUnbinnedData(ZfitData, metaclass=abc.ABCMeta):
    @property
    @abstractmethod
    def has_weights(self): ...

class ZfitLimit(abc.ABC, metaclass=ABCMeta):
    @property
    @abstractmethod
    def rect_limits(self) -> ztyping.RectLimitsReturnType: ...
    @property
    @abstractmethod
    def rect_limits_np(self) -> ztyping.RectLimitsNPReturnType: ...
    @property
    @abstractmethod
    def rect_lower(self) -> ztyping.RectLowerReturnType: ...
    @property
    @abstractmethod
    def rect_upper(self) -> ztyping.RectUpperReturnType: ...
    @abstractmethod
    def rect_area(self) -> float | np.ndarray | tf.Tensor: ...
    @abstractmethod
    def inside(self, x: ztyping.XTypeInput, guarantee_limits: bool = False) -> ztyping.XTypeReturn: ...
    @abstractmethod
    def filter(self, x: ztyping.XTypeInput, guarantee_limits: bool = False, axis: int | None = None) -> ztyping.XTypeReturnNoData: ...
    @property
    @abstractmethod
    def has_rect_limits(self) -> bool: ...
    @property
    def rect_limits_are_tensors(self) -> bool: ...
    @property
    @abstractmethod
    def limits_are_set(self) -> bool: ...
    @property
    @abstractmethod
    def limits_are_false(self) -> bool: ...
    @property
    @abstractmethod
    def has_limits(self) -> bool: ...
    def get_subspace(self, *_, **__) -> None: ...
    @property
    @abstractmethod
    def n_obs(self) -> int: ...
    @property
    def n_events(self) -> int | None: ...
    @abstractmethod
    def equal(self, other: object, allow_graph: bool) -> bool | tf.Tensor: ...
    @abstractmethod
    def __eq__(self, other: object) -> bool: ...
    @abstractmethod
    def less_equal(self, other: object, allow_graph: bool = True) -> bool | tf.Tensor: ...
    @abstractmethod
    def __le__(self, other: object) -> bool: ...
    @abstractmethod
    def get_sublimits(self): ...
    @abstractmethod
    def __hash__(self): ...

class ZfitSpace(ZfitLimit, ZfitOrderableDimensional, ZfitObject, metaclass=ABCMeta):
    V1: ZfitSpace
    V0: ZfitSpace
    @property
    def is_binned(self) -> None: ...
    @property
    def binning(self) -> None: ...
    @property
    @abstractmethod
    def n_limits(self) -> int: ...
    @property
    @abstractmethod
    def limits(self) -> tuple[ztyping.LowerTypeReturn, ztyping.UpperTypeReturn]: ...
    @property
    @abstractmethod
    def lower(self) -> ztyping.LowerTypeReturn: ...
    @property
    @abstractmethod
    def upper(self) -> ztyping.UpperTypeReturn: ...
    @abstractmethod
    def with_limits(self, limits: ztyping.LimitsTypeInput = None, rect_limits: ztyping.RectLimitsInputType | None = None, name: str | None = None) -> ZfitSpace: ...
    def with_binning(self, binning: ztyping.BinningTypeInput) -> ZfitSpace: ...
    @abstractmethod
    def get_subspace(self, obs, axes, name): ...
    @abstractmethod
    def with_coords(self, coords: ZfitOrderableDimensional, allow_superset: bool = True, allow_subset: bool = True) -> object: ...
    @abstractmethod
    def __iter__(self): ...
    @abstractmethod
    def __len__(self): ...

class ZfitParametrized(ZfitObject, metaclass=abc.ABCMeta):
    @abstractmethod
    def get_params(self, floating: bool | None = True, is_yield: bool | None = None, extract_independent: bool | None = True, *, autograd: bool | None = None) -> set[ZfitParameter]: ...
    @property
    @abstractmethod
    def params(self) -> ztyping.ParameterType: ...

class ZfitNumericParametrized(ZfitParametrized, metaclass=abc.ABCMeta):
    @property
    @abstractmethod
    def dtype(self) -> tf.DType: ...

class ZfitParameter(ZfitNumericParametrized, metaclass=abc.ABCMeta):
    @property
    @abstractmethod
    def name(self) -> str: ...
    @property
    @abstractmethod
    def shape(self): ...
    @property
    @abstractmethod
    def floating(self) -> bool: ...
    @floating.setter
    @abstractmethod
    def floating(self, value: bool): ...
    @abstractmethod
    def value(self) -> tf.Tensor: ...
    @property
    @abstractmethod
    def independent(self) -> bool: ...

class ZfitIndependentParameter(ZfitParameter, metaclass=ABCMeta):
    @abstractmethod
    def randomize(self, minval, maxval, sampler): ...
    @abstractmethod
    def set_value(self, value): ...
    @property
    @abstractmethod
    def has_limits(self) -> bool: ...
    @property
    @abstractmethod
    def at_limit(self) -> tf.Tensor: ...
    @property
    def stepsize(self) -> tf.Tensor: ...

class ZfitLoss(ZfitObject, metaclass=ABCMeta):
    @abstractmethod
    def gradient(self, params: ztyping.ParamTypeInput = None) -> list[tf.Tensor]: ...
    @abstractmethod
    def value(self, *, full: bool | None = None) -> ztyping.NumericalTypeReturn: ...
    @property
    @abstractmethod
    def model(self) -> list[ZfitModel]: ...
    @property
    @abstractmethod
    def data(self) -> list[ZfitData]: ...
    @property
    @abstractmethod
    def fit_range(self) -> list[ZfitSpace]: ...
    @abstractmethod
    def add_constraints(self, constraints: list[tf.Tensor]): ...
    @property
    @abstractmethod
    def errordef(self) -> float: ...
    def hessian(self, params) -> None: ...
    @abstractmethod
    def value_gradient(self, params): ...
    @abstractmethod
    def value_gradient_hessian(self, params, hessian=None): ...
    @abstractmethod
    def create_new(self, **kwargs): ...
    @property
    @abstractmethod
    def is_weighted(self): ...

class ZfitModel(ZfitNumericParametrized, ZfitDimensional, metaclass=abc.ABCMeta):
    @abstractmethod
    def update_integration_options(self, *args, **kwargs): ...
    @abstractmethod
    def integrate(self, limits: ztyping.LimitsType, norm: ztyping.LimitsType = None, *, options=None) -> ztyping.XType: ...
    @classmethod
    @abstractmethod
    def register_analytic_integral(cls, func: Callable, limits: ztyping.LimitsType = None, priority: int = 50, *, supports_norm: bool = False, supports_multiple_limits: bool = False): ...
    @abstractmethod
    def partial_integrate(self, x: ztyping.XType, limits: ztyping.LimitsType, *, norm=None, options=None, norm_range: ztyping.LimitsType = None) -> ztyping.XType: ...
    @classmethod
    @abstractmethod
    def register_inverse_analytic_integral(cls, func: Callable): ...
    @abstractmethod
    def sample(self, n: int, limits: ztyping.LimitsType = None) -> ztyping.XType: ...

class ZfitFunc(ZfitModel, metaclass=abc.ABCMeta):
    @abstractmethod
    def func(self, x: ztyping.XType, name: str = 'value') -> ztyping.XType: ...
    @abstractmethod
    def as_pdf(self): ...

class ZfitPDF(ZfitModel, metaclass=abc.ABCMeta):
    @abstractmethod
    def pdf(self, x: ztyping.XType, norm: ztyping.LimitsType = None, params=None) -> ztyping.XType: ...
    @property
    @abstractmethod
    def is_extended(self) -> bool: ...
    @abstractmethod
    def set_norm_range(self): ...
    @abstractmethod
    def create_extended(self, yield_: ztyping.ParamTypeInput, name: str | None = None) -> ZfitPDF: ...
    @abstractmethod
    def get_yield(self) -> ZfitParameter | None: ...
    @abstractmethod
    def normalization(self, norm: ztyping.LimitsType, *, options) -> ztyping.NumericalTypeReturn: ...
    @abstractmethod
    def as_func(self, norm: ztyping.LimitsType = False): ...
    @property
    def plot(self) -> zfit.util.plotter.ZfitPDFPlotter: ...

class ZfitFunctorMixin(metaclass=abc.ABCMeta):
    @property
    @abstractmethod
    def models(self) -> dict[float | int | str, ZfitModel]: ...
    @abstractmethod
    def get_models(self) -> list[ZfitModel]: ...

class ZfitConstraint(abc.ABC, metaclass=abc.ABCMeta):
    @abstractmethod
    def value(self): ...

class ZfitMinimalHist(PlottableHistogram):
    @property
    def kind(self) -> None: ...
    def values(self) -> None: ...
    def variances(self) -> np.ndarray | None: ...
    def axes(self) -> None: ...

class ZfitBinnedData(ZfitDimensional, ZfitMinimalHist, metaclass=ABCMeta):
    @abstractmethod
    def variances(self): ...
    def with_obs(self, obs) -> ZfitBinnedData: ...
    @abstractmethod
    def binning(self): ...
    @abstractmethod
    def to_hist(self): ...

class ZfitBinnedPDF(ZfitPDF, metaclass=ABCMeta):
    @abstractmethod
    def counts(self, x, norm): ...
    @abstractmethod
    def rel_counts(self, x, norm): ...

class ZfitBinning: ...

class ZfitRectBinning(ZfitBinning, metaclass=abc.ABCMeta):
    @abstractmethod
    def get_edges(self): ...
