run_grid_hierarchical.m

function run_grid_hierarchical(S, outDir, maxConcurrent, poolSizePerJob)
% Parallelize grid runs across separate MATLAB jobs (best for heavy work).
% Each job:
%   - creates its own run directory under outDir
%   - writes synthetic_data.mat for that (mu,G)
%   - copies model files into runDir/synthetic_sims
%   - calls hierarchical_imr_compare (which uses its own pool)
%
% Usage:
%   run_grid_hierarchical_batch(S);
%   run_grid_hierarchical_batch('synthetic_grid.mat', 'grid_results', 4, 4)
%
% Args:
%   S               : (nMu x nG struct) OR .mat filename containing S
%   outDir          : results root (default 'grid_results')
%   maxConcurrent   : #simultaneous jobs (default: max(1, floor(numcores/2)))
%   poolSizePerJob  : workers inside each job (default: 0 = let code decide)

    % ----- load S if filename provided -----
    if nargin < 1 || isempty(S)
        if isfile('synthetic_grid.mat')
            L = load('synthetic_grid.mat'); S = L.S;
        else
            error('Provide S or synthetic_grid.mat');
        end
    elseif ischar(S) || isstring(S)
        L = load(S); S = L.S;
    end

    if nargin < 2 || isempty(outDir), outDir = 'grid_results'; end
    if ~exist(outDir,'dir'), mkdir(outDir); end

    if nargin < 3 || isempty(maxConcurrent)
        try ncores = feature('numcores'); catch, ncores = 4; end
        maxConcurrent = max(1, floor(ncores/2));   % conservative default
    end
    if nargin < 4 || isempty(poolSizePerJob), poolSizePerJob = 0; end

    baseDir = pwd;               % project root (where your code lives)
    c = parcluster();            % default cluster/profile
    jobs = parallel.Job.empty(0,1);

    [nMu,nG] = size(S);
    fprintf('Submitting %d x %d jobs (max concurrent = %d, pool per job = %d)\n', ...
        nMu, nG, maxConcurrent, poolSizePerJob);

    submitted = 0;
    for i = 1:nMu
        for j = 1:nG
            muVal = S(i,j).mu; GVal = S(i,j).G;
            runDir = fullfile(outDir, sprintf('i%02d_j%02d_mu%.6g_G%.6g', i, j, muVal, GVal));
            outMat = fullfile(runDir, 'hierarchical_results.mat');

            if isfile(outMat)
                fprintf('Skip existing: %s\n', outMat);
                continue;
            end

            if ~exist(runDir,'dir'), mkdir(runDir); end

            % Throttle to maxConcurrent
            while countActive(jobs) >= maxConcurrent
                pause(1); jobs = pruneFinished(jobs);
            end

            % Submit job; each job copies models & writes its own synthetic_data.mat
            args = {S(i,j), runDir, baseDir};
            jh = batch(c, @imr_job_one_cell, 0, args, 'Pool', poolSizePerJob);
            jobs(end+1) = jh; %#ok<AGROW>
            submitted = submitted + 1;
            fprintf('Submitted (%d,%d)\n', i, j);
        end
    end

    % Wait for all remaining jobs
    jobs = pruneFinished(jobs);
    if ~isempty(jobs)
        fprintf('Waiting on %d jobs...\n', numel(jobs));
        wait(jobs);
    end
    fprintf('All done. Submitted %d jobs.\n', submitted);
end

% -------- helpers --------
function n = countActive(jobs)
    if isempty(jobs), n=0; return; end
    st = string({jobs.State});
    n = nnz(st=="queued" | st=="running");
end

function jobs = pruneFinished(jobs)
    if isempty(jobs), return; end
    st = string({jobs.State});
    jobs = jobs(~(st=="finished" | st=="failed"));
end

% -------- job payload (runs in its own MATLAB) --------
function imr_job_one_cell(Cell, runDir, baseDir)
    % Prepare isolated run directory
    if ~exist(runDir,'dir'), mkdir(runDir); end
    % Make a local 'synthetic_sims' so hierarchical_imr_compare finds model .mat files
    srcModels = fullfile(baseDir, 'synthetic_sims');
    dstModels = fullfile(runDir, 'synthetic_sims');
    if ~exist(dstModels,'dir')
        if exist(srcModels,'dir')
            copyfile(srcModels, dstModels);   % small folder; safe & simple
        else
            error('Model folder not found: %s', srcModels);
        end
    end
    % Satisfy expDataPrep addpath('./experiments/') (ok if empty)
    expDir = fullfile(runDir, 'experiments');
    if ~exist(expDir,'dir'), mkdir(expDir); end

    % Write synthetic_data.mat expected by expDataPrep (case 0)
    synthetic_data = Cell.synthetic_data; %#ok<NASGU>
    if isfield(Cell,'Rmax') && isfinite(Cell.Rmax), Rmax = Cell.Rmax; else, Rmax = 100e-6; end %#ok<NASGU>
    if isfield(Cell,'Rmax_std') && isfinite(Cell.Rmax_std), Rmax_std = Cell.Rmax_std; else, Rmax_std = 0.1*Rmax; end %#ok<NASGU>
    if isfield(Cell,'lambda') && isfinite(Cell.lambda), lambda = Cell.lambda; else, lambda = 5; end %#ok<NASGU>
    save(fullfile(runDir,'synthetic_data.mat'), 'synthetic_data','Rmax','Rmax_std','lambda');

    % Run the pipeline inside runDir
    old = pwd; cleaner = onCleanup(@() cd(old));
    cd(runDir);
    addpath(baseDir);                    % to see your functions
    addpath(fullfile(baseDir,'experiments'));  % if needed

    % hierarchical_imr_compare will:
    %   - read local synthetic_data.mat
    %   - read local synthetic_sims/<model>.mat
    %   - save hierarchical_results.mat into runDir
    hierarchical_imr_compare;

    % Add a tiny meta file
    meta = rmfield(Cell, 'synthetic_data'); %#ok<NASGU>
    save(fullfile(runDir, 'meta.mat'), 'meta');
end