# score.py {{SCORE_VERSION}} generated by ADS {{ADS_VERSION}} on {{time_created}}
import os
import sys
from functools import lru_cache
import torch
import json
from typing import Dict, List
import numpy as np
import pandas as pd
from io import BytesIO
import base64
import logging
from random import randint


def get_torch_device():
    num_devices = torch.cuda.device_count()
    if num_devices == 0:
        return "cpu"
    if num_devices == 1:
        return "cuda:0"
    else:
        return f"cuda:{randint(0, num_devices-1)}"


model_name = '{{model_file_name}}'
device = torch.device(get_torch_device())

"""
   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_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 == "torchscript" %}
    print(f'Start loading {model_file_name} from model directory {model_dir} ...')
    the_model = torch.jit.load(os.path.join(model_dir, model_file_name))
    print(f"loading {model_file_name} is complete.")

{% elif model_serializer == "torch" %}
    print(f'Start loading {model_file_name} from model directory {model_dir} ...')
    model_state_dict = torch.load(os.path.join(model_dir, model_file_name))
    print(f"loading {model_file_name} is complete.")

    # User would need to provide reference to the TheModelClass and
    # construct the the_model instance first before loading the parameters.
    # the_model = TheModelClass(*args, **kwargs)
    try:
        the_model.load_state_dict(model_state_dict)
    except NameError as e:
        raise NotImplementedError("TheModelClass instance must be constructed before loading the parameters. Please modify the load_model() function in score.py." )
    except Exception as e:
        raise e

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

{% endif %}
    print("Model is successfully loaded.")
    the_model = the_model.to(device)
    return the_model

@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.pytorch_model.PyTorchModel.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-serialized 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 "torch.Tensor" in data_type:
        load_bytes = BytesIO(base64.b64decode(json_data.encode('utf-8')))
        return torch.load(load_bytes)
    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))

    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 json-serialized data to feed into predict function.

    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)

    # Add further data preprocessing if needed
    data = data.to(device)
    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.

    """
    if isinstance(yhat, torch.Tensor):
        return yhat.tolist()

    # Add further data postprocessing if needed
    return yhat

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
    input_schema_path: path of input schema.

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

    """
    inputs = pre_inference(data, input_schema_path)

    with torch.no_grad():
        yhat = post_inference(
        model(inputs).to("cpu")
    )
    return {'prediction': yhat}
