"""
Code for processing datasets using scikit-learn.
"""
import inspect
import logging
from typing import List, Optional

import numpy as np
from sklearn.base import BaseEstimator

from deepchem.models import Model
from deepchem.data import Dataset
from deepchem.trans import Transformer
from deepchem.utils.data_utils import load_from_disk, save_to_disk
from deepchem.utils.typing import OneOrMany

logger = logging.getLogger(__name__)


class SklearnModel(Model):
  """Wrapper class that wraps scikit-learn models as DeepChem models.

  When you're working with scikit-learn and DeepChem, at times it can
  be useful to wrap a scikit-learn model as a DeepChem model. The
  reason for this might be that you want to do an apples-to-apples
  comparison of a scikit-learn model to another DeepChem model, or
  perhaps you want to use the hyperparameter tuning capabilities in
  `dc.hyper`. The `SklearnModel` class provides a wrapper around scikit-learn
  models that allows scikit-learn models to be trained on `Dataset` objects
  and evaluated with the same metrics as other DeepChem models.

  Example
  ------
  >>> import deepchem as dc
  >>> import numpy as np
  >>> from sklearn.linear_model import LinearRegression
  >>> # Generating a random data and creating a dataset
  >>> X, y = np.random.randn(5, 1), np.random.randn(5)
  >>> dataset = dc.data.NumpyDataset(X, y)
  >>> # Wrapping a Sklearn Linear Regression model using DeepChem models API
  >>> sklearn_model = LinearRegression()
  >>> dc_model = dc.models.SklearnModel(sklearn_model)
  >>> dc_model.fit(dataset)  # fitting dataset

  Notes
  -----
  All `SklearnModels` perform learning solely in memory. This means that it
  may not be possible to train `SklearnModel` on large `Dataset`s.
  """

  def __init__(self,
               model: BaseEstimator,
               model_dir: Optional[str] = None,
               **kwargs):
    """
    Parameters
    ----------
    model: BaseEstimator
      The model instance which inherits a scikit-learn `BaseEstimator` Class.
    model_dir: str, optional (default None)
      If specified the model will be stored in this directory. Else, a
      temporary directory will be used.
    model_instance: BaseEstimator (DEPRECATED)
      The model instance which inherits a scikit-learn `BaseEstimator` Class.
    kwargs: dict
      kwargs['use_weights'] is a bool which determines if we pass weights into
      self.model.fit().
    """
    if 'model_instance' in kwargs:
      model_instance = kwargs['model_instance']
      if model is not None:
        raise ValueError(
            "Can not use both model and model_instance argument at the same time."
        )
      logger.warning(
          "model_instance argument is deprecated and will be removed in a future version of DeepChem."
          "Use model argument instead.")
      model = model_instance

    super(SklearnModel, self).__init__(model, model_dir, **kwargs)
    if 'use_weights' in kwargs:
      self.use_weights = kwargs['use_weights']
    else:
      self.use_weights = True

    if self.use_weights and self.model is not None:
      # model is None when reloading a model
      if 'sample_weight' not in inspect.getfullargspec(self.model.fit).args:
        self.use_weights = False
        logger.info("The model does not support training with weights."
                    "Hence, not using weight of datapoint for training")

  def fit(self, dataset: Dataset) -> None:
    """Fits scikit-learn model to data.

    Parameters
    ----------
    dataset: Dataset
      The `Dataset` to train this model on.
    """
    X = dataset.X
    y = np.squeeze(dataset.y)
    w = np.squeeze(dataset.w)
    # Some scikit-learn models don't use weights.
    if self.use_weights:
      self.model.fit(X, y, w)
      return
    self.model.fit(X, y)

  def predict_on_batch(self, X: np.typing.ArrayLike) -> np.ndarray:
    """Makes predictions on batch of data.

    Parameters
    ----------
    X: np.ndarray
      A numpy array of features.

    Returns
    -------
    np.ndarray
      The value is a return value of `predict_proba` or `predict` method
      of the scikit-learn model. If the scikit-learn model has both methods,
      the value is always a return value of `predict_proba`.
    """
    try:
      return self.model.predict_proba(X)
    except AttributeError:
      return self.model.predict(X)

  def predict(self,
              X: Dataset,
              transformers: List[Transformer] = []) -> OneOrMany[np.ndarray]:
    """Makes predictions on dataset.

    Parameters
    ----------
    dataset: Dataset
      Dataset to make prediction on.
    transformers: List[Transformer]
      Transformers that the input data has been transformed by. The output
      is passed through these transformers to undo the transformations.
    """
    return super(SklearnModel, self).predict(X, transformers)

  def save(self):
    """Saves scikit-learn model to disk using joblib."""
    save_to_disk(self.model, self.get_model_filename(self.model_dir))

  def reload(self):
    """Loads scikit-learn model from joblib file on disk."""
    self.model = load_from_disk(self.get_model_filename(self.model_dir))
