import os
import pickle as pkl
import json
import numpy as np
import xgboost as xgb

from sagemaker_containers.beta.framework import (
    content_types,
    encoders,
    env,
    modules,
    transformer,
    worker,
)

from sagemaker_xgboost_container import encoder as xgb_encoders


def input_fn(input_data, content_type):
    if content_type == content_types.JSON:
        obj = json.loads(input_data)
        features = obj["instances"][0]["features"]
        array = np.array(features).reshape((1, -1))
        return xgb.DMatrix(array)
    else:
        return xgb_encoders.decode(input_data, content_type)


def model_fn(model_dir):
    model_file = model_dir + "/model.bin"
    model = pkl.load(open(model_file, "rb"))
    return model


def output_fn(prediction, accept):

    pred_array_value = np.array(prediction)
    score = pred_array_value[0]

    if accept == "application/json":
        predicted_label = 1 if score > 0.5 else 0
        return_value = {
            "predictions": [{"score": score.astype(float), "predicted_label": predicted_label}]
        }
        return worker.Response(json.dumps(return_value), mimetype=accept)
    elif accept == "text/csv":
        return_value = "yes" if score > 0.5 else "no"
        return worker.Response(encoders.encode(prediction, accept), mimetype=accept)
    else:
        raise RuntimeException("{} accept type is not supported.".format(accept))