"""
model.py
This is a model template that shows the interface expected by MLOps framework
for an example model (scikit-learn's DummyClassifier).

Internal elements start by "_", as usually done in Python. This points to parts
of this file that are specific to the example and not required by the framework.

We will develop a single service, namely prediction/inference. Refer to the history for code needed to include a training service (was never used)
"""
import os
from typing import Any, Dict, List

from joblib import load

from akerbp.mlops.core.logger import get_logger

logger = get_logger(__name__)

_model_name = "dummy.joblib"


class ModelException(Exception):
    """
    This exception should be raised when there are errors (required by the
    framework). It can be subclassed if more specific exceptions are needed, but
    this isn't required by the framework.
    """

    pass


class ModelIO(ModelException):
    pass


class ModelDataValidation(ModelException):
    pass


def _load_model(folder_path: str) -> Any:
    """
    Load model artifact generated by `train`
    Input:
      - folder path (string) that contains model artifacts
    Output:
      - clf: model object required by `predict`. If there are multiple objects,
      you can return a tuple e.g. `(model_1, model_2)`
    """
    source_path = os.path.join(folder_path, _model_name)
    try:
        clf = load(source_path)
    except Exception as error:
        error_type = type(error).__name__
        error_message = f"Could not load model. {error_type}: {error}"
        logger.error(error_message)
        raise ModelIO(error_message) from error
    logger.info(f"Read model {clf=} from {source_path=}")
    return clf


def initialization(folder_path: str, secrets=None) -> Any:
    """
    Initialize objects required by `predict`.
    This can involve e.g. loading model artifacts generated by `train`,
    instantiate objects required to read data, etc.

    Input:
      - folder path (string) that contains model artifacts
      - secrets: dictionary with api keys, keys are "files" and "data"

    Output: object required by predict. If there are multiple objects,
      this could be a tuple, e.g. `(model, client)`
    """
    model = _load_model(folder_path)
    return model


def _validate_input(data: Dict):
    """
    Validate features sent by prediction or training service API.
    Generate ModelDataValidation with a useful message for the user.
    """
    if not isinstance(data, dict):
        raise ModelDataValidation("Data should be a dictionary")
    if "input" not in data.keys():
        raise ModelDataValidation("Input dictionary should have a column 'input'")
    if not isinstance(data["input"], list):
        raise ModelDataValidation("field 'input' should contain an array")


def _process(data_dict: Dict) -> List:
    """
    Transform data provided by the prediction service API into the data
    expected by the model.
    """
    input_dict = data_dict["input"][0]["input_logs"]
    return [input_dict[key][0] for key in input_dict.keys()]


def predict(data_dict, init_object, secrets=None) -> Any:
    """
    Predict data provided through the model service API

    Inputs:
      - data_dict: data (dictionary)
      - init_object: object provided by `load_model`. If `initialization` returns
          a tuple, you can unpack it, e.g. `model, client = init_object`
      - secrets: dictionary with api keys, keys are "files" and "data"
    Output:
      - y: prediction provided by `model` for `x`
        Note: prediction should be jsonifyable, which means for example that any
        np array should be converted to list
    """
    model = init_object
    _validate_input(data_dict)
    x = _process(data_dict)
    y = model.predict(x)
    logger.info(f"Prediction for {x=} was {y=}")
    prediction = y.tolist()
    return prediction
