"""Implementations of different data feeders to provide data for TF trainer."""

#  Copyright 2015-present The Scikit Flow Authors. All Rights Reserved.
#
#  Licensed under the Apache License, Version 2.0 (the "License");
#  you may not use this file except in compliance with the License.
#  You may obtain a copy of the License at
#
#   http://www.apache.org/licenses/LICENSE-2.0
#
#  Unless required by applicable law or agreed to in writing, software
#  distributed under the License is distributed on an "AS IS" BASIS,
#  WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
#  See the License for the specific language governing permissions and
#  limitations under the License.
from __future__ import absolute_import
from __future__ import division
from __future__ import print_function

import itertools
import math

import six
from six.moves import xrange  # pylint: disable=redefined-builtin

import numpy as np

from tensorflow.python.ops import array_ops
from tensorflow.python.framework import dtypes

from .pandas_io import HAS_PANDAS, extract_pandas_data, extract_pandas_matrix, extract_pandas_labels
from .dask_io import HAS_DASK, extract_dask_data, extract_dask_labels


def _get_in_out_shape(x_shape, y_shape, n_classes, batch_size):
  """Returns shape for input and output of the data feeder."""
  x_shape = list(x_shape[1:]) if len(x_shape) > 1 else [1]
  input_shape = [batch_size] + x_shape
  if y_shape is None:
    return input_shape, None
  y_shape = list(y_shape[1:]) if len(y_shape) > 1 else []
  # Skip first dimension if it is 1.
  if y_shape and y_shape[0] == 1:
    y_shape = y_shape[1:]
  if n_classes > 1:
    output_shape = [batch_size] + y_shape + [n_classes]
  else:
    output_shape = [batch_size] + y_shape
  return input_shape, output_shape


def _data_type_filter(X, y):
  """Filter data types into acceptable format"""
  if HAS_DASK:
    X = extract_dask_data(X)
    if y is not None:
      y = extract_dask_labels(y)
  if HAS_PANDAS:
    X = extract_pandas_data(X)
    if y is not None:
      y = extract_pandas_labels(y)
  return X, y


def _is_iterable(X):
  return hasattr(X, 'next') or hasattr(X, '__next__')


def setup_train_data_feeder(X, y, n_classes, batch_size):
  """Create data feeder, to sample inputs from dataset.

  If X and y are iterators, use StreamingDataFeeder.

  Args:
    X: numpy, pandas or Dask matrix or iterable.
    y: numpy, pandas or Dask array or iterable.
    n_classes: number of classes.
    batch_size: size to split data into parts.

  Returns:
    DataFeeder object that returns training data.
  """
  X, y = _data_type_filter(X, y)
  if HAS_DASK:
    import dask.dataframe as dd
    allowed_classes = (dd.Series, dd.DataFrame)
    if (isinstance(X, allowed_classes) and
        (y is None or isinstance(y, allowed_classes))):
      data_feeder_cls = DaskDataFeeder
    else:
      data_feeder_cls = DataFeeder
  else:
    data_feeder_cls = DataFeeder

  if _is_iterable(X):
    if y is not None and not _is_iterable(y):
      raise ValueError('Both X and y should be iterators for '
                       'streaming learning to work.')
    data_feeder_cls = StreamingDataFeeder
  return data_feeder_cls(X, y, n_classes, batch_size)


def _batch_data(X, batch_size):
  chunk = []
  for data in X:
    chunk.append(data)
    if batch_size > 0 and len(chunk) >= batch_size:
      yield np.matrix(chunk)
      chunk = []
  yield np.matrix(chunk)


def setup_predict_data_feeder(X, batch_size=-1):
  """Returns an iterable for feeding into predict step.

  Args:
    X: numpy, pandas, Dask array or iterable.
    batch_size: Size of batches to split data into.
      If negative, returns one batch of full size.

  Returns:
    List or iterator of parts of data to predict on.
  """
  if HAS_DASK:
    X = extract_dask_data(X)
  if HAS_PANDAS:
    X = extract_pandas_data(X)
  if _is_iterable(X):
    return _batch_data(X, batch_size)
  if len(X.shape) == 1:
    X = np.reshape(X, (-1, 1))
  if batch_size > 0:
    n_batches = int(math.ceil(float(len(X)) / batch_size))
    return [X[i * batch_size:(i + 1) * batch_size] for i in xrange(n_batches)]
  return [X]


def setup_processor_data_feeder(X):
  """Sets up processor iterable.

  Args:
    X: numpy, pandas or iterable.

  Returns:
    Iterable of data to process.
  """
  if HAS_PANDAS:
    X = extract_pandas_matrix(X)
  return X


def check_array(array, dtype):
  """Checks array on dtype and convers it if different.

  Args:
    array: Input array.
    dtype: Expected dtype.

  Returns:
    Original array or converted.
  """
  array = np.array(array, dtype=dtype, order=None, copy=False)
  return array


class DataFeeder(object):
  """Data feeder is an example class to sample data for TF trainer.

  Parameters:
    X: feature Nd numpy matrix of shape [n_samples, n_features, ...].
    y: target vector, either floats for regression or class id for
      classification. If matrix, will consider as a sequence
      of targets. Can be None for unsupervised setting.
    n_classes: number of classes, 0 and 1 are considered regression.
    batch_size: mini batch size to accumulate.
    random_state: numpy RandomState object to reproduce sampling.

  Attributes:
    X: input features.
    y: input target.
    n_classes: number of classes.
    batch_size: mini batch size to accumulate.
    input_shape: shape of the input.
    output_shape: shape of the output.
    input_dtype: dtype of input.
    output_dtype: dtype of output.
  """

  def __init__(self, X, y, n_classes, batch_size, random_state=None):
    x_dtype = np.int64 if X.dtype == np.int64 else np.float32
    y_dtype = np.int64 if n_classes > 1 else np.float32
    self.X = check_array(X, dtype=x_dtype)
    self.y = (None if y is None else check_array(y, dtype=y_dtype))
    self.n_classes = n_classes
    self.batch_size = batch_size
    self.input_shape, self.output_shape = _get_in_out_shape(self.X.shape, None
                                                            if self.y is None
                                                            else self.y.shape,
                                                            n_classes,
                                                            batch_size)
    # Input dtype matches dtype of X.
    self.input_dtype = self.X.dtype
    # Output dtype always float32 (because for classification we use
    # one-hot vectors.
    self.output_dtype = np.float32
    self.random_state = np.random.RandomState(
        42) if random_state is None else random_state
    self.indices = self.random_state.permutation(self.X.shape[0])
    self.offset = 0
    self.epoch = 0
    self._epoch_placeholder = None

  def make_epoch_variable(self):
    """Adds a placeholder variable for the epoch to the graph.

    Returns:
      The epoch placeholder.
    """
    self._epoch_placeholder = array_ops.placeholder(dtypes.int32, [1],
                                                    name='epoch')
    return self._epoch_placeholder

  def input_builder(self):
    """Builds inputs in the graph.

    Returns:
      Two placeholders for inputs and outputs.
    """
    input_shape = [None] + self.input_shape[1:]
    self._input_placeholder = array_ops.placeholder(
        dtypes.as_dtype(self.input_dtype),
        input_shape,
        name='input')
    if self.output_shape is None:
      self._output_placeholder = None
    else:
      output_shape = [None] + self.output_shape[1:]
      self._output_placeholder = array_ops.placeholder(
          dtypes.as_dtype(self.output_dtype),
          output_shape,
          name='output')
    return self._input_placeholder, self._output_placeholder

  def set_placeholders(self, input_placeholder, output_placeholder):
    """Sets placeholders for this data feeder.

    Args:
      input_placeholder: Placeholder for `X` variable. Should match shape
        of the examples in the X dataset.
      output_placeholder: Placeholder for `y` variable. Should match
        shape of the examples in the y dataset. Can be None.
    """
    self._input_placeholder = input_placeholder
    self._output_placeholder = output_placeholder

  def get_feed_params(self):
    """Function returns a dict with data feed params while training.

    Returns:
      A dict with data feed params while training.
    """
    return {
        'epoch': self.epoch,
        'offset': self.offset,
        'batch_size': self.batch_size
    }

  def get_feed_dict_fn(self):
    """Returns a function, that will sample data and provide it to given
    placeholders.

    Args:
      input_placeholder: tf.Placeholder for input features mini batch.
      output_placeholder: tf.Placeholder for output targets.
    Returns:
      A function that when called samples a random subset of batch size
      from X and y.
    """
    assert self._input_placeholder != None

    def _feed_dict_fn():
      feed_dict = {}
      if self._epoch_placeholder is not None:
        feed_dict[self._epoch_placeholder.name] = self.epoch

      # take random indices
      batch_indices = self.indices[self.offset:self.offset + self.batch_size]

      # assign input features from random indices
      inp = np.array(self.X[batch_indices]).reshape((batch_indices.shape[0], 1)) \
          if len(self.X.shape) == 1 else self.X[batch_indices]
      feed_dict[self._input_placeholder.name] = inp

      # move offset and reset it if necessary
      self.offset += self.batch_size
      if self.offset >= self.X.shape[0]:
        self.indices = self.random_state.permutation(self.X.shape[0])
        self.offset = 0
        self.epoch += 1

      # return early if there are no labels
      if self._output_placeholder is None:
        return feed_dict

      # assign labels from random indices
      self.output_shape[0] = batch_indices.shape[0]
      out = np.zeros(self.output_shape, dtype=self.output_dtype)
      for i in xrange(out.shape[0]):
        sample = batch_indices[i]
        if self.n_classes > 1:
          if len(self.output_shape) == 2:
            out.itemset((i, self.y[sample]), 1.0)
          else:
            for idx, value in enumerate(self.y[sample]):
              out.itemset(tuple([i, idx, value]), 1.0)
        else:
          out[i] = self.y[sample]
      feed_dict[self._output_placeholder.name] = out

      return feed_dict

    return _feed_dict_fn


class StreamingDataFeeder(DataFeeder):
  """Data feeder for TF trainer that reads data from iterator.

  Streaming data feeder allows to read data as it comes it from disk or
  somewhere else. It's custom to have this iterators rotate infinetly over
  the dataset, to allow control of how much to learn on the trainer side.

  Parameters: