import contextlib
import logging
import math
import shutil
import tempfile
import uuid

import numpy as np
import pandas as pd
import tlz as toolz

from .. import base, config
from ..base import compute, compute_as_if_collection, is_dask_collection, tokenize
from ..highlevelgraph import HighLevelGraph
from ..layers import ShuffleLayer, SimpleShuffleLayer
from ..sizeof import sizeof
from ..utils import M, digit
from . import methods
from .core import DataFrame, Series, _Frame, map_partitions, new_dd_object
from .dispatch import group_split_dispatch, hash_object_dispatch

logger = logging.getLogger(__name__)


def _calculate_divisions(
    df,
    partition_col,
    repartition,
    npartitions,
    upsample=1.0,
    partition_size=128e6,
):
    """
    Utility function to calculate divisions for calls to `map_partitions`
    """
    sizes = df.map_partitions(sizeof) if repartition else []
    divisions = partition_col._repartition_quantiles(npartitions, upsample=upsample)
    mins = partition_col.map_partitions(M.min)
    maxes = partition_col.map_partitions(M.max)
    divisions, sizes, mins, maxes = base.compute(divisions, sizes, mins, maxes)
    divisions = methods.tolist(divisions)
    if type(sizes) is not list:
        sizes = methods.tolist(sizes)
    mins = methods.tolist(mins)
    maxes = methods.tolist(maxes)

    empty_dataframe_detected = pd.isnull(divisions).all()
    if repartition or empty_dataframe_detected:
        total = sum(sizes)
        npartitions = max(math.ceil(total / partition_size), 1)
        npartitions = min(npartitions, df.npartitions)
        n = len(divisions)
        try:
            divisions = np.interp(
                x=np.linspace(0, n - 1, npartitions + 1),
                xp=np.linspace(0, n - 1, n),
                fp=divisions,
            ).tolist()
        except (TypeError, ValueError):  # str type
            indexes = np.linspace(0, n - 1, npartitions + 1).astype(int)
            divisions = [divisions[i] for i in indexes]

    mins = remove_nans(mins)
    maxes = remove_nans(maxes)
    if pd.api.types.is_categorical_dtype(partition_col.dtype):
        dtype = partition_col.dtype
        mins = pd.Categorical(mins, dtype=dtype).codes.tolist()
        maxes = pd.Categorical(maxes, dtype=dtype).codes.tolist()

    return divisions, mins, maxes


def sort_values(
    df,
    by,
    npartitions=None,
    ascending=True,
    na_position="last",
    upsample=1.0,
    partition_size=128e6,
    sort_function=None,
    sort_function_kwargs=None,
    **kwargs,
):
    """See DataFrame.sort_values for docstring"""
    if na_position not in ("first", "last"):
        raise ValueError("na_position must be either 'first' or 'last'")
    if not isinstance(by, str):
        # support ["a"] as input
        if isinstance(by, list) and len(by) == 1 and isinstance(by[0], str):
            by = by[0]
        else:
            raise NotImplementedError(
                "Dataframe only supports sorting by a single column which must "
                "be passed as a string or a list of a single string.\n"
                "You passed %s" % str(by)
            )
    if npartitions == "auto":
        repartition = True
        npartitions = max(100, df.npartitions)
    else:
        if npartitions is None:
            npartitions = df.npartitions
        repartition = False

    sort_by_col = df[by]

    divisions, mins, maxes = _calculate_divisions(
        df, sort_by_col, repartition, npartitions, upsample, partition_size
    )

    if len(divisions) == 2:
        return df.repartition(npartitions=1).map_partitions(
            sort_function, **sort_function_kwargs
        )

    if not isinstance(ascending, bool):
        # support [True] as input
        if (
            is
instance(ascending, list)
            and len(ascending) == 1
            and isinstance(ascending[0], bool)
        ):
            ascending = ascendin
g[0]
        else:
            raise NotImplementedError(
                f"Dask currently only supports a single boolean for ascending. You passed {str(ascending)}"
            )

    if (
        all(not pd.isna(x) for x in divisions)
        and mins == sorted(mins, reverse=not ascending)
        and maxes == sorted(maxes, reverse=not ascending)
        and all(
            mx < mn
            for mx, mn in zip(
                maxes[:-1] if ascending else maxes[1:],
                mins[1:] if ascending else mins[:-1],
            )
        )
        and npartitions == df.npartitions
    ):
        # divisions are in the right place
        return df.map_partitions(sort_function, **sort_function_kwargs)

    df = rearrange_by_divisions(
        df,
        by,
        divisions,
        ascending=ascending,
        na_position=na_position,
        duplicates=False,
    )
    df = df.map_partitions(sort_function, **sort_function_kwargs)
    return df


def set_index(
    df,
    index,
    npartitions=None,
    shuffle=None,
    compute=False,
    drop=True,
    upsample=1.0,
    divisions=None,
    partition_size=128e6,
    **kwargs,
):
    """See _Frame.set_index for docstring"""
    if isinstance(index, Series) and index._name == df.index._name:
        return df
    if isinstance(index, (DataFrame, tuple, list)):
        # Accept ["a"], but not [["a"]]
        if (
            isinstance(index, list)
            and len(index) == 1
            and not isinstance(index[0], list)  # if index = [["a"]], leave it that way
        ):
            index = index[0]
        else:
            raise NotImplementedError(
                "Dask dataframe does not yet support multi-indexes.\n"
                "You tried to index with this index: %s\n"
                "Indexes must be single columns only." % str(index)
            )

    if npartitions == "auto":
        repartition = True
        npartitions = max(100, df.npartitions)
    else:
        if npartitions is None:
            npartitions = df.npartitions
        repartition = False

    if not isinstance(index, Series):
        index2 = df[index]
    else:
        index2 = index

    if divisions is None:
        divisions, mins, maxes = _calculate_divisions(
            df, index2, repartition, npartitions, upsample, partition_size
        )

        if (
            mins == sorted(mins)
            and maxes == sorted(maxes)
            and all(mx < mn for mx, mn in zip(maxes[:-1], mins[1:]))
            and npartitions == df.npartitions
        ):
            divisions = mins + [maxes[-1]]
            result = set_sorted_index(df, index, drop=drop, divisions=divisions)
            return result.map_partitions(M.sort_index)

    return set_partition(
        df, index, divisions, shuffle=shuffle, drop=drop, compute=compute, **kwargs
    )


def remove_nans(divisions):
    """Remove nans from divisions

    These sometime pop up when we call min/max on an empty partition

    Examples
    --------
    >>> remove_nans((np.nan, 1, 2))
    [1, 1, 2]
    >>> remove_nans((1, np.nan, 2))
    [1, 2, 2]
    >>> remove_nans((1, 2, np.nan))
    [1, 2, 2]
    """
    divisions = list(divisions)

    for i in range(len(divisions) - 2, -1, -1):
        if pd.isnull(divisions[i]):
            divisions[i] = divisions[i + 1]

    for i in range(len(divisions) - 1, -1, -1):
        if not pd.isnull(divisions[i]):
            for j in range(i + 1, len(divisions)):
                divisions[j] = divisions[i]
            break

    return divisions


def set_partition(
    df, index, divisions, max_branch=32, drop=True, shuffle=None, compute=None
):
    """Group DataFrame by index

    Sets a new index and partitions data along that index according to
    divisions.  Divisions are often found by computing approximate quantiles.
    The function ``set_index`` will do both of these steps.

    Parameters
    ----------
    df: DataFrame/Series
        Data that we want to re-partition
    index: string or Series
        Column to become the new index
    divisions: list
        Values to form new divisions between partitions
    drop: bool, default True
        Whether to delete columns to be used as the new index
    shuffle: str (optional)
     