classdef Scheduler < internal.stats.bigdata.JobContainer
    %Scheduler class.
    %  Scheduler implements a gather-free scheduler. It contains tall array jobs
    %  and runs them together by synchronizing data passes from different
    %  jobs. To use a tall method it must have scheduler support built-in,
    %  e.g. tall.fitclinear and tall.fitckernel. Scheduler has the capacity
    %  of adding more jobs on-the-fly.
    %
    %  Scheduler properties:
    %      MarkForReuseFlag - Dictates cache policy
    %      NumJobs  - Number of contained jobs
    %      State    - Current state of the contained jobs ('running','finished','unavailable')
    %
    %  Scheduler methods:
    %      Scheduler           - Constructor.
    %      addJob              - Add a job to the scheduler.
    %      fetchOutput         - Get the output of specified job
    %      fetchNext           - Get output of next available job.    
    %      run                 - Run all jobs until all are finished.    
   
    %   Copyright 2018 The MathWorks, Inc.
    
    %  Scheduler hidden properties (advanced use):
    %      Jobs     - Jobs to be run or currently running.    
    %
    %  Scheduler hidden methods (advanced use):
    %      addCachingJob       - Add a caching job to the scheduler.    
    %      runStep             - Run all jobs until the next gather.
    %      runUntilOneJobReady - Run all jobs until one is finished. 

    
    properties(SetAccess=private,Hidden)
        Jobs = {}           % cell array of jobs
    end

    properties(Dependent=true,SetAccess=private)
        NumJobs             % number of jobs
        State               % state of contained job ('running','finished','unavailable')
    end
    
    properties(Access=private)
        VarsPerJob = [];    % used internally to figure out which variables are pushed to each job workspace   
    end
    
    methods
        function n = get.NumJobs(s)
            n = length(s.Jobs);
        end
        function st = get.State(s)
            st = cell(1,s.NumJobs);
            for i = 1:s.NumJobs
                st{i} = s.Jobs{i}.State;
            end
        end
    end
    
    methods (Hidden)
        function registerChildJob(s,j)
            i = s.NumJobs+1;
            s.VarsPerJob(i) = 0;
            s.Jobs{i} = j; % Connect to parent scheduler
        end
        
        function runStep(s)
        % runStep gathers and runs until next gather needed.
        %
        % runStep(Sch) triggers all pending gathers from all jobs by optimizing
        % big data read operations and pushes back the value of the
        % gathered variables back to the workspace of each corresponding job.
        % Then, runStep continues running each individual job until a new
        % gather is required. Some jobs may not require more gathers, for
        % these, their state changes to 'finished' and the output will be
        % ready for pulling it out.
            vars = {};
            for i = 1:s.NumJobs
                tmp = getVarsToGather(s.Jobs{i});
                s.VarsPerJob(i) = numel(tmp);
                vars = [vars,tmp];
            end
            if ~isempty(vars)
                [vars{:}] = gather(vars{:});
                j = 0;
                for i = 1:s.NumJobs
                    if s.VarsPerJob(i)>0
                        pushGatheredVars(s.Jobs{i},vars(j+(1:s.VarsPerJob(i))))
                        j = j+s.VarsPerJob(i);
                    end
                    runUntilGatherNeeded(s.Jobs{i})
                end
            end
        end
        
        function idx = runUntilOneJobReady(s)
        % runUntilOneJobReady run all jobs until one is finished.
        %
        % Index = runUntilOneJobReady(Sch) run all jobs synchronizing gathers
        % until at least one is finished and waiting to have its results
        % pulled out. runUntilOneJobReady returns the Index to finished job.
        % It is possible that there are more than one finished jobs,
        % runUntilOneJobReady returns only one index; in this case by calling
        % runUntilOneJobReady subsequently, nothing is actually run but 
        % you get an index to another finished job. idx is 0 if all jobs are
        % inactive. 
            if all(strcmp(s.State,'unavailable'))
                idx = 0;
                return
            end
            while ~any(strcmp(s.State,'finished'))
               runStep(s)
            end        
            idx = find(strcmp(s.State,'finished'),1);
        end  
        
        function addCachingJob(s,fh,varargin)
        % addCachingJob adds a Job to the scheduler.
        %
        % Sch = addCachingJob(Sch,FunctionHandle,Input1,Input2,...) adds a
        % job FunctionHandle to the scheduler Sch with inputs Input1, Input2,
        % etc. % Job runs until a gather is needed (or runs until the output is
        % ready in the rare case that a particular job did not need any
        % gather to complete). Under normal behavior (when using addJob) an
        % internal flag is set indicating that tall data was already flagged with
        % markForReuse, so tall method implementations can check this flag
        % to avoid double caching. When a job is added using addCachingJob
        % this mechanism is disabled, forcing a tall method implementation
        % to cache the data again.
            pf = s.MarkForReuseFlag;
            s.MarkForReuseFlag = true;
            addJob(s,fh,varargin{:})
            s.MarkForReuseFlag = pf;
        end
    end
    
    methods
        function s = Scheduler(varargin)
        % Scheduler Constructor. 
        %
        % Sch = Scheduler() returns an empty scheduler. 
        %
        % Sch = Scheduler(FunctionHandle,Input1,Input2,...) returns a
        % scheduler with one job for the function FunctionHandle with
        % inputs Input1, Input2, etc. FunctionHandle is a function with
        % gather-free scheduler support.  Job runs until a gather is needed
        % (or runs until the output is ready in the rare case that a
        % particular job did not need any gather to complete).
            s.MarkForReuseFlag = false;
            if nargin>0
                varargin{1}(varargin{2:end},s);
            end
        end
        
        function addJob(s,fh,varargin)
        % addJob adds a Job to the scheduler.
        %
        % Sch = addJob(Sch,FunctionHandle,Input1,Input2,...) adds a job
        % FunctionHandle to the scheduler Sch with inputs Input1, Input2, etc. 
        % FunctionHandle is a function with gather-free scheduler support.
        % Job runs until a gather is needed (or runs until the output is
        % ready in the rare case that a particular job did not need any
        % gather to complete).
             fh(varargin{:},s);
        end
        
        function run(s)
        % run run all jobs until all are finished.
        %
        % run(Sch) run all jobs synchronizing gathers until all are
        % finished and waiting to have its results pulled out. 
            while any(strcmp(s.State,'running'))
               runStep(s)
            end
        end

        function varargout = fetchNext(s)
        % fetchNext fetches the output of the next available finished job.
        %
        % [i,Out1,Out2,...] = fetchNext(Sch) gets the index and the outputs
        % of the next finished job. After pulling out the outputs the job
        % is no longer active, job's state will change to 'unavailable'.
        % fetchOutput waits until at least one job finishes. In case two or
        % more job have finished, fetchNext will select the first one.
        % Subsequent calls to fetchNext should be used to get outputs of
        % the other finished jobs. If all jobs have finished and all
        % outputs have been fetched, fetchNext returns 0 in the first
        % output argument.
            varargout = cell(1,nargout);
            varargout{1} = runUntilOneJobReady(s);
            if varargout{1}>0
                [varargout{2:end}] = fetchOutput(s,varargout{1});
            end 
        
        end
                    
        function varargout = fetchOutput(s,idx)
        % fetchOutput fetches the output of a specified finished job.
        %
        % [Out1,Out2,...] = fetchOutput(Sch,i) gets the outputs of the i-th
        % job. After pulling out the outputs the job is no longer
        % active, job's state will change to 'unavailable'. fetchOutput
        % waits until the i-th job finishes, so you should only fetch
        % outputs of jobs with 'finished' state. 
            if strcmp(s.State{idx},'unavailable')
                [varargout{1:nargout}] = getOutput(s.Jobs{idx});
            end
            while ~strcmp(s.State{idx},'finished')
                runStep(s)
            end
            [varargout{1:nargout}] = getOutput(s.Jobs{idx});
        end
        
    end
    
    methods(Static)
        function varargout = createAndRun(fh,varargin)
        % createAndRun runs a function with gather-free scheduler support.
        %
        % [Out1,Out2,...] = createAndRun(FunctionHandle,Input1,Input2,...)
        % runs FunctionHandle with input arguments Input1, Input2, etc.
        % When the job is finished output are fetched and assigned to Out1,
        % Out2, etc. When createAndRun is utilized, it can not be combined
        % with other jobs, as a new scheduler will be created and
        % finalized; createAndRun is a helper function that allows to call 
        % functions with gather-free scheduler support directly without
        % specifying any scheduler.
            s = internal.stats.bigdata.Scheduler;
            s.MarkForReuseFlag = true; % Scheduler created with only one job, most likely when
                                       % user called directly a tall method, therefore allow
                                       % the tall method to do further caching
            fh(varargin{:},s); % call the tall method with the scheduler as in input
            run(s)             % intruct the scheduler to run all thunks
            [varargout{1:nargout}] = getOutput(s.Jobs{1}); % fetch output arguments
        end
    end
    
end
