# *******************************************************
#   Mosel Example Problems
#   ======================
#
#   file mosel_numpy.py
#   ```````````````````
#   Helper module for converting dictionaries from Mosel
#   to NumPy arrays and vice versa.
#
#   (c) 2018 Fair Isaac Corporation
#       author: J. Müller, April 2018
# *******************************************************
import numpy as np


def dict_to_array_2d(dictionary, indices1, indices2):
    """ Transform a dictionary from Mosel to a two dimensional NumPy array.

    Note that sets are not ordered. For that reason the array rows and columns
    are not ordered if the indices are sets.
    Use ordered lists as indices to order the rows and columns of the array.

    Parameters
    ----------
    dictionary : dict
        Two dimensional dictionary.
    indices1 : iterable
        List or set of all indices in dimension one.
    indices2 : iterable
        List or set of all indices in dimension two.

    Returns
    -------
    np.array
        Two dimensional NumPy array.
    """
    return np.array([[dictionary[key1, key2] for key2 in indices2] for key1 in indices1])


def array_to_dict_2d(np_array, indices1, indices2):
    """ Transform a two dimensional NumPy array to a dictionary for Mosel.

    Note that the order of the index elements is important. For that reason,
    it is necessary to use the same index sets or lists as in dict_to_array_2d.

    Parameters
    ----------
    np_array : np.array
        Two dimensional array.
    indices1 : iterable
        List or set of all indices in dimension one.
    indices2 : iterable
        List or set of all indices in dimension two.

    Returns
    -------
    dict
        Two dimensional dictionary for Mosel.
    """
    return {(key1, key2): np_array[id1, id2]
            for id2, key2 in enumerate(indices2)
            for id1, key1 in enumerate(indices1)}


def print_dict_as_table(dictionary, indices1, indices2):
    """ Print a two dimensional dictionary as table. """
    width = 6
    print(''.rjust(width), end=' ')

    for j in indices2:
        print(repr(j).rjust(width), end=' ')

    print()

    for i in indices1:
        print(repr(i).rjust(width), end=' ')

        for j in indices2:
            print((("%" + str(width) + ".2f") % dictionary[i, j]), end=' ')

        print()