import collections
import importlib
import itertools
import re
import reprlib
import secrets
from collections.abc import Iterable

import numpy as np
from sklearn.base import BaseEstimator, clone

from skrub import _dataframe as sbd


class LRUDict:
    """Dict with limited capacity.

    Using LRU eviction avoids memorizing a full dataset.
    """

    def __init__(self, capacity):
        self.capacity = capacity
        self.cache = collections.OrderedDict()

    def __getitem__(self, key):
        try:
            value = self.cache.pop(key)
            self.cache[key] = value
            return value
        except KeyError:
            return -1

    def __setitem__(self, key, value):
        try:
            self.cache.pop(key)
        except KeyError:
            if len(self.cache) >= self.capacity:
                self.cache.popitem(last=False)
        self.cache[key] = value

    def __contains__(self, key):
        return key in self.cache


def unique_strings(values, is_null):
    """Unique values, accounting for nulls.

    This is like np.unique except
    - it is only for 1d arrays of strings
    - caller must pass a boolean array indicating which values are null: ``is_null``
    - null values are considered to be the same as the empty string.

    >>> import numpy as np
    >>> import pandas as pd
    >>> from skrub._utils import unique_strings
    >>> a = np.asarray(['paris', '', 'berlin', None, 'london'])
    >>> values, idx = unique_strings(a, pd.isna(a))
    >>> values
    array(['', 'berlin', 'london', 'paris'], dtype=object)
    >>> values[idx]
    array(['paris', '', 'berlin', '', 'london'], dtype=object)
    """
    not_null_values = values[~is_null]
    unique, idx = np.unique(not_null_values, return_inverse=True)
    if not is_null.any():
        return unique, idx
    if not len(unique) or unique[0] != "":
        unique = np.concatenate([[""], unique])
        idx += 1
    full_idx = np.empty(values.shape, dtype=idx.dtype)
    full_idx[is_null] = 0
    full_idx[~is_null] = idx
    return unique, full_idx


def import_optional_dependency(name, extra=""):
    """Import an optional dependency.

    By default, if a dependency is missing an ImportError with a nice
    message will be raised.

    Parameters
    ----------
    name : str
        The module name.
    extra : str
        Additional text to include in the ImportError message.

    Returns
    -------
    maybe_module : Optional[ModuleType]
        The imported module when found.
    """

    msg = (
        f"Missing optional dependency '{name}'. {extra} "
        f"Use pip or conda to install {name}."
    )
    try:
        module = importlib.import_module(name)
    except ImportError as exc:
        raise ImportError(msg) from exc

    return module


def atleast_1d_or_none(x):
    """``np.atleast_1d`` helper returning an empty list when x is None."""
    if x is None:
        return []
    return np.atleast_1d(x).tolist()


def _is_array_like(x):
    return (
        isinstance(x, Iterable)
        and not isinstance(x, (str, bytes))
        and not hasattr(x, "__dataframe__")
    )


def clone_if_default(estimator, default_estimator):
    return clone(estimator) if estimator is default_estimator else estimator


def random_string():
    return secrets.token_hex()[:8]


def get_duplicates(values):
    counts = collections.Counter(values)
    duplicates = [k for k, v in counts.items() if v > 1]
    return duplicates


def renaming_func(renaming):
    if isinstance(renaming, str):
        return renaming.format
    return renaming


class Repr(reprlib.Repr):
    fillvalue = "..."

    def repr_dict(self, x, level):
        # (probably for historical reasons) reprlib sorts dict keys, destroying
        # the actual order of the dictionary which is misleading.
        # We adapt the stdlib implementation:
        # https://github.com/python/cpython/blob/630dc2bd6422715f848b76d7950919daa8c44b99/Lib/reprlib.py#L156 # noqa
        # to remove the sorting
        n = len(x)
        if n == 0:
            return "{}"
        if level <= 0:
            return "{" + self.fillvalue + "}"
        newlevel = level - 1
        repr1 = self.repr1
        pieces = []
        for key in itertools.islice(x, self.maxdict):
            keyrepr = repr1(key, newlevel)
            valrepr = repr1(x[key], newlevel)
            pieces.append(f"{keyrepr}: {valrepr}")
        if n > self.maxdict:
            pieces.append(self.fillvalue)

        # in stdlib this line is:
        # s = self._join(pieces, level)
        # but we do not use indentation
        # so we can just ', '.join to avoid using a private method
        s = ", ".join(pieces)

        return f"{{{s}}}"


class _ShortRepr(Repr):
    def repr_instance(self, instance, level):
        if hasattr(instance, "__skrub_short_repr__"):
            return instance.__skrub_short_repr__()

        r = repr(instance)
        if "\n" in r or len(r) > 50:
            if (m := re.match(r"^(\w{1,25})\((.*)\)$", r)) is not None:
                return f"{m.group(1)}({m.group(2)[:20]}...)"
            return f"{instance.__class__.__name__}(...)"
        return super().repr_instance(instance, level)


def short_repr(obj):
    r = _ShortRepr()
    r.maxother = 50
    return r.repr(obj)


def repr_args(args, kwargs, defaults={}):
    return ", ".join(
        [short_repr(a) for a in args]
        + [
            f"{k}={short_repr(v)}"
            for k, v in kwargs.items()
            if k not in defaults or defaults[k] != v
        ]
    )


def set_output(transformer, X):
    if not hasattr(transformer, "set_output"):
        return
    target_module = sbd.dataframe_module_name(X)
    try:
        transformer.set_output(transform=target_module)
    except Exception:
        # Some scikit-learn estimators have a set_output method, but it can
        # fail -- for example a Pipeline containing a step that doesn't have
        # set_output. The pipeline may still produce the right output type if
        # it does it by default, without the set_output call. So we allow
        # set_output to fail and attempt the transform, and an error is raised
        # if the output of transform has the wrong type.
        pass


def check_output(
    transformer, transform_input, transform_output, allow_column_list=True
):
    target_module = sbd.dataframe_module_name(transform_input)

    def has_correct_module(obj):
        return sbd.dataframe_module_name(obj) == target_module

    if sbd.is_dataframe(transform_output) and has_correct_module(transform_output):
        return transform_output
    if (
        allow_column_list
        and sbd.is_column(transform_output)
        and has_correct_module(transform_output)
    ):
        return transform_output
    if (
        allow_column_list
        and sbd.is_column_list(transform_output)
        and (not len(transform_output) or has_correct_module(transform_output[0]))
    ):
        return transform_output
    message = (
        f"{transformer.__class__.__name__}.fit_transform returned a result of type"
        f" {transform_output.__class__.__name__}, but a {target_module} DataFrame was"
        f" expected. If {transformer.__class__.__name__} is a custom transformer class,"
        f" please make sure that the output is a {target_module} container when the"
        f" input is a {target_module} container."
    )
    if not hasattr(transformer, "set_output"):
        message += (
            f" One way of enabling a transformer to output {target_module} DataFrames"
            " is inheriting from the sklearn.base.TransformerMixin class and defining"
            " the 'get_feature_names_out' method. See"
            " https://scikit-learn.org/stable/auto_examples/miscellaneous/plot_set_output.html"
            " for details."
        )
    raise TypeError(message)


class PassThrough(BaseEstimator):
    def fit(self, X, y=None):
        return self

    def fit_transform(self, X, y=None):
        return X

    def transform(self, X):
        return X


def format_duration(seconds):
    if seconds < 0:
        raise ValueError(
            f"format_duration only handles non-negative durations, got: {seconds}"
        )
    hours, rest = divmod(seconds, 3600)
    minutes, rest = divmod(rest, 60)
    return f"{hours:.0f}h {minutes:.0f}m {rest:.2g}s"
