overlayComparison.m

function outFile = overlayCorrection(noMocoInput, mocoInput, z, opts)
%OVERLAYCORRECTION  Grayscale anatomy with the moco correction painted in color.
%   Shows one fused image: the anatomy in grey, and wherever motion correction
%   CHANGED the image (moco vs no-moco), that change is overlaid as a colored
%   heatmap. Lets you see at a glance WHERE on the anatomy correction acted.
%
%   outFile = overlayCorrection(noMocoFile, mocoFile)
%   outFile = overlayCorrection(noMocoFile, mocoFile, z, opts)
%   (arrays accepted; pass PATHS to auto-save into the moco config folder)
%
%   INPUTS
%     noMocoInput, mocoInput : .mat paths (var auto-detected) or arrays.
%     z    : slice (default central).
%     opts : struct, all optional:
%            .signed   true = show + and - changes in two colors (default true).
%                      false = show magnitude of change in one hot colormap.
%            .thresh   fraction (0..1) of the max change below which nothing is
%                      overlaid, so noise stays grey (default 0.10).
%            .alpha    overlay opacity 0..1 (default 0.6).
%            .base     'moco' or 'nomoco' -- which image is the grey backdrop
%                      (default 'moco').
%            .outDir   save folder (default: moco file's folder; pwd if array).
%            .show     true to display on screen (default false = save only).
%
%   OUTPUT
%     outFile : cell of 3 PNG paths, one per plane (axial/coronal/sagittal).
%               Each PNG is a 1x3 panel: no-moco | moco | correction overlay.
%
%   COLOR MEANING (signed mode):
%     warm (red/yellow) = moco INCREASED intensity here
%     cool (blue/cyan)  = moco DECREASED intensity here
%     grey (no color)   = moco left it essentially unchanged
%   On a still subject the overlay is nearly absent (little correction).
%   On a mover, colored regions trace where ghosting/blur was removed.

    if nargin < 3, z = []; end
    if nargin < 4, opts = struct(); end
    o = setdefaults(opts, struct('signed',true,'thresh',0.10,'alpha',0.6, ...
                                 'base','moco','outDir','','show',false));

    [A,~]        = loadVol(noMocoInput);   % no-moco
    [B,mocoPath] = loadVol(mocoInput);     % moco
    assert(isequal(size(A),size(B)),'Images must match in size (%s vs %s).', ...
        mat2str(size(A)), mat2str(size(B)));
    % Magnitude images on a COMMON scale. Both are divided by their own 99th
    % percentile so intensity differences reflect real structural change, not a
    % global brightness offset between the two reconstructions.
    A = abs(double(A)); B = abs(double(B));
    A = A/prctile(A(:),99); 
    B = B/prctile(B(:),99);   % common scale
    sz = size(A);

    hiBright = prctile(B(:), 99.9);

    % slice indices for each plane. z (if given) sets the AXIAL slice; the
    % coronal/sagittal slices are taken at the centre of their axes (or scale
    % z proportionally if you passed one).
    if isempty(z)
        iAx = round(sz(3)/2); iCor = round(sz(2)/2); iSag = round(sz(1)/2);
    else
        iAx = z;  iCor = round(sz(2)/2); iSag = round(sz(1)/2);
    end

    % --- output location ---
    if isempty(o.outDir)
        if ~isempty(mocoPath), o.outDir = fileparts(mocoPath); else, o.outDir = pwd; end
    end
    if ~exist(o.outDir,'dir'), mkdir(o.outDir); end
    LSstr = parseLS(mocoPath); if isempty(LSstr), LSstr='config'; end

    % planes: name + how to pull the 2-D slice for A and B
    planes(1).name='axial';    planes(1).i=iAx;
    planes(1).get=@(V,i) squeeze(V(:,:,i));
    planes(2).name='coronal';  planes(2).i=iCor;
    planes(2).get=@(V,i) squeeze(V(:,i,:));
    planes(3).name='sagittal'; planes(3).i=iSag;
    planes(3).get=@(V,i) squeeze(V(i,:,:));

    % --- one 1x3 figure PER plane: no-moco | moco | overlay ---
    vis = 'off'; if o.show, vis='on'; end
    outFile = {};
    for pp = 1:3
        P = planes(pp);
        sA = P.get(A, clampIdx(P.i, planeDim(pp), sz));
        sB = P.get(B, clampIdx(P.i, planeDim(pp), sz));
        rgb = makeOverlay(sA, sB, o, hiBright);      % the colored fusion for this plane

        fig = figure('Visible',vis,'Color','w','Position',[40 60 1500 500]);
        tl = tiledlayout(fig,1,3,'Padding','compact','TileSpacing','compact');
        nexttile(tl); imagesc(sA, [0 hiBright]); axis image off; colormap(gca,'gray'); title('no-moco');
        nexttile(tl); imagesc(sB, [0 hiBright]); axis image off; colormap(gca,'gray'); title('moco');
        nexttile(tl); image(rgb); axis image off;
        if o.signed, title('correction overlay: red=+  blue=-');
        else,        title('correction overlay: |change|'); end
        sgtitle(sprintf('no-moco vs moco + overlay — %s plane — %s', P.name, LSstr), ...
            'Interpreter','none');

        f = fullfile(o.outDir, sprintf('overlayCorrection_%s_%s.png', P.name, LSstr));
        exportgraphics(fig, f, 'Resolution',200);
        if ~o.show, close(fig); end
        outFile{end+1} = f; %#ok<AGROW>
        fprintf('Saved %s-plane overlay:\n%s\n', P.name, f);
    end
end

% ============================================================
% HEATMAP CALCULATION for one 2-D slice pair (see comments in header).
% ============================================================
function rgb = makeOverlay(sA, sB, o, hiBright)
    % STEP 1 - per-pixel change:  D = moco - noMoco
    D = sB - sA;
    % which grayscale image sits UNDER the colour
    if strcmpi(o.base,'nomoco'), base = sA; else, base = sB; end
    % STEP 2 - grey anatomy backdrop replicated into R,G,B
    baseRGB = repmat(mat2gray(base, [0 hiBright]), [1 1 3]);
    % STEP 3 - normalize the change to [-1,1] by its own peak magnitude
    Dn = D / (max(abs(D(:))) + eps);
    % STEP 4 - threshold: only meaningful changes get coloured
    m  = abs(Dn) >= o.thresh;
    rgb = baseRGB;
    if o.signed
        % STEP 5a - split by sign; STEP 6 - opacity ~ |change|*alpha
        pos = m & (Dn>0);  neg = m & (Dn<0);
        w   = min(abs(Dn),1) * o.alpha;
        % STEP 7 - alpha-blend colour into grey
        rgb = blendColor(rgb, pos, w, [1.0 0.2 0.0]);   % + -> red-orange
        rgb = blendColor(rgb, neg, w, [0.0 0.4 1.0]);   % - -> blue
    else
        % STEP 5b - magnitude only, single hot colour
        w = min(abs(Dn),1) .* m * o.alpha;
        rgb = blendColor(rgb, m, w, [1.0 0.85 0.0]);
    end
    rgb = min(max(rgb,0),1);   % keep in [0,1] for image()
end

function d = planeDim(pp)
    switch pp, case 1, d=3; case 2, d=2; case 3, d=1; end
end
function i = clampIdx(i, dim, sz)
    i = max(1, min(i, sz(dim)));
end

% ================= helpers =================
function rgb = blendColor(rgb, maskPix, w, color)
    % Alpha-blend a solid 'color' into the RGB image at the masked pixels.
    % For each colour channel c:  out = grey*(1-w) + color(c)*w
    % where w is the per-pixel opacity (0 = keep grey, 1 = full color).
    % Pixels outside maskPix get w=0, so they stay exactly the grey backdrop.
    for c = 1:3
        ch = rgb(:,:,c);
        wc = w; wc(~maskPix) = 0;          % zero opacity outside the masked region
        ch = ch.*(1-wc) + color(c).*wc;    % the alpha blend
        rgb(:,:,c) = ch;
    end
end

function o = setdefaults(in, def)
    o = def; fn = fieldnames(in);
    for i=1:numel(fn), o.(fn{i}) = in.(fn{i}); end
end

function [V,pth] = loadVol(f)
    pth='';
    if isnumeric(f), V=f; if iscell(V), V=V{1}; end; return; end
    pth=f; S=load(f); fn=fieldnames(S); V=[];
    prefer={'x_refined_moco','x_nomoco','x0','x_refined_moco','x'};
    for k=1:numel(prefer), if isfield(S,prefer{k}), V=S.(prefer{k}); break; end; end
    if isempty(V)
        for i=1:numel(fn)
            v=S.(fn{i}); if iscell(v)&&~isempty(v), v=v{1}; end
            if isnumeric(v)&&ndims(v)==3&&min(size(v))>8, V=v; break; end
        end
    end
    if iscell(V), V=V{1}; end
    assert(~isempty(V),'No 3-D array found in %s.',f);
end

function ls = parseLS(pth)
    ls=''; if isempty(pth), return; end
    tok=regexp(pth,'(L\d+(?:\.\d+)?_S\d+(?:\.\d+)?)','tokens','once');
    if ~isempty(tok), ls=tok{1}; end
end

mocoFile = 'I:\yannic.delisle\scripts\sub-09\sliding_L3.5_S2.1\x_moco_mathilda_smoothedMotion_from_x_cs_tres2.1s_Nu120x120x120_delta2_sliding.mat'; 
noMocoFile = 'I:\yannic.delisle\scripts\sub-09\baseline_nomoco\x_nomoco_1bin_mathilda_manual_finalNu240.mat';

% overlayCorrection(noMocoFile, mocoFile);          % defaults, saves to config folder

% see it on screen, tune options:
% signed (default true) — two colors for +/− change. 
% Set false for a single hot colormap showing just the magnitude of change 
% (simpler, if the direction doesn't matter).
% 
% thresh (default 0.10) — changes below 10% of the max stay grey, so voxel 
% noise doesn't paint the whole image. Raise it to show only the biggest 
% changes, lower it to catch subtle ones.
% 
% alpha (default 0.6) — overlay opacity. Higher = more saturated color, 
% lower = anatomy shows through more.
% 
% base — whether the grey backdrop is the moco or no-moco image 
% (default moco).
opts.show = false; opts.thresh = 0.1; opts.alpha = 3;
overlayCorrection(noMocoFile, mocoFile, [], opts);