plot_beta_uq.m

function plot_beta_uq(resultsFile, expDataFunc)
% Clean, publication-ready top panel:
%  - marginal p(beta | D, Mwin)
%  - conditional p(beta | D, theta_MAP, Mwin)  (if theta_MAP exists)
%  - MAP markers + 95% HPD band
% Usage: plot_beta_uq_top('hierarchical_results.mat', @expDataPrep)

if nargin<1 || isempty(resultsFile), resultsFile = 'results_UM3.mat'; end
if nargin<2 || isempty(expDataFunc), expDataFunc = @expDataPrep; end

S = load(resultsFile);
Models=S.Models; post=S.post(:).'; beta=S.betaGrid(:).'; 
kappa=S.kappa; m_floor=S.m_floor; useHetero=S.useHetero;
Thresholds=S.Thresholds; mask_saved=S.mask; folder=S.folder;

[~,best] = max(post); Mw = Models(best); modelName = Mw.name;

% --- Data & exact mask (match BIMR) ---
expDataFunc(); [T,J]=size(Rmatrix);
mask = reconcile(mask_saved,T,J);

% Baseline sigmas (match BIMR fallback)
sigmaR0 = sigcol(R_stats,    Rmatrix,    1e-12);
sigmaRd = sigcol(Rdot_stats, Rdotmatrix, 1e-12);

% Hetero weights (same sign as BIMR scorer)
tcrow = repmat(tc(:).',T,1);
epsdot = -2.*(Rdotmatrix./max(Rmatrix,1e-12)).*tcrow;
sr_th  = srstar(tc(:).',T,Thresholds);
a_full = 1./(1+exp(-kappa.*(sr_th-abs(epsdot))./max(sr_th,eps)));
w_full = m_floor + (1-m_floor).*a_full;

% --- Marginal posterior from BIMR ---
p_marg = Mw.post_beta(:).'; p_marg = p_marg./max(sum(p_marg),eps);
[bMAP,HPD] = map_hpd(beta,p_marg,0.95);

% --- Conditional posterior at theta_MAP (optional) ---
p_cond = []; bMAPc = NaN;
if ~isempty(Mw.map_subs)
    [p_cond,bMAPc] = cond_beta_at_thetaMAP(Mw.map_subs,folder,Mw.name, ...
        Rmatrix,Rdotmatrix,tmatrix,mask,sigmaR0,sigmaRd,w_full, ...
        Rmax_range,tc,useHetero,m_floor,beta);
end

% --- Figure (linear axis, clean style) ---
set(0,'defaultTextInterpreter','latex');
set(0,'defaultAxesTickLabelInterpreter','latex');
set(0,'defaultLegendInterpreter','latex');

fs = 12;                % base font size for labels
lw = 4;                 % line width for plots
colBlue = [0 114 178]/255;  % MATLAB blue
colRed  = [1 0 0];           % MATLAB red
colHPD  = [0.85 0.80 0.90];  % light gray fill (unused but safe to keep)

% Auto-focus on HPD with margin
xlo = max(min(beta), HPD(1) - 0.25*(HPD(2)-HPD(1)));
xhi = min(max(beta), HPD(2) + 0.25*(HPD(2)-HPD(1)));
if ~isfinite(xlo) || ~isfinite(xhi) || xlo>=xhi, xlo=min(beta); xhi=max(beta); end
ymax = max([p_marg, p_cond], [], 'omitnan')*1.15; if ~isfinite(ymax) || ymax<=0, ymax=1; end

f=figure('Color','w','Position',[120 130 640 360]); ax=gca; hold on; box off;
% patch([HPD(1) HPD(2) HPD(2) HPD(1)], [0 0 ymax ymax], colHPD, 'EdgeColor','none','FaceAlpha',0.25);
h1 = plot(beta, p_marg, '-',  'Color',colBlue,'LineWidth',lw);      % blue solid
h2 = [];
if ~isempty(p_cond), h2 = plot(beta, p_cond, '--', 'Color',colRed,'LineWidth',lw); end  % red dashed
% plot([bMAP bMAP],[0 ymax], ':',  'Color',colBlue,'LineWidth',1.2);         % blue dotted
% if ~isnan(bMAPc), plot([bMAPc bMAPc],[0 ymax], '--', 'Color',colRed,'LineWidth',1.2); end  % red dashed

xlim([xlo xhi]); ylim([0 ymax]); grid off; ax.GridAlpha=0.12;

xlabel('$\beta$','FontSize',fs);
ylabel('$P(\beta \mid D, \cdot)$','FontSize',fs);
% leg = {'marginal $p(\beta \mid D,M_{\mathrm{win}})$'};
% if ~isempty(p_cond), leg{end+1}='conditional $p(\beta \mid D,\theta_{\mathrm{MAP}},M_{\mathrm{win}})$'; end
% if ~isempty(p_cond)
%     L=legend([h1 h2], leg, 'Location','northoutside');
% else
%     L=legend(h1, leg, 'Location','northoutside');
% end
% L.Box='off';

if exist('formatPlot','file')==2
    formatPlot();
end

% Export (vector)
% out = sprintf('beta_uq_%s_top.pdf', modelName);
% exportgraphics(f,out,'ContentType','vector');
% fprintf('Saved: %s\n', out);

% Console summary
fprintf('beta_MAP (marg) = %.4g | 95%% HPD = [%.4g, %.4g]\n', bMAP, HPD(1), HPD(2));
if ~isnan(bMAPc), fprintf('beta_MAP (cond@thetaMAP) = %.4g\n', bMAPc); end
end

% ===== helpers (compact) =====
function [pcond,bmap]=cond_beta_at_thetaMAP(subs,folder,modelName, ...
    R,Rd,t,mask,sR0,sRd,w, Rmax_range,tc,useHetero,m_floor,beta)
M=load(fullfile(folder,[modelName '.mat'])); C=M.Sims{subs{:}};
t0=C(:,1);[t0,uq]=unique(t0,'stable'); Rm=C(uq,2); Rdm=C(uq,3);
Rmax_mean=mean(Rmax_range,'omitnan'); tc_mean=mean(tc,'omitnan');
s = Rmax_range./max(Rmax_mean,eps); tau = tc./max(tc_mean,eps);
[T,J]=size(R); ft=isfinite(t); fR=isfinite(R); fRd=isfinite(Rd);
sR0=repmat(sR0,1,J); sRd=repmat(sRd,1,J); ll=zeros(1,numel(beta));
for j=1:J
    use0=mask(:,j)&ft(:,j)&fR(:,j)&fRd(:,j); if ~any(use0), continue; end
    tj=t0./tau(j); Rth=interp1(tj, Rm./s(j),      t(use0,j),'linear');
                   Rdth=interp1(tj, Rdm.*(tau(j)/s(j)), t(use0,j),'linear');
    rR=R(use0,j)-Rth; rRd=Rd(use0,j)-Rdth;
    vR0=max(sR0(use0,j).^2,1e-24); vRd0=max(sRd(use0,j).^2,1e-24);
    vscale = (useHetero) .* (1./max(w(use0,j),m_floor)) + (~useHetero);
    vR=vR0.*vscale; vRd=vRd0.*vscale;
    for b=1:numel(beta)
        vb=beta(b)^2; 
        ll(b)=ll(b) -0.5*( nansum((rR.^2)./(vb*vR)) + nansum(log(2*pi*vb*vR)) ...
                          +nansum((rRd.^2)./(vb*vRd))+ nansum(log(2*pi*vb*vRd)) );
    end
end
halfCauchy = @(b) (2/pi) * (1 ./ (1 + b.^2)) .* (b>0);
logw = trapezoid_log_weights(beta, halfCauchy);
lp = ll + logw; pcond = exp(lp - lse(lp));
[~,ix]=max(pcond); bmap=beta(ix);
end

function [map,hpd]=map_hpd(x,p,alpha)
[~,ix]=max(p); map=x(ix);
[ps,idx]=sort(p(:),'descend'); cs=cumsum(ps); k=find(cs>=alpha,1); keep=idx(1:k);
hpd=[min(x(keep)) max(x(keep))];
end

function L=srstar(tcrow,T,Th)
if iscolumn(tcrow), tcrow=tcrow.'; end
m=lower(string(Th.mode));
switch m
    case "dim", L = max(Th.sr_dim,0).*tcrow;
    case "nd",  L = Th.sr_nd .* ones(size(tcrow));
    otherwise,  base=Th.auto_base; if isempty(base)||~isfinite(base)||base<=0, base=1e5; end
               L = base .* tcrow;
end
L=repmat(L,T,1);
end

function s = sigcol(stats,M,floorv)
if exist('stats','var') && ~isempty(stats) && size(stats,2)>=2
    s = max(stats(:,2), floorv);
else
    s = max(std(M,0,2,'omitnan'), floorv);
end
end

function m = lse(a), mx=max(a(:)); if ~isfinite(mx), m=mx; else, m=mx+log(sum(exp(a(:)-mx))); end, end

function M=reconcile(Min,T,J)
if isempty(Min), M=false(T,J); return; end, [t0,j0]=size(Min);
if t0==T && j0==J, M=logical(Min); return; end
M=false(T,J); M(1:min(t0,T),1:min(j0,J)) = logical(Min(1:min(t0,T),1:min(j0,J)));
end

function logw_beta = trapezoid_log_weights(betaGrid, pdfFcn)
B = numel(betaGrid);
if B==1, logw_beta = 0; return; end
d = zeros(size(betaGrid));
d(1)   = 0.5*(betaGrid(2)-betaGrid(1));
d(end) = 0.5*(betaGrid(end)-betaGrid(end-1));
for b=2:B-1
    d(b) = 0.5*(betaGrid(b+1)-betaGrid(b-1));
end
raw = pdfFcn(betaGrid) .* max(d, eps);
m   = max(log(raw));
logw_beta = log(raw) - (m + log(sum(exp(log(raw)-m))));
end

function formatPlot()
    ax  = gca;
    fig = gcf;

    % ---------- Make window as large as possible with fixed aspect ratio ----------
    ar = 1.5;                                 % desired width/height (match pbaspect)
    scr = get(0,'ScreenSize');                % [left bottom width height] (px)
    pad = 60;                                 % margin for OS chrome

    scrW = scr(3) - 2*pad;
    scrH = scr(4) - 2*pad;

    W = scrW; H = W/ar;                       % try width-limited
    if H > scrH                                % if too tall, limit by height instead
        H = scrH; W = H*ar;
    end

    left   = pad + (scr(3) - 2*pad - W)/2;    % center on screen
    bottom = pad + (scr(4) - 2*pad - H)/2;

    set(fig, 'Units','pixels', 'Position',[left bottom W H]);

    % ---------- Axes styling ----------
    box(ax,'on');
    set(ax,'LineWidth',2,'FontSize',36,'TickLabelInterpreter','latex');
    pbaspect(ax,[ar 1 1]);                    % keep same visual ratio as window

    % ---------- Tighten horizontally only (keep vertical unchanged) ----------
    drawnow;                                  % update layout metrics
    set(ax,'Units','normalized');

    % --- Avoid cutting off labels or title (leave small margin) ---
    ti  = get(ax,'TightInset');               % [left bottom right top]
    pos = get(ax,'Position');                 % [x y w h]
    
    marginTop = 0.05;                         % small top margin (normalized units)
    marginBottom = 0.02;                      % small bottom margin
    
    pos(1) = ti(1);                           % left padding
    pos(3) = 1 - ti(1) - ti(3);               % width
    pos(2) = ti(2) + marginBottom;            % lift up slightly
    pos(4) = 1 - ti(2) - ti(4) - marginTop;   % reduce height slightly for top room
    
    set(ax,'Position',pos);

    xlim([0,10])
    ylim([0,0.953])%0.765])

    % Automatically detect current upper limit and add a small buffer
    yl = get(ax,'YLim');
    ymax_pad = yl(2) * 1.05;   % add 5% extra space at the top
    ylim([yl(1), ymax_pad]);

    % Respect limits set earlier and add x-ticks every 0.2
    xl = get(ax,'XLim');
    set(ax,'XTick', xl(1):2:xl(2));


    % ---------- Clean vector export defaults ----------
    set(fig,'Renderer','painters','Color','w');
    set(ax,'Color','w');                   % transparent axes bg
end