import ast
import json
import os

import numpy as np
import pandas as pd
from autogluon.multimodal import MultiModalPredictor
from constants import (
    ALLOWED_INPUT_FORMATS,
    ALLOWED_OUTPUT_FORMATS,
    BRACKET_FORMATTER,
    COMMA_DELIMITER,
    IMAGE_COLUMN_NAME,
    JSON_FORMAT,
    LABELS,
    NUM_GPU,
    PREDICTED_LABEL,
    PROBABILITIES,
    PROBABILITY,
    SAGEMAKER_INFERENCE_OUTPUT,
)
from utils import infer_type_and_cast_value

INFERENCE_OUTPUT = (
    infer_type_and_cast_value(os.getenv(SAGEMAKER_INFERENCE_OUTPUT))
    if SAGEMAKER_INFERENCE_OUTPUT in os.environ
    else [PREDICTED_LABEL]
)
NUM_GPUS = infer_type_and_cast_value(os.getenv(NUM_GPU))


def generate_single_csv_line_inference_selection(data):
    """Generate a single csv line response.

    :param data: list of output generated from the model
    :return: csv line for the predictions
    """
    contents: str
    for single_prediction in data:
        contents = (
            BRACKET_FORMATTER.format(single_prediction)
            if isinstance(single_prediction, list)
            else str(single_prediction)
        )
    return contents


def model_fn(model_dir):
    """Load model from previously saved artifact.

    :param model_dir: local path to the model directory
    :return: loaded model
    """
    predictor = MultiModalPredictor.load(model_dir)
    if NUM_GPUS is not None:
        predictor._config.env.num_gpus = NUM_GPUS

    return predictor


def convert_to_json_compatible_type(value):
    """Convert the input value to a JSON compatible type.

    :param value: input value
    :return: JSON compatible value
    """
    string_value = "{}".format(value)
    try:
        return ast.literal_eval(string_value)
    except Exception:
        return string_value


def transform_fn(model, request_body, input_content_type, output_content_type):
    """Transform function for serving inference requests.

    If INFERENCE_OUTPUT is provided, then the predictions are generated in the requested format and concatenated in the
    same order. Otherwise, prediction_labels are generated by default.

    :param model: loaded model
    :param request_body: request body
    :param input_content_type: content type of the input
    :param output_content_type: content type of the response
    :return: prediction response
    """
    if input_content_type.lower() not in ALLOWED_INPUT_FORMATS:
        raise Exception(
            f"{input_content_type} input content type not supported. Supported formats are {ALLOWED_INPUT_FORMATS}"
        )

    if output_content_type.lower() not in ALLOWED_OUTPUT_FORMATS:
        raise Exception(
            f"{output_content_type} output content type not supported. Supported formats are {ALLOWED_OUTPUT_FORMATS}"
        )

    data = pd.DataFrame({IMAGE_COLUMN_NAME: [request_body]})
    result_dict = dict()

    result = []
    inference_output_list = (
        INFERENCE_OUTPUT if isinstance(INFERENCE_OUTPUT, list) else [INFERENCE_OUTPUT]
    )
    for output_type in inference_output_list:
        if output_type == PREDICTED_LABEL:
            prediction = model.predict(data)
            result_dict[PREDICTED_LABEL] = convert_to_json_compatible_type(prediction.squeeze())
        elif output_type == PROBABILITIES:
            predict_probs = model.predict_proba(data)
            prediction = predict_probs.to_numpy()
            result_dict[PROBABILITIES] = predict_probs.squeeze().tolist()
        elif output_type == LABELS:
            labels = model.class_labels
            prediction = np.array([labels]).astype("str")
            result_dict[LABELS] = labels.tolist()
        else:
            predict_probabilities = model.predict_proba(data).to_numpy()
            prediction = np.max(predict_probabilities, axis=1)
            result_dict[PROBABILITY] = prediction.squeeze().tolist()
        result.append(generate_single_csv_line_inference_selection(prediction.tolist()))
    response = COMMA_DELIMITER.join(result)

    if output_content_type == JSON_FORMAT:
        response = json.dumps(result_dict)

    return response, output_content_type