from typing import (
    Any,
    Literal,
    overload,
)

from pandas.api.extensions import ExtensionArray
from pandas.core.arrays.categorical import Categorical
from pandas.core.indexes.base import Index
from pandas.core.indexes.category import CategoricalIndex
from pandas.core.indexes.datetimes import DatetimeIndex
from pandas.core.indexes.interval import IntervalIndex
from pandas.core.indexes.multi import MultiIndex
from pandas.core.indexes.period import PeriodIndex
from pandas.core.indexes.range import RangeIndex
from pandas.core.indexes.timedeltas import TimedeltaIndex
from pandas.core.series import Series

from pandas._typing import (
    T_EXTENSION_ARRAY,
    GenericT,
    IntervalT,
    ShapeT,
    TakeIndexer,
    np_1darray,
    np_1darray_int64,
    np_ndarray,
)

# These are type: ignored because the Index types overlap due to inheritance but indices
# with extension types return the same type while standard type return ndarray
@overload
def unique(values: CategoricalIndex) -> CategoricalIndex: ...
@overload
def unique(values: IntervalIndex[IntervalT]) -> IntervalIndex[IntervalT]: ...
@overload
def unique(values: PeriodIndex) -> PeriodIndex: ...
@overload
def unique(values: DatetimeIndex) -> DatetimeIndex: ...
@overload
def unique(values: TimedeltaIndex) -> TimedeltaIndex: ...
@overload
def unique(values: RangeIndex) -> Index[int]: ...
@overload
def unique(values: MultiIndex) -> MultiIndex: ...
@overload
def unique(values: Index) -> Index: ...
@overload
def unique(values: Categorical) -> Categorical: ...

# @overload
# def unique(values: Series[Never]) -> np_1darray | ExtensionArray: ...
# TODO: DatetimeArray python/mypy#19952
# @overload
# def unique(values: Series[Timestamp]) -> np_ndarray_dt | ExtensionArray: ...
# @overload
# def unique(values: Series[int]) -> np_1darray_anyint | ExtensionArray: ...
@overload
def unique(values: Series) -> np_1darray | ExtensionArray: ...
@overload
def unique(values: np_ndarray[ShapeT, GenericT]) -> np_1darray[GenericT]: ...
@overload
def unique(values: T_EXTENSION_ARRAY) -> T_EXTENSION_ARRAY: ...
@overload
def factorize(
    values: np_ndarray[ShapeT, GenericT],
    sort: bool = False,
    use_na_sentinel: bool = True,
    size_hint: int | None = None,
) -> tuple[np_1darray_int64, np_1darray[GenericT]]: ...
@overload
def factorize(
    values: Index | Series,
    sort: bool = False,
    use_na_sentinel: bool = True,
    size_hint: int | None = None,
) -> tuple[np_1darray_int64, Index]: ...
@overload
def factorize(
    values: Categorical,
    sort: bool = False,
    use_na_sentinel: bool = True,
    size_hint: int | None = None,
) -> tuple[np_1darray_int64, Categorical]: ...

# The following unpublished function is added to reduce type checking ignores
def value_counts(
    values: np_ndarray,
    sort: bool = True,
    ascending: bool = False,
    normalize: bool = False,
    bins: int | None = None,
    dropna: bool = True,
) -> Series[int]: ...
def take(
    arr: np_ndarray[Any] | ExtensionArray | Index | Series,
    indices: TakeIndexer,
    axis: Literal[0, 1] = 0,
    allow_fill: bool = False,
    fill_value: Any = None,
) -> np_1darray | ExtensionArray: ...
