modelProbComp_wrapper.m

function nll = modelProbComp_wrapper(xlog_tblrow, model, paramSpec, dataPkg, opts)
% Wrapper used by bayesopt. Runs IMR, computes β-marginalized log-like with
% heteroscedastic variance, adds log prior on theta, returns -log10 posterior.

% xlog_tblrow is a 1xP numeric row (log10) from bayesopt table
xlog = xlog_tblrow(:).';
theta = paramTransform(model, xlog, paramSpec);

% Unpack dataPkg (align with your solver + data)
radial     = dataPkg.radial;
stress     = dataPkg.stress;
bubtherm   = dataPkg.bubtherm;
medtherm   = dataPkg.medtherm;
masstrans  = dataPkg.masstrans;
vapor      = dataPkg.vapor;
tvector    = dataPkg.tvector;
R0         = dataPkg.R0;
Req        = dataPkg.Req;
Nt         = dataPkg.Nt;
Mt         = dataPkg.Mt;
Rexps      = dataPkg.Rmatrix;      % T x Ntrials
Rdotexps   = dataPkg.Rdotmatrix;   % T x Ntrials
tmatrix    = dataPkg.tmatrix;      % T x Ntrials
maxTime    = dataPkg.maxTime;
% Sigmas: pass either explicit columns or stats-derived
if isfield(dataPkg,'Rsigma'),      Rsigma = dataPkg.Rsigma;
else,                               Rsigma = get_sigma_column(dataPkg.R_stats, Rexps); end
if isfield(dataPkg,'Rdotsigma'), Rdotsigma = dataPkg.Rdotsigma;
else,                               Rdotsigma = get_sigma_column(dataPkg.Rdot_stats, Rdotexps); end

% Optional masks/ranges (use if provided)
if isfield(dataPkg,'indicator'),         indicator = dataPkg.indicator; else, indicator = true(size(tmatrix)); end
if isfield(dataPkg,'highStrainRanges'),  highStrainRanges = dataPkg.highStrainRanges; else, highStrainRanges = []; end
if isfield(dataPkg,'overlapRanges'),     overlapRanges = dataPkg.overlapRanges; else, overlapRanges = []; end%#ok<NASGU>
% per-trial nondim mapping (if provided)
if isfield(dataPkg,'s_vec'), s_vec = dataPkg.s_vec; else, s_vec = ones(1,size(tmatrix,2)); end
if isfield(dataPkg,'tau_vec'), tau_vec = dataPkg.tau_vec; else, tau_vec = ones(1,size(tmatrix,2)); end

try
    % --- Run solver (on the same tvector grid) ---
    [t, R, U, p] = f_imr_fd( ...
        'radial', radial, 'stress', stress, ...
        'bubtherm', bubtherm, 'medtherm', medtherm, ...
        'masstrans', masstrans, 'vapor', vapor, ...
        'tvector', tvector, 'r0', R0, 'req', Req, ...
        'mu', thetaNamed('mu',theta,paramSpec), ...
        'g',  thetaNamed('G',theta,paramSpec), ...
        'lambda1', thetaNamed('lambda1',theta,paramSpec), ...
        'lambda2', 0, ...
        'alphax', thetaNamed('alpha',theta,paramSpec), ...
        'collapse', 1, 'nt', Nt, 'mt', Mt);

    Sim = [t(:), R(:), U(:), p(:)];
catch
    % Failed sim => huge NLL
    nll = 1e6; return;
end

% --- β-marginalized log-likelihood with heteroscedastic variance ---
loglikeMarg = loglike_theta_betaMarg(Sim, Rexps, Rdotexps, Rsigma, Rdotsigma, tmatrix, maxTime, ...
                                     indicator, highStrainRanges, s_vec, tau_vec, opts);

% --- log prior on theta (optional, set in paramSpec.priors) ---
logprior = logPrior_theta(theta, paramSpec);

% Posterior objective
logpost = loglikeMarg + logprior;
% Return -log10 posterior (stable)
nll = -logpost / log(10);
% clamp if needed
if ~isfinite(nll), nll = 1e6; end
end

%% ======= Likelihood pieces =======
function loglikeMarg = loglike_theta_betaMarg(Sim, Rexps, Rdotexps, Rsigma, Rdotsigma, tmatrix, maxTime, ...
                                              indicator, highStrainRanges, s_vec, tau_vec, opts)
% Computes log p(D | theta) with β marginalized over opts.betaGrid & prior opts.logp_beta
% Includes heteroscedastic variance shaping based on strain-rate.

betaGrid  = opts.betaGrid;
logp_beta = opts.logp_beta;
useHetero = opts.useHetero;
srThr     = opts.srThr;
sharp     = opts.sharp;
minAct    = opts.minAct;

[T, Ntr] = size(Rexps);
if size(Sim,2) < 3
    loglikeMarg = -inf; return;
end

% If Sim's time grid equals tmatrix(:,j), we can index directly; else interpolate.
t_sim = Sim(:,1); R_sim = Sim(:,2); V_sim = Sim(:,3);

% Precompute SR on experiments for hetero shaping
SR_full = abs(-2 .* Rdotexps ./ max(Rexps, 1e-12));

nBeta = numel(betaGrid);
ll_theta_beta = zeros(1, nBeta);

for j = 1:Ntr
    valid = isfinite(tmatrix(:,j)) & isfinite(Rexps(:,j)) & isfinite(Rdotexps(:,j)) ...
          & (tmatrix(:,j) <= maxTime) & logical(indicator(:,j));
    if ~any(valid), continue; end

    tj = tmatrix(valid,j);

    % If your solver already returned on tvector==tj, skip interp
    if numel(tj) <= size(Sim,1) && isequal(tj, t_sim(1:numel(tj)))
        Rs = R_sim(1:numel(tj));
        Vs = V_sim(1:numel(tj));
    else
        % Map SIM mean frame -> trial frame if s_vec/tau_vec provided externally
        s   = s_vec(j); tau = tau_vec(j);
        t_map    = t_sim ./ tau;
        R_map    = R_sim ./ s;
        V_map    = V_sim .* (tau ./ s);
        Rs = interp1(t_map, R_map, tj, 'linear', 'extrap');
        Vs = interp1(t_map, V_map, tj, 'linear', 'extrap');
    end

    Re = Rexps(valid,j);
    Ve = Rdotexps(valid,j);

    sigR  = max(Rsigma(valid),    1e-12);
    sigV  = max(Rdotsigma(valid), 1e-12);

    % Heteroscedastic variance scaling
    if useHetero
        sr = SR_full(valid,j);
        activity = 1 ./ (1 + exp(-sharp * (sr - srThr))); % (0,1)
        activity = max(activity, minAct);
        vscale   = 1 ./ activity;
    else
        vscale = 1;
    end
    vR0 = (sigR.^2)  .* vscale;
    vV0 = (sigV.^2)  .* vscale;

    rR = (Re - Rs);
    rV = (Ve - Vs);

    % Per-β log-like for this trial
    for b = 1:nBeta
        beta = betaGrid(b);
        vR_b = (beta^2) * vR0;
        vV_b = (beta^2) * vV0;

        ll_R = -0.5*( nansum((rR.^2)./vR_b) + nansum(log(2*pi*vR_b)) );
        ll_V = -0.5*( nansum((rV.^2)./vV_b) + nansum(log(2*pi*vV_b)) );
        ll_theta_beta(b) = ll_theta_beta(b) + (ll_R + ll_V);
    end
end

% Marginalize β
if all(~isfinite(ll_theta_beta))
    loglikeMarg = -inf;
else
    loglikeMarg = logsumexp(ll_theta_beta + arrayfun(logp_beta, betaGrid)) - log(nBeta);
end
end

%% ======= Priors =======
function lp = logPrior_theta(theta, paramSpec)
lp = 0;
if ~isfield(paramSpec,'priors') || isempty(paramSpec.priors), return; end
for j = 1:numel(theta)
    pr = paramSpec.priors(j);
    switch lower(pr.type)
        case 'uniform'
            a = pr.params(1); b = pr.params(2);
            if theta(j) < a || theta(j) > b, lp = -inf; return;
            else, lp = lp - log(b - a);
            end
        case 'loguniform'
            a = pr.params(1); b = pr.params(2);
            if theta(j) <= 0 || theta(j) < a || theta(j) > b, lp = -inf; return; end
            lp = lp - log(theta(j)) - log(log(b) - log(a));
        case 'normal'
            mu = pr.params(1); sig = pr.params(2);
            lp = lp - 0.5*((theta(j)-mu)/sig)^2 - log(sig*sqrt(2*pi));
        otherwise
            error('Unknown prior type: %s', pr.type);
    end
end
end

%% ======= Utilities =======
function val = thetaNamed(name, theta, spec)
idx = find(strcmp(spec.names, name), 1);
if isempty(idx), val = 0; else, val = theta(idx); end
end

function theta = paramTransform(model, xlog, paramSpec) %#ok<INUSD>
% Map log10-space vector xlog -> physical params theta (same order as names).
theta = zeros(size(xlog));
for j=1:numel(xlog)
    if paramSpec.is_pos(j)
        theta(j) = 10.^xlog(j);
    else
        theta(j) = xlog(j); % linear
    end
end
end

function sigma_base = get_sigma_column(statsMaybe, M)
    if exist('statsMaybe','var') && ~isempty(statsMaybe) && size(statsMaybe,2) >= 2
        sigma_base = max(statsMaybe(:,2), 1e-12);
    else
        sigma_base = max(std(M,0,2,'omitnan'), 1e-12);
    end
end

function s = logsumexp(a)
    amax = max(a(:)); if ~isfinite(amax), s = amax; return; end
    s = amax + log( sum( exp(a(:) - amax) ) );
end