# pylint: disable=unused-import
"""For compatibility and optional dependencies."""

import functools
import importlib.util
import logging
import types
from typing import TYPE_CHECKING, Any, Sequence, TypeGuard, cast

import numpy as np

from ._typing import _T, DataType

if TYPE_CHECKING:
    import pandas as pd
    import pyarrow as pa


def py_str(x: bytes | None) -> str:
    """convert c string back to python string"""
    assert x is not None  # ctypes might return None
    return x.decode("utf-8")  # type: ignore[union-attr]


def lazy_isinstance(instance: Any, module: str, name: str) -> bool:
    """Use string representation to identify a type."""

    # Notice, we use .__class__ as opposed to type() in order
    # to support object proxies such as weakref.proxy
    cls = instance.__class__
    is_same_module = cls.__module__ == module
    has_same_name = cls.__name__ == name
    return is_same_module and has_same_name


# sklearn
try:
    from sklearn import __version__ as _sklearn_version
    from sklearn.base import BaseEstimator as XGBModelBase
    from sklearn.base import ClassifierMixin as XGBClassifierBase
    from sklearn.base import RegressorMixin as XGBRegressorBase
    from sklearn.model_selection import StratifiedKFold as XGBStratifiedKFold

    # sklearn.utils Tags types can be imported unconditionally once
    # xgboost's minimum scikit-learn version is 1.6 or higher
    try:
        from sklearn.utils import Tags as _sklearn_Tags
    except ImportError:
        _sklearn_Tags = object

    SKLEARN_INSTALLED = True

except ImportError:
    SKLEARN_INSTALLED = False

    # used for compatibility without sklearn
    class XGBModelBase:  # type: ignore[no-redef]
        """Dummy class for sklearn.base.BaseEstimator."""

    class XGBClassifierBase:  # type: ignore[no-redef]
        """Dummy class for sklearn.base.ClassifierMixin."""

    class XGBRegressorBase:  # type: ignore[no-redef]
        """Dummy class for sklearn.base.RegressorMixin."""

    XGBStratifiedKFold = None

    _sklearn_Tags = object
    _sklearn_version = object


_logger = logging.getLogger(__name__)


@functools.cache
def is_cudf_available() -> bool:
    """Check cuDF package available or not"""
    if importlib.util.find_spec("cudf") is None:
        return False
    try:
        import cudf

        return True
    except ImportError:
        _logger.exception("Importing cuDF failed, use DMatrix instead of QDM")
        return False


@functools.cache
def is_cupy_available() -> bool:
    """Check cupy package available or not"""
    if importlib.util.find_spec("cupy") is None:
        return False
    try:
        import cupy

        return True
    except ImportError:
        return False


@functools.cache
def import_cupy() -> types.ModuleType:
    """Import cupy."""
    if not is_cupy_available():
        raise ImportError("`cupy` is required for handling CUDA buffer.")

    import cupy

    return cupy


@functools.cache
def is_pyarrow_available() -> bool:
    """Check pyarrow package available or not"""
    if importlib.util.find_spec("pyarrow") is None:
        return False
    return True


@functools.cache
def import_pyarrow() -> types.ModuleType:
    """Import pyarrow with memory cache."""
    import pyarrow as pa

    return pa


@functools.cache
def import_pandas() -> types.ModuleType:
    """Import pandas with memory cache."""
    import pandas as pd

    return pd


@functools.cache
def import_polars() -> types.ModuleType:
    """Import polars with memory cache."""
    import polars as pl

    return pl


@functools.cache
def is_pandas_available() -> bool:
    """Check the pandas package is available or not."""
    if importlib.util.find_spec("pandas") is None:
        return False
    return True


try:
    import scipy.sparse as scipy_sparse
    from scipy.sparse import csr_matrix as scipy_csr
except ImportError:
    scipy_sparse = False
    scipy_csr = object


def _is_polars_lazyframe(data: DataType) -> bool:
    return lazy_isinstance(data, "polars.lazyframe.frame", "LazyFrame")


def _is_polars_series(data: DataType) -> bool:
    return lazy_isinstance(data, "polars.series.series", "Series")


def _is_polars(data: DataType) -> bool:
    lf = _is_polars_lazyframe(data)
    df = lazy_isinstance(data, "polars.dataframe.frame", "DataFrame")
    return lf or df


def _is_arrow(data: DataType) -> TypeGuard["pa.Table"]:
    return lazy_isinstance(data, "pyarrow.lib", "Table")


def _is_cudf_df(data: DataType) -> bool:
    return lazy_isinstance(data, "cudf.core.dataframe", "DataFrame")


def _is_cudf_ser(data: DataType) -> bool:
    return lazy_isinstance(data, "cudf.core.series", "Series")


def _is_cudf_pandas(data: DataType) -> bool:
    """Must go before both pandas and cudf checks."""
    return (_is_pandas_df(data) or _is_pandas_series(data)) and lazy_isinstance(
        type(data), "cudf.pandas.fast_slow_proxy", "_FastSlowProxyMeta"
    )


def _is_pandas_df(data: DataType) -> TypeGuard["pd.DataFrame"]:
    return lazy_isinstance(data, "pandas.core.frame", "DataFrame") or lazy_isinstance(
        data, "pandas", "DataFrame"
    )


def _is_pandas_series(data: DataType) -> TypeGuard["pd.Series"]:
    return lazy_isinstance(data, "pandas.core.series", "Series") or lazy_isinstance(
        data, "pandas", "Series"
    )


def _is_modin_df(data: DataType) -> bool:
    return lazy_isinstance(data, "modin.pandas.dataframe", "DataFrame")


def _is_modin_series(data: DataType) -> bool:
    return lazy_isinstance(data, "modin.pandas.series", "Series")


def is_dataframe(data: DataType) -> bool:
    """Whether the input is a dataframe. Currently supported dataframes:

    - pandas
    - cudf
    - cudf.pandas
    - polars
    - pyarrow
    - modin


    """
    return any(
        p(data)
        for p in (
            _is_polars,
            _is_polars_series,
            _is_arrow,
            _is_cudf_df,
            _is_cudf_ser,
            _is_cudf_pandas,
            _is_pandas_df,
            _is_pandas_series,
            _is_modin_df,
            _is_modin_series,
        )
    )


def _is_cupy_alike(data: DataType) -> bool:
    return hasattr(data, "__cuda_array_interface__")


def concat(value: Sequence[_T]) -> _T:  # pylint: disable=too-many-return-statements
    """Concatenate row-wise."""
    if isinstance(value[0], np.ndarray):
        value_arr = cast(Sequence[np.ndarray], value)
        return cast(_T, np.concatenate(value_arr, axis=0))
    if scipy_sparse and isinstance(value[0], scipy_sparse.csr_matrix):
        return scipy_sparse.vstack(value, format="csr")
    if scipy_sparse and isinstance(value[0], scipy_sparse.csc_matrix):
        return scipy_sparse.vstack(value, format="csc")
    if scipy_sparse and isinstance(value[0], scipy_sparse.spmatrix):
        # other sparse format will be converted to CSR.
        return scipy_sparse.vstack(value, format="csr")
    if _is_pandas_df(value[0]) or _is_pandas_series(value[0]):
        from pandas import concat as pd_concat

        return pd_concat(value, axis=0)
    if lazy_isinstance(value[0], "cudf.core.dataframe", "DataFrame") or lazy_isinstance(
        value[0], "cudf.core.series", "Series"
    ):
        from cudf import concat as CUDF_concat

        return CUDF_concat(value, axis=0)
    if _is_cupy_alike(value[0]):
        import cupy

        # pylint: disable=c-extension-no-member,no-member
        d = cupy.cuda.runtime.getDevice()
        for v in value:
            arr = cast(cupy.ndarray, v)
            d_v = arr.device.id
            assert d_v == d, "Concatenating arrays on different devices."
        return cupy.concatenate(value, axis=0)
    raise TypeError(f"Unknown type: {type(value[0])}")
