classdef Job  < internal.stats.bigdata.JobContainer
    % Job is a class that represents a tall method or a function that
    % has gather-free scheduler support. Job can be contained by
    % another Job or by the Scheduler. 
    %
    % The public interface of Job must be used from inside the
    % implementation of a tall method with gather-free scheduler support.
    %
    %  Job properties:
    %      MarkForReuseFlag      - Dictates cache policy
    %
    %  Job methods:
    %      Job                   - Constructor.
    %      addJob                - Add a nested job.
    %      deferGather           - Gather variables after thunk.
    %      deferOutputAssignment - Assign variables after thunk.
    %      setOutput             - Set output arguments.
    %
    % A tall method with gather-free support should be implemented
    % with nested functions following this scheleton:
    %
    %     function varargout = foo(X,...,varargin)
    %
    %         [iContainer,varargin] = internal.stats.bigdata.Scheduler.catchJobContainerFromInputArgs(varargin);
    %         if isempty(iContainer)
    %             varargout{:} = internal.stats.bigdata.Scheduler.createAndRun(@foo,X,...,varargin{:});
    %             return
    %         end
    %
    %         % Prelude code
    %         ...
    %         % Initialization of nested variables
    %         ...
    %
    %         iJob = Job(iContainer,@thunk_1,@pushVar)
    %
    %         function iNext = thunk_1()
    %            ...   
    %         end
    %         function iNext = thunk_2()
    %            ...
    %         end
    %
    %         ...
    %
    %         function iNext = thunk_last()
    %             ...
    %             iNext = [];
    %             setOutput(iJob,varargot{:})
    %         end
    %
    %         function pushVar(iJob_vn,iJob_v)
    %             eval([iJob_vn '=iJob_v;']);
    %         end
    %
    %     end   
    
    %   Copyright 2018 The MathWorks, Inc.
    
    properties (Hidden)
        
        NestedJob = {};         % contains nested jobs if any
        State     = '';         % current state of the job ('running','finished','unavailable')
        NextThunk = [];
        ImportDataHelper = [];
        OutputArguments = [];
        DeferredGatherVariableValues = [];
        DeferredGatherVariableNames  = cell(1,0);
        DeferredNestedVariableNames  = cell(1,0);
    end
    
    methods
        function j = Job(firstThunk,importDataHelper,Container)
        % Job constructor.
        %
        %   Job(FirstThunk,ImportDataHelper,Container) creates a new Job to
        %   represent the context of the actual workspace. The Job is
        %   linked to the Container, which can be another Job or a
        %   Scheduler. After this, the scheduler will take control of the
        %   code execution, starting from firsThunk. importDataHelper is a
        %   function handle to a nested function that helps to copy data
        %   back to the current workspace. 
        
            registerChildJob(Container,j);   % Connect to parent (may be a Job or the Scheduler)
            j.MarkForReuseFlag = Container.MarkForReuseFlag;
            importDataHelper('iJob',j);      % Place a handle in the function workspace
            j.NextThunk = firstThunk;
            j.ImportDataHelper = importDataHelper;
            j.State = 'running';
            runUntilGatherNeeded(j);
        end  
        function addJob(j,fh,varargin)
        % addJob adds a nested Job.
        %
        % addJob(Job,FunctionHandle,Input1,Input2,...) adds a nested job
        % FunctionHandle to the current Job with inputs Input1, Input2, etc. 
        % FunctionHandle is another function with gather-free scheduler support.
             fh(varargin{:},j);
        end               
        function deferGather(j,varNames,varValues)
        % deferGather marks a variable to be gathered after finishing a thunk.
        %         
        %   deferGather(Job,VariableName,VariableValue) indicates to the
        %   Scheduler that the tall variable VariableValue must be gathered
        %   after the thunk finishes and placed back into the context of this
        %   Job by assigning its gathered value to VariableName. VariableName
        %   and VariableValues may also be lists passed as a cell array, both
        %   must have the same length. deferGather must only be used inside a 
        %   thunk and the gathered variables cannot be used in the remmanining 
        %   of the current thunk.
        
            % For convinience, when only a single variable needs to be
            % gather we allow to pass it without being a cell:
            if ~iscell(varNames)
                varNames = {varNames};
                varValues = {varValues};
            end
            % Even if marked for gathering some variables may not be tall,
            % if so then push them in to the function context right away,
            % so this block will not be holding the scheduler.
            for k = numel(varValues):-1:1
                if ~istall(varValues{k})
                    j.ImportDataHelper(varNames{k},varValues{k})
                    varNames(k)  = [];
                    varValues(k) = [];
                end
            end
            j.DeferredGatherVariableNames = varNames;
            j.DeferredGatherVariableValues = varValues;
        end
        function deferOutputAssignment(j,varNames)
        % deferOutputAssignment marks a variable to be assigned after finishing a thunk.
        %         
        %   deferGather(Job,VariableName) indicates to the Scheduler that a
        %   variable VariableName must be copied into the current workspace
        %   from a finished nested job. VariableName may also be a list
        %   passed as a cell array. deferOutputAssignment must only be used
        %   inside a thunk and the copied variables cannot be used in the
        %   remmanining of the current thunk.
        
            % For convinience, when only a single variable needs to be
            % assigned we allow to pass it without being a cell:
            if ~iscell(varNames)
                varNames = {varNames};
            end
            j.DeferredNestedVariableNames = varNames;
        end       
        function setOutput(j,varValues)
        % setOutput set output arguments.
        %
        %    setOutput(Job,VariableValues) indicates which variables of the
        %    current workspace should be kept as output arguments to be
        %    fetched by a container (either another Job or the Scheduler). 
        %    VariableValues may also be a list passed as a cell array.
            if ~iscell(varValues)
                varValues = {varValues};
            end
            j.OutputArguments = varValues;
        end
 
        
    end
    
    methods (Hidden)
        function registerChildJob(iJob,j)
            iJob.NestedJob{end+1} = j; % Connect to parent job
        end

        function tf = moreThunks(j)
           tf = ~isempty(j.NextThunk);
        end
        
        function tf = hasNestedJob(j)
            tf = ~isempty(j.NestedJob);
        end
        
        function tf = waitingForGather(j)
            if hasNestedJob(j)
                tf = any(cellfun(@(x) waitingForGather(x),j.NestedJob));
            else
                tf = ~isempty(j.DeferredGatherVariableValues);
            end
        end

        
        function runUntilGatherNeeded(j)
            
            while moreThunks(j)
                
                if  hasNestedJob(j)
                    % First continue with any nested job
                    cellfun(@(x) runUntilGatherNeeded(x),j.NestedJob)
                    % by now, the nested job is either waiting for a gather
                    % or it is finalized and outputs are ready to be
                    % collected and bring them into the context of the
                    % parent job
                    n = numel(j.NestedJob);
                    for k=1:n
                        if waitingForGather(j.NestedJob{k})
                            return % there is something to gather one level before
                        end
                    end
                    nOutArgs = cellfun(@(x) numel(x.OutputArguments),j.NestedJob);
                    vtmp = cell(n,max(nOutArgs));
                    for k=1:n
                        [vtmp{k,1:nOutArgs(k)}] = getOutput(j.NestedJob{k});
                    end
                    
                    for i = 1:numel(j.DeferredNestedVariableNames)
                        if n==1
                            j.ImportDataHelper(j.DeferredNestedVariableNames{i},vtmp{1,i})
                        else
                            j.ImportDataHelper(j.DeferredNestedVariableNames{i},vtmp(:,i))
                        end
                    end
                    j.NestedJob = {};
                    j.DeferredNestedVariableNames = cell(1,0);
                else
                    % Evaluate the Thunk
                    out = feval(j.NextThunk); 
                    if iscell(out) 
                        j.OutputArguments = out;
                        out = [];
                    end
                    if isempty(out)
                        j.State = 'finished';
                        j.NextThunk = [];
                        j.ImportDataHelper = [];
                    else
                        j.NextThunk = out;
                    end
                end

                if waitingForGather(j)
                    return
                end
            end
            
        end
        
        function vars = getVarsToGather(j)
            if hasNestedJob(j)
                vars = cellfun(@(x) getVarsToGather(x),j.NestedJob,'Uniform',false);
                vars = [vars{:}];
            else
                vars = j.DeferredGatherVariableValues;
            end
        end
        
        function pushGatheredVars(j,vars)
            if ~waitingForGather(j)
                return
            end
            if hasNestedJob(j)
                n = numel(j.NestedJob);
                sel = repelem(1:n,cellfun(@(x) numelDeferredGatherVariableNames(x),j.NestedJob));
                for i = 1:n
                    pushGatheredVars(j.NestedJob{i},vars(sel==i));
                end
            else
                n = numel(j.DeferredGatherVariableNames);
                for i = 1:n
                    j.ImportDataHelper(j.DeferredGatherVariableNames{i}, vars{i});
                end
                j.DeferredGatherVariableNames  = cell(1,0);
                j.DeferredGatherVariableValues = cell(1,0);                
            end
        end
        
        function d = numelDeferredGatherVariableNames(j)
            if hasNestedJob(j)
                d = sum(cellfun(@(x) numelDeferredGatherVariableNames(x),j.NestedJob));
            else
                d = numel(j.DeferredGatherVariableNames);
            end
        end
        
        function [varargout] = getOutput(j)
            if strcmp(j.State,'unavailable')
                error('Job is no longer in context')
            end
            if moreThunks(j) || waitingForGather(j)
                error('Job has not finished')
            end
            if nargout>numel(j.OutputArguments)
                error('Too many output arguments.')
            end
            varargout = cell(nargout,1);
            for i = 1:nargout
                varargout{i} = j.OutputArguments{i};
            end
            j.OutputArguments = [];
            j.State = 'unavailable';
        end
        
    end
end