import pandas as pd
import logging
import os
import joblib

"""
This file was autogenerated by: drum new model --language python
Generation date: Tue Aug 11 16:53:33 2020

Note: this is an example of custom.py file.
   Below are all the hooks you can use to provide your own implementation.
   All hooks are currently commented out so uncomment a hook function in 
   order to use it.
"""


def init(**kwargs):
   """
   This hook can be implemented to adjust logic in the training and scoring mode.
   init is called once the code is started.

   :param kwargs: additional keyword arguments to the function.
   code_dir - code folder passed in --code_dir argument
   """
   pass


def load_model(input_dir):
   """
   This hook can be implemented to adjust logic in the scoring mode.

   load_model hook provides a way to implement model loading your self.
   This function should return an object that represents your model. This object will
   be passed to the predict hook for performing predictions.
   This hook can be used to load supported models if your model has multiple artifacts, or
   for loading models that drum does not natively support

   :param input_dir: the directory to load serialized models from
   :returns: Object containing the model - the predict hook will get this object as a parameter
   """
   with open(os.path.join(input_dir, "catboost_trained.pkl"), "rb") as picklefile:
      try:
         custom_model = joblib.load(picklefile)
      except TypeError:
         custom_model = joblib.load(picklefile)
   return custom_model


def transform(data, model):
   """
   This hook can be implemented to adjust logic in the scoring mode.

   transform(data: DataFrame, model: Any) -> DataFrame

   Intended to apply transformations to the prediction data before making predictions.
   This is most useful if drum supports the model's library, but your model requires additional
   data processing before it can make predictions

   :param data: dataframe given to drum to make predictions on
   :param model: is the deserialized model loaded by drum or by load_model hook , if supplied
   :returns: a dataframe after transformation needed
   """
   return data


def score(data, model, **kwargs):
   """
   This hook can be implemented to adjust logic in the scoring mode.

   This method should return predictions as a dataframe with the following format:

   Binary Classification:
   Must have columns for each class label with floating-point class probabilities as values.
   Each row should sum to 1.0

   Regression:
   Must have a single column called "Predictions" with numerical values

   This hook is only needed if you would like to use drum with a framework not natively
   supported by the tool.

   :param data: the dataframe to make predictions against. If transform is supplied, data
       will be the transformed data.
   :param model: is the deserialized model loaded by drum or by load_model hook, if supplied
   :param kwargs: additional keyword arguments to the function. If model is classification,
   positive_class_label and negative_class_label will be provided in kwargs.
   :returns: a dataframe, see documentation above on the structure of the dataframe to return.
   """

   # Creating a dataframe with a column named Predictions with a constant value of 1
   predictions = pd.DataFrame(model.predict_proba(data), columns=["True", "False"])
   return predictions


def post_process(predictions, model):
   """
   This hook can be implemented to adjust logic in the scoring mode.

   This method should return predictions as a dataframe with the following format:

   Binary Classification:
   Must have columns for each class label with floating- point class probabilities as values.
       Each row should sum to 1.0

   Regression:
   Must have a single column called `Predictions` with numerical values

   This method is only needed if your model's output does not match the above expectations

   :param predictions: is the dataframe of predictions produced by `cmrun` or by the
       `score` hook, if supplied
   :param model: the deserialized model loaded by `cmrun` or by `load_model`, if supplied
   :returns: dataframe with the results to return
   """
   return predictions

def fit(X, y, output_dir, **kwargs):

   """
   This hook must be implemented with your fitting code, for running drum in the fit mode.

   Parameters
   ----------
   :param X: pd.DataFrame - training data to perform fit on
   :param y: pd.Series - target data to perform fit on
   :param output_dir: str - the path to write output. This is the path provided in '--output' parameter of the 'drum fit' command.
   :param kwargs: Dict[str, Any] - additional optional keyword arguments to the function:
       - 'class_order': [List[str] - a two element long list dictating the order of classes which should be used for modeling
       - 'row_weights': np.ndarray - an array of non-negative numeric values which can be used to dictate how important a row is

   :returns: nothing
   """
   return None