% CachedMJSAccess - Decorator for MJSAccess which has a local data cache to
% minimize remote network calls. Data not available in the local cache will
% be fetched using the using MJSAccess call.
% The local cache is used as follows:
%
% Cluster Properties:
%   1. Return any property values which exist in the local cache
%   2. Fetch remaining properties from the job manager
%
% Work Unit Properties:
%   1. Use the cache for any work units in the cache with values for all 
%   required properties
%   2. Use the cache for any properties in the cache with values for all
%   required work units
%   3. Fetch all remaining work unit and property combinations from the job
%   manager
%
%   The above strategy means we make at most 1 remote call for a
%   vectorised get. Values in the cache will be ignored if the remote call
%   involves getting that value.

% Copyright 2020-2021 The MathWorks, Inc.

classdef (Sealed, Hidden) CachedMJSAccess < parallel.internal.cluster.ForwardingMJSAccess
        
    properties (GetAccess = private, SetAccess = private)
        % Local map of cluster Property-Value pairs
        ClusterPropertyCache
             
        % Local map of WorkUnit key to a map of work unit Property-Value pairs
        WorkUnitPropertyCache
    end
    
    methods
        function obj = CachedMJSAccess(access)            
            obj@parallel.internal.cluster.ForwardingMJSAccess(access);
            obj.ClusterPropertyCache = containers.Map('KeyType', 'char', 'ValueType', 'any');
            obj.WorkUnitPropertyCache = containers.Map('KeyType', 'char', 'ValueType', 'any');
        end
        
        function clear(obj)
            obj.ClusterPropertyCache = containers.Map('KeyType', 'char', 'ValueType', 'any');
            obj.WorkUnitPropertyCache = containers.Map('KeyType', 'char', 'ValueType', 'any');
        end

        function clearClusterProperties(obj)
            obj.ClusterPropertyCache = containers.Map('KeyType', 'char', 'ValueType', 'any');
        end
        
        function clearWorkUnit(obj, workUnitID)
            iSafeRemove(obj.WorkUnitPropertyCache, iGetKey(workUnitID));
        end
        
        function cacheClusterProperties(obj, propertyMap)
            % Add cluster properties to local cache. Any property which
            % already existed in the cache is replaced
            obj.ClusterPropertyCache = iCombineMaps(obj.ClusterPropertyCache, propertyMap);
        end
        
        function cacheWorkUnitProperties(obj, workUnitID, propertyMap)
            % Add work unit properties to local cache. If the work unit
            % already exists in the cache, the new property map and
            % existing one will be combined. For any properties which exist in both
            % the old and new map, the value in the new map will be cached.
            workUnitMap = containers.Map(iGetKey(workUnitID), propertyMap);
            obj.WorkUnitPropertyCache = iCombineMapsOfMaps(obj.WorkUnitPropertyCache, workUnitMap);
        end
    end
    
    % Override parallel.internal.cluster.IMJSAccess methods to use local
    % property cache
    methods        
        function values = getClusterProperties(obj, propertyNames, hashedCredentials)
            isCached = isKey(obj.ClusterPropertyCache, propertyNames);
            if any(isCached)
                cachedValues = obj.ClusterPropertyCache.values(propertyNames(isCached));
                cachedValues = cellfun(@(val) {val}, cachedValues, 'UniformOutput', false);
                remoteValues = obj.Access.getClusterProperties(propertyNames(~isCached), hashedCredentials);
                
                % Combine cached and remote values
                values = cell(size(propertyNames));
                values(isCached) = cachedValues;
                values(~isCached) = remoteValues;
            else
                values = obj.Access.getClusterProperties(propertyNames, hashedCredentials);
            end
        end
        
        function values = getTaskProperties(obj, workUnitIDs, propertyNames, hashedCredentials)
            accessFcn = @(wid, props, auth) obj.Access.getTaskProperties(wid, props, auth);
            values = obj.getWorkUnitProperties(accessFcn, workUnitIDs, propertyNames, hashedCredentials);
        end
        
        function values = getJobProperties(obj, workUnitIDs, propertyNames, hashedCredentials)
            accessFcn = @(wid, props, auth) obj.Access.getJobProperties(wid, props, auth);
            values = obj.getWorkUnitProperties(accessFcn, workUnitIDs, propertyNames, hashedCredentials);
        end
    end
    
    methods (Access = private)
        function values = getWorkUnitProperties(obj, accessFcn, workUnitIDs, propertyNames, hashedCredentials)
            workUnitIdentifierKeys = arrayfun(@iGetKey, workUnitIDs, 'UniformOutput', false);
            isWorkUnitCached = isKey(obj.WorkUnitPropertyCache, workUnitIdentifierKeys);
            
            if any(isWorkUnitCached)
                values = cell(size(propertyNames));
                                
                % Some work units are cached.
                cachedMaps = obj.WorkUnitPropertyCache.values(workUnitIdentifierKeys(isWorkUnitCached));
                                               
                % 1. Check which properties are cached for all work units
                % (if any)
                if all(isWorkUnitCached)
                    isPropertyCached = true(size(propertyNames));
                    for idx = 1:numel(cachedMaps)
                        isPropertyCached = isPropertyCached & isKey(cachedMaps{idx}, propertyNames);
                        if ~any(isPropertyCached)
                            % None of the properties are cached for all the
                            % work units.
                            break
                        end
                    end
                else
                    isPropertyCached = false(size(propertyNames));
                end
                cachedProperties = propertyNames(isPropertyCached);
                
                % Collect these cached properties from the cache
                if any(isPropertyCached)
                    cachedPropertiesValues = cellfun(@(map) map.values(cachedProperties), cachedMaps, 'UniformOutput', false);
                    cachedPropertiesValues = vertcat(cachedPropertiesValues{:})';
                    cachedPropertiesValues = cellfun(@iResolveAnyDeferedEvaluations, cachedPropertiesValues, 'UniformOutput', false);
                    values(isPropertyCached) = arrayfun(@(ii) cachedPropertiesValues(ii, :), 1:size(cachedPropertiesValues, 1), 'UniformOutput', false);
                end
                                
                % 2. Find which work units are cached for all properties
                propertiesToGet = propertyNames(~isPropertyCached);
                cachesAllProperties = cellfun(@(map) all(isKey(map, propertiesToGet)), cachedMaps);
                isWorkUnitCached(isWorkUnitCached) = cachesAllProperties;
                workUnitsToGet = workUnitIDs(~isWorkUnitCached);

                % Collect these cached work units values from the cache
                if any(isWorkUnitCached)
                    cachedWorkUnitValues = cellfun(@(map) map.values(propertiesToGet), cachedMaps(isWorkUnitCached), 'UniformOutput', false);
                else
                    cachedWorkUnitValues = {};
                end
                
                % Fetch remote values for remaining property and work
                % unit combinations
                remoteValues = accessFcn(workUnitsToGet, propertiesToGet, hashedCredentials);
                
                % Combine cached work unit values and remote work unit
                % values
                combinedWorkUnitValues = cell(size(propertiesToGet));
                for idx = 1:numel(propertiesToGet)
                    resultForProp = cell(size(workUnitIDs));
                    resultForProp(isWorkUnitCached) = cellfun(@(c) c{idx}, cachedWorkUnitValues, 'UniformOutput', false);
                    resultForProp(~isWorkUnitCached) = remoteValues{idx};
                    combinedWorkUnitValues{idx} = resultForProp;
                end
                
                values(~isPropertyCached) = combinedWorkUnitValues;
                
                return
            end
            
            % No work units are cached
            values = accessFcn(workUnitIDs, propertyNames, hashedCredentials);
        end
    end
end

function key = iGetKey(id)
    key = string(id.long1) + string(id.long2);
end

function combinedMap = iCombineMaps(oldMap, newMap)
% Combine oldMap and new Map such that:
% - combinedMap contains the union of oldMap and new Map
% - combinedMap contains entries of newMap for the intersect of oldMap and
% newMap keys

newKeys = newMap.keys;
oldKeys = oldMap.keys;

[~, newIndex, oldIndex] = union(newKeys, oldKeys);
keys = [newKeys(newIndex) oldKeys(oldIndex)];
newValues = newMap.values(newKeys(newIndex));
oldValues = oldMap.values(oldKeys(oldIndex));
combinedMap = containers.Map(keys, [newValues, oldValues], 'UniformValues', false);
end

function combinedMap = iCombineMapsOfMaps(oldMap, newMap)
% Combine oldMap and new Map such that:
% - combinedMap contains the union of oldMap and newMap
% - entries in both oldMap and newMap are combined using iCombineMaps

newKeys = newMap.keys;
oldKeys = oldMap.keys;

% Collect entries unique to each map
[~, newIndex, oldIndex] = setxor(newKeys, oldKeys);
xorKeys = [newKeys(newIndex) oldKeys(oldIndex)];
xorValues = [newMap.values(newKeys(newIndex)), oldMap.values(oldKeys(oldIndex))];

% Combine the intersection
intKeys = intersect(newKeys, oldKeys);
intnewValues = newMap.values(intKeys);
intoldValues = oldMap.values(intKeys);
intValues = cellfun(@iCombineMaps, intoldValues, intnewValues, 'UniformOutput', false);

keys = [xorKeys, intKeys];
values = [xorValues, intValues];

combinedMap = containers.Map(keys, values, 'UniformValues', false);
end

% Helper method to avoid warning when removing keys which don't exist
function keysToRemove = iSafeRemove(map, keysToRemove)
    allKeys = keys(map);
    keysToRemove = intersect(keysToRemove, allKeys);
    remove(map, keysToRemove);
end

function value = iResolveAnyDeferedEvaluations(value)
    if isa(value, 'parallel.internal.cluster.DeferredEvaluation')
        value = value.get();
    end
end
