import pickle
from functools import partial

import pandas as pd
import pytest
from pandas import testing as pdt

from kartothek.io.dask.dataframe import read_dataset_as_ddf
from kartothek.io.eager import store_dataframes_as_dataset
from kartothek.io.testing.read import *  # noqa
from kartothek.io_components.metapartition import SINGLE_TABLE


@pytest.fixture()
def output_type():
    return "table"


def _read_as_ddf(
    dataset_uuid,
    store,
    factory=None,
    categoricals=None,
    tables=None,
    dataset_has_index=False,
    **kwargs,
):
    table = tables or SINGLE_TABLE
    if categoricals:
        categoricals = categoricals[table]
    ddf = read_dataset_as_ddf(
        dataset_uuid=dataset_uuid,
        store=store,
        factory=factory,
        categoricals=categoricals,
        table=table,
        **kwargs,
    )
    if categoricals:
        assert ddf._meta.dtypes["P"] == pd.api.types.CategoricalDtype(
            categories=["__UNKNOWN_CATEGORIES__"], ordered=False
        )
        if dataset_has_index:
            assert ddf._meta.dtypes["L"] == pd.api.types.CategoricalDtype(
                categories=[1, 2], ordered=False
            )
        else:
            assert ddf._meta.dtypes["L"] == pd.api.types.CategoricalDtype(
                categories=["__UNKNOWN_CATEGORIES__"], ordered=False
            )

    s = pickle.dumps(ddf, pickle.HIGHEST_PROTOCOL)
    ddf = pickle.loads(s)

    ddf = ddf.compute().reset_index(drop=True)

    def extract_dataframe(ix):
        df = ddf.iloc[[ix]].copy()
        for col in df.columns:
            if pd.api.types.is_categorical(df[col]):
                df[col] = df[col].cat.remove_unused_categories()
        return df.reset_index(drop=True)

    return [extract_dataframe(ix) for ix in ddf.index]


@pytest.fixture()
def bound_load_dataframes():
    return _read_as_ddf


def test_load_dataframe_categoricals_with_index(dataset_with_index_factory):
    func = partial(_read_as_ddf, dataset_has_index=True)
    test_read_dataset_as_dataframes(  # noqa: F405
        dataset_factory=dataset_with_index_factory,
        dataset=dataset_with_index_factory,
        store_session_factory=dataset_with_index_factory.store_factory,
        use_dataset_factory=True,
        bound_load_dataframes=func,
        use_categoricals=True,
        output_type="table",
        label_filter=None,
        dates_as_object=False,
    )


def test_read_ddf_from_categorical_partition(store_factory):
    df = pd.DataFrame({"x": ["a"]}).astype({"x": "category"})
    store_dataframes_as_dataset(
        dfs=[df], dataset_uuid="dataset_uuid", store=store_factory
    )
    ddf = read_dataset_as_ddf(
        dataset_uuid="dataset_uuid", store=store_factory, table="table"
    )
    df_expected = pd.DataFrame({"x": ["a"]})
    df_actual = ddf.compute(scheduler="sync")
    pdt.assert_frame_equal(df_expected, df_actual)

    ddf = read_dataset_as_ddf(
        dataset_uuid="dataset_uuid",
        store=store_factory,
        categoricals=["x"],
        table="table",
    )
    df_actual = ddf.compute(scheduler="sync")
    pdt.assert_frame_equal(df, df_actual)

pandas.DataFrame pandas.testing.assert_frame_equal pandas.api.types.CategoricalDtype pandas.api.types.is_categorical