callMCMC.m

close all; clear; clc;

% ===================== settings =====================
N              = 10000;                   % MH steps per (model, beta)
betaVals       = 0.1:0.1:1.0;             % must match how modelProb was computed
modelProbsFile = 'UM1_Prob.mat';         % file with *Prob vars (one per model)
samplesFile    = 'UM1_Samples.mat';      % where to save MH samples
posteriorFile  = 'UM1_Posterior.mat';    % where to save evidence/posterior

% Parameter boxes per model (rows: [lower; upper], columns: dims)
X = {[-4;0]; ...                     % newtonian (1D)
     [2;5]; ...                      % NH (1D)
     [-4, 2;   0, 5]; ...            % NHKV (2D)
     [ 2,-3;   5, 1]; ...            % qNH (2D)
     [-4,-7;   0,-3]; ...            % linmax (2D)
     [-4, 2,-3; 0, 5, 1]; ...        % qKV (3D)
     [-4, 2,-7; 0, 5,-3]};           % SLS (3D)

modelLabels = {'newtonian','NH','NHKV','qNH','linmax','qKV','SLS'};

% ===================== load probs =====================
m = load(modelProbsFile);

% Stack into nModels x nBeta struct array with .prob ND cell grids
probs = [m.newtonianProb; ...
         m.NHProb; ...
         m.NHKVProb; ...
         m.qNHProb; ...
         m.linmaxProb; ...
         m.qKVProb; ...
         m.SLSProb];

nModels = size(probs,1);
nBeta   = size(probs,2);
if nBeta ~= numel(betaVals)
    warning('betaVals length (%d) != probs second dim (%d). Adjusting betaVals.', numel(betaVals), nBeta);
    betaVals = 1:nBeta;  % fallback for x-axis labeling
end

% ===================== pool & attach =====================
pool = gcp('nocreate');
if isempty(pool)
    pool = parpool('local');  % or parpool(32) etc.
end
% Ensure workers can see mh.m (uncomment if needed)
% addAttachedFiles(pool, {'mh.m'});

% Optional reproducible randomness (add a seed param to mh if you want)
% baseSeed = 12345;

% ===================== submit jobs =====================
nJobs   = nModels * nBeta;
futures = parallel.FevalFuture.empty(nJobs, 0);
results = cell(nJobs, 1);

job = 0;
for i = 1:nModels
    for j = 1:nBeta
        job = job + 1;
        probGrid = probs(i,j).prob;   % ND cell grid (>= realmin already in your pipeline)
        box      = X{i};               % 2 x dim

        % If you add a seed to mh: futures(job) = parfeval(@mh, 1, N, probGrid, box, baseSeed+job);
        futures(job) = parfeval(@mh, 1, N, probGrid, box);
    end
end

% ===================== collect =====================
completed = 0;
while completed < nJobs
    [doneIdx, out] = fetchNext(futures);   % blocks until a job finishes
    results{doneIdx} = out;                % out is samples [N x (dim+1)]
    completed = completed + 1;
    fprintf('Completed %d/%d\n', completed, nJobs);
end

% Reshape to (nModels x nBeta)
samples = reshape(results, [nModels, nBeta]);

% Save chains
save(samplesFile, 'samples', '-v7.3');

% ===================== evidence & posterior =====================
% Evidence(i,j) = naive MC = mean of likelihoods in the chain for model i, beta j
evidence = zeros(nModels, nBeta);
for i = 1:nModels
    for j = 1:nBeta
        Sij = samples{i,j};
        if isempty(Sij)
            evidence(i,j) = realmin;
        else
            lik = Sij(:, end);                 % last col = likelihood
            lik(~isfinite(lik) | lik <= 0) = realmin;
            evidence(i,j) = mean(lik);
        end
    end
end

% Normalize across models for each beta (columns): posterior sums to 1 per beta
posterior = zeros(nModels, nBeta);
for j = 1:nBeta
    s = sum(evidence(:, j));
    if s<=0 || ~isfinite(s)
        posterior(:, j) = 1 / nModels;
    else
        posterior(:, j) = evidence(:, j) ./ s;
    end
end

save(posteriorFile, 'evidence', 'posterior', 'betaVals', 'modelLabels');

%% ===================== plot =====================
figure; hold on;
for i = 1:nModels
    plot(betaVals, posterior(i, :), '-o', 'DisplayName', modelLabels{i});
end
xlim([min(betaVals) max(betaVals)]); ylim([0 1]);
xlabel('\beta'); ylabel('Posterior probability');
title('Model posterior vs. noise scale \beta (MH over spline, parfeval)');
legend('Location','best'); grid on;