from typing import OrderedDict
import json


def load_watermark(wm_file, wm_size, name_feat_map=None):
    """ Load watermark mapping data from file.

    :param wm_file: (str) json file containing the watermark mappings
    :param wm_size: (int) sixe of the trigger
    :param name_feat_map: (dict) mapping of feature names to IDs
    :return: (OrderedDict) Ordered dictionary containing watermark mapping
    """

    wm = OrderedDict()
    loaded_json = json.load(open(wm_file, 'r'))
    ordering = loaded_json['order']
    mapping = loaded_json['map']

    i = 0
    for ind in sorted(ordering.keys()):
        feat = ordering[ind]

        if name_feat_map is not None:
            key = name_feat_map[feat]
        else:
            key = feat

        wm[key] = mapping[feat]

        i += 1
        if i == wm_size:
            break

    return wm