# score.py {{SCORE_VERSION}} generated by ADS {{ADS_VERSION}} on {{time_created}}
import json
import os
import sys
import xgboost as xgb
import pandas as pd
import numpy as np
from functools import lru_cache
from io import BytesIO
import base64
import logging

model_name = '{{model_file_name}}'


"""
   Inference script. This script is used for prediction by scoring server when schema is known.
"""


@lru_cache(maxsize=10)
def load_model(model_file_name=model_name):
    """
    Loads model from the serialized format

    Returns
    -------
    model:  a model instance on which predict API can be invoked
    """
    model_xgb = xgb.Booster()
    model_dir = os.path.dirname(os.path.realpath(__file__))
    if model_dir not in sys.path:
        sys.path.insert(0, model_dir)
    contents = os.listdir(model_dir)
    if model_file_name not in contents:
        raise Exception(f'{model_file_name} is not found in model directory {model_dir}')

{% if model_serializer == "xgboost" %}
    print(f'Start loading {model_file_name} from model directory {model_dir} ...')
    model_xgb.load_model(os.path.join(os.path.dirname(os.path.realpath(__file__)), model_file_name))
    print("Model is successfully loaded.")
    return model_xgb

{% else %}
    raise NotImplementedError("`load_model()` needs to be implemented.")
{% endif %}

@lru_cache(maxsize=1)
def fetch_data_type_from_schema(input_schema_path=os.path.join(os.path.dirname(os.path.realpath(__file__)), "input_schema.json")):
    """
    Returns data type information fetch from input_schema.json.

    Parameters
    ----------
    input_schema_path: path of input schema.

    Returns
    -------
    data_type: data type fetch from input_schema.json.

    """
    data_type = {}
    if os.path.exists(input_schema_path):
        schema = json.load(open(input_schema_path))
        for col in schema['schema']:
            data_type[col['name']] = col['dtype']
    else:
        print("input_schema has to be passed in in order to recover the same data type. pass `X_sample` in `ads.model.framework.xgboost_model.XGBoostModel.prepare` function to generate the input_schema. Otherwise, the data type might be changed after serialization/deserialization.")
    return data_type


def deserialize(data, input_schema_path):
    """
    Deserialize json serialization data to data in original type when sent to predict.

    Parameters
    ----------
    data: serialized input data.
    input_schema_path: path of input schema.

    Returns
    -------
    data: deserialized input data.

    """
{% if data_deserializer == "json" %}
    if isinstance(data, bytes):
        logging.warning(
            "bytes are passed directly to the model. If the model expects a specific data format, you need to write the conversion logic in `deserialize()` yourself."
        )
        return data

    data_type = data.get('data_type', '') if isinstance(data, dict) else ''
    json_data = data.get('data', data) if isinstance(data, dict) else data

    if "numpy.ndarray" in data_type:
        load_bytes = BytesIO(base64.b64decode(json_data.encode('utf-8')))
        return np.load(load_bytes, allow_pickle=True)
    if "pandas.core.series.Series" in data_type:
        return pd.Series(json_data)
    if "pandas.core.frame.DataFrame" in data_type or isinstance(json_data, str):
        return pd.read_json(json_data, dtype=fetch_data_type_from_schema(input_schema_path))
    if isinstance(json_data, dict):
        return pd.DataFrame.from_dict(json_data)

    return json_data

{% elif data_deserializer == "cloudpickle" %}
    import cloudpickle
    from pickle import UnpicklingError
    deserialized_data = data
    try:
        deserialized_data = cloudpickle.loads(data)
    except TypeError:
        pass
    except UnpicklingError:
        logger.warning(
            "bytes are passed directly to the model. If the model expects a specific data format, you need to write the conversion logic in `deserialize()` yourself."
        )

    return deserialized_data

{% else %}
    # Add further data deserialization if needed
    return data
{% endif %}

def pre_inference(data, input_schema_path):
    """
    Preprocess data

    Parameters
    ----------
    data: Data format as expected by the predict API of the core estimator.
    input_schema_path: path of input schema.

    Returns
    -------
    data: Data format after any processing.

    """
    data = deserialize(data, input_schema_path)
    return data

def post_inference(yhat):
    """
    Post-process the model results

    Parameters
    ----------
    yhat: Data format after calling model.predict.

    Returns
    -------
    yhat: Data format after any processing.

    """
    return yhat.tolist()

def predict(data, model=load_model(), input_schema_path=os.path.join(os.path.dirname(os.path.realpath(__file__)), "input_schema.json")):
    """
    Returns prediction given the model and data to predict

    Parameters
    ----------
    model: Model instance returned by load_model API
    data: Data format as expected by the predict API of the core estimator. For eg. in case xgboost models it could be numpy array/List of list/Pandas DataFrame
    input_schema_path: path of input schema.

    Returns
    -------
    predictions: Output from scoring server
        Format: {'prediction': output from model.predict method}

    """
    data = pre_inference(data, input_schema_path)
    data = xgb.DMatrix(data)
    return {'prediction': post_inference(model.predict(data))}
