# coding:utf-8
import hashlib
import inspect
import logging
import numpy as np

import pandas as pd
from scipy.sparse import csr_matrix, csc_matrix

from sklearn.model_selection import train_test_split, StratifiedKFold, KFold
from .cache import Cache, numpy_buffer
from .utils.main import idx, concat

logger = logging.getLogger("heamy.dataset")


class Dataset(object):
    """Dataset wrapper.

    Parameters
    ----------
    X_train : pd.DataFrame or np.ndarray, optional
    y_train : pd.DataFrame, pd.Series or np.ndarray, optional
    X_test : pd.DataFrame or np.ndarray, optional
    y_test : pd.DataFrame, pd.Series or np.ndarray, optional
    preprocessor : function, optional
        A callable function that returns preprocessed data.
    use_cache : bool, default True
        If `use_cache=True` then preprocessing step will be cached until function code is changed.

    Examples
    ----------
    >>> # function-based definition
    >>> from sklearn.datasets import load_boston
    >>> def boston_dataset():
    >>>     data = load_boston()
    >>>     X, y = data['data'], data['target']
    >>>     X_train, X_test, y_train, y_test = train_test_split(X, y, test_size=0.1, random_state=111)
    >>>     return X_train, y_train, X_test, y_test
    >>> dataset = Dataset(preprocessor=boston_dataset)

    >>> # class-based definition
    >>> class BostonDataset(Dataset):
    >>> def preprocess(self):
    >>>     data = load_boston()
    >>>     X, y = data['data'], data['target']
    >>>     X_train, X_test, y_train, y_test = train_test_split(X, y, test_size=0.1, random_state=111)
    >>>     return X_train, y_train, X_test, y_test
    """

    def __init__(
        self,
        X_train=None,
        y_train=None,
        X_test=None,
        y_test=None,
        preprocessor=None,
        use_cache=True,
    ):
        self._hash = None
        self.use_cache = use_cache
        self._preprocessor = preprocessor
        self._setup_data(X_train, y_train, X_test, y_test)

        if self.loaded:
            self._check_input()
            self._setup_columns()

        if hasattr(self.__class__, "preprocess"):
            self._preprocessor = self.preprocess

    def _setup_columns(self):
        if isinstance(self._X_train, (pd.Series, pd.DataFrame)):
            self.columns = self.X_train.columns.tolist()
        else:
            self.columns = []

    def _check_input(self):
        if not self.loaded:
            raise ValueError("Missing 2 required arrays: X_train and y_train.")

        if self._X_train.shape[0] != self._y_train.shape[0]:
            raise ValueError(