import mxnet as mx
import numpy as np
import cv2
import logging
from collections import namedtuple
Batch = namedtuple('Batch', ['data'])

DEFAULT_INPUT_SHAPE = 512

def get_ctx():
    """
    Automatically choose the device (CPU or GPU) for running inference
    :return: A list of GPUs available. If no GPU is available, return a list with just the CPU device.
    """
    try:
        gpus = mx.test_utils.list_gpus()
        if len(gpus) > 0:
            ctx = [mx.gpu(gpu) for gpu in gpus]
        else:
            ctx = [mx.cpu()]
    except:
        ctx = [mx.cpu()]
    return ctx


class MLModel(object):
    """
    Loads the pre-trained model which can be found in /ml/od when running on greengrass core or
    from a different path for testing locally.
    """
    def __init__(self, param_path, label_names=[], input_shapes=[('data', (1, 3, DEFAULT_INPUT_SHAPE, DEFAULT_INPUT_SHAPE))]):

        context = get_ctx()[0]
        # Load the network parameters from default epoch 0
        logging.info('Load network parameters from default epoch 0 with prefix: {}'.format(param_path))
        sym, arg_params, aux_params = mx.model.load_checkpoint(param_path, 0)

        # Load the network into an MXNet module and bind the corresponding parameters
        logging.info('Loading network into mxnet module and binding corresponding parameters: {}'.format(arg_params))
        self.mod = mx.mod.Module(symbol=sym, label_names=label_names, context=context)
        self.mod.bind(for_training=False, data_shapes=input_shapes)
        self.mod.set_params(arg_params, aux_params)

    """
    Takes in an image, reshapes it, and runs it through the loaded MXNet graph for inference returning the top label from the softmax
    """
    def predict_from_file(self, filepath, reshape=(DEFAULT_INPUT_SHAPE, DEFAULT_INPUT_SHAPE)):
        # Switch RGB to BGR format (which ImageNet networks take)
        img = cv2.cvtColor(cv2.imread(filepath), cv2.COLOR_BGR2RGB)
        if img is None:
            return []

        # Resize image to fit network input
        img = cv2.resize(img, reshape)
        img = np.swapaxes(img, 0, 2)
        img = np.swapaxes(img, 1, 2)
        img = img[np.newaxis, :]

        self.mod.forward(Batch([mx.nd.array(img)]))
        prob = self.mod.get_outputs()[0].asnumpy()
        prob = np.squeeze(prob)

        # Grab top result, convert to python list of lists and return
        results = [prob[0].tolist()]
        return results