completePipelineCommented.m

%% ============================================================
%  IMAGE_4D.M
%  Build a 4D image (3D + time) using MonaLisa sliding window
%
%  Goal:
%  - Read the raw Siemens dataset
%  - Split the acquisition into temporal windows using timestamps
%  - Reconstruct one 3D image per temporal window
%  - Stack all 3D images into a 4D array x4D
%
%  This 4D image can later be used for rigid motion estimation with SPM.
%% ============================================================

clc

%% 0) Paths
% Add all required toolboxes / repos to MATLAB
addpath(genpath('/usr/src/app/mapVBVD'))
addpath(genpath('/usr/src/app/pulseq'))
addpath(genpath('/usr/src/app/rigid_motion_correction_monalisa_spm'))
addpath(genpath('/usr/src/app/spm'))
addpath(genpath('/usr/src/app/coilCompress'))
% Compile MEX only if missing. For PARALLEL sweeps, build ONCE beforehand
% so every concurrent job hits this guard and skips — concurrent rebuilds
% into the shared monalisa tree cause the 0-byte-mex / segfault race.
if ~exist('doCompileMex','var')
    doCompileMex = (exist('bmGridder_n2u_leight3_mex','file') ~= 3);
end
if doCompileMex
    compile_mex_for_monalisa();
else
    fprintf('MEX already present - skipping compile.\n');
end

disp(which('estimate_rigid_motion_params_spm'))
disp(which('apply_rigid_motion_to_kspace'))
disp(which('spm'))
disp(which('spm_jobman'))
disp(which('spm_select'))
%% 1) Dataset paths
% Raw Siemens data (.dat)
pathToSiemensRd = ...
    '/usr/src/app/Data/data/datasets/meas_MID00391_FID35695_HCP_HiFi_Flexy.dat';

% Main sequence file (.seq)
pathToPulseqSeq = ...
    '/usr/src/app/Data/data/HCP_HiFi_Flexy_Test_NSeg33_NShot686_BW500_SpoilFactor1_RfIncre117_TR2820.seq';

% Coil sensitivity file generated earlier by the main reconstruction script
coilSenseFile = ...
    '/usr/src/app/Data/data/coil_sens_240';

rootDir = '/usr/src/app/Data/scripts';
subject = 'sub-05-HiFi';

%% 2) Build raw data reader
fprintf('\n=== Step 2: Creating raw data reader ===\n');

% --- createRawDataReader: open the Siemens .dat and build a reader object.
%     Returns acquisitionParams (nSeg, nShot, timestamps, coils...). 'true' = auto-detect.
reader = createRawDataReader(pathToSiemensRd, true);
% Tell the reader the trajectory is defined by a Pulseq .seq file (non-Cartesian).
reader.acquisitionParams.traj_type = 'pulseq';
reader.acquisitionParams.pulseqTrajFile_name = pathToPulseqSeq;

% ------------------------------------------------------------
% IMPORTANT DEBUGGING NOTE:
% Earlier I had reshape errors when forcing nShot / nSeg from the .seq.
% For the first clean version, we do NOT overwrite them yet.
% We first inspect what the reader already thinks the dataset looks like.
% ------------------------------------------------------------

disp('Reader acquisition parameters:')
disp(reader.acquisitionParams)

% --- extract_seq_params: parse nSeg/nShot (and others) from the .seq filename/contents,
%     used to override whatever the reader auto-guessed.
seqParams = extract_seq_params(pathToPulseqSeq);

disp('Extracted sequence parameters:')
disp(seqParams)

if isfield(seqParams, 'n_seg')
    reader.acquisitionParams.nSeg = seqParams.n_seg;
elseif isfield(seqParams, 'nseg')
    reader.acquisitionParams.nSeg = seqParams.nseg;
end

if isfield(seqParams, 'n_shot')
    reader.acquisitionParams.nShot = seqParams.n_shot;
elseif isfield(seqParams, 'nshot')
    reader.acquisitionParams.nShot = seqParams.nshot;
end
% number of shots to ignore based on the non-steady state
reader.acquisitionParams.nShot_off = 15;

fprintf('After overwrite:\n');
fprintf('nSeg = %d\n', reader.acquisitionParams.nSeg);
fprintf('nShot = %d\n', reader.acquisitionParams.nShot);
fprintf('nSeg*nShot = %d\n', reader.acquisitionParams.nSeg * reader.acquisitionParams.nShot);

%% 3) Load raw data and trajectory
fprintf('\n=== Step 3: Loading raw data and trajectory ===\n');

% --- readRawData(flagSS, flagExcludeSI): load raw k-space [nCoils,nRead,nLines].
%     flagSS=true  -> reader drops nShot_off startup shots (here nShot_off=15).
%     flagExcludeSI=false -> keep the self-nav line (removed manually below).
y = reader.readRawData(true, false); % raw k-space data

disp('Original size of y:')
disp(size(y))

disp('expected trajectory lines from nSeg*nShot:')
disp(reader.acquisitionParams.nSeg * reader.acquisitionParams.nShot)

% ------------------------------------------------------------
% If the raw data contains 44 segments per shot, but the trajectory
% contains only 43 segments per shot.
%
% Difference:
%   8404 - 8213 = 191 = nShot
%
% So we likely have 1 self-navigation line per shot in y
% that is not present in t.
%
% We remove the first segment of each shot from y.
% ------------------------------------------------------------
nShotOff = reader.acquisitionParams.nShot_off;

nSeg  = reader.acquisitionParams.nSeg;   % should be 44
nShot = reader.acquisitionParams.nShot - nShotOff;  % should be 191
reader.acquisitionParams.nShot = nShot;

% y is [nCoils, nRead, nLines]
[nCoils, nRead, nLines] = size(y);

if nLines ~= nSeg * nShot
    error('Cannot reshape y into [nCoils, nRead, nSeg, nShot]: line count mismatch.');
end

% ---------------------------------------------------------------------
% SELF-NAVIGATOR REMOVAL — done manually (not via readRawData's
% flagExcludeSI) ON PURPOSE.
%
% Each shot's first segment is a self-navigator: a readout repeated at a
% FIXED k-space location every shot. It is NOT part of the imaging
% trajectory t (t has nSeg-1 segments/shot), so it must be removed or y
% and t won't align and gridding would misplace it. It would also
% over-weight k-space centre if left in.
%
% We remove it HERE, manually, rather than with readRawData(_, true)
% because:
%   (1) it keeps the navigator available for extraction/QC (the reader's
%       flagExcludeSI deletes it before we can use it), and
%   (2) the reader's flagExcludeSI reduces only a LOCAL nSeg and does not
%       trim the timestamp vector, so we would have to reconcile nSeg and
%       timestamps by hand anyway — the manual reshape below does both in
%       one place, keeping y and the timestamps consistent.
% ---------------------------------------------------------------------
removeSelfNav = true;

if removeSelfNav
% Self-navigator removal: reshape lines into seg x shot, delete segment 1 of every
% shot (the repeated k-space-center navigator, which the trajectory t does NOT contain),
% then flatten back. Reconciles y's line count with t's.
    y4 = reshape(y, [nCoils, nRead, nSeg, nShot]);
    y4(:,:,1,:) = [];
    y = reshape(y4, [nCoils, nRead, (nSeg-1)*nShot]);
else
    fprintf('Self-navigation removal disabled for this dataset.\n');
end

disp('Filtered size of y (after removing self-nav lines):')
disp(size(y))

fprintf('nSeg  = %d\n', reader.acquisitionParams.nSeg);
fprintf('nShot = %d\n', reader.acquisitionParams.nShot);
fprintf('Expected lines = %d\n', reader.acquisitionParams.nSeg * reader.acquisitionParams.nShot);

% Now build trajectory
% bmTraj: build the (nominal) k-space trajectory [3,nRead,nLines] from the Pulseq
% sequence. NOTE: this is the IDEAL trajectory; no gradient-delay/GIRF correction.
t = bmTraj(reader.acquisitionParams);

fprintf('Raw data lines   : %d\n', size(y,3));
fprintf('Trajectory lines : %d\n', size(t,3));

if size(y,3) ~= size(t,3)
    fprintf('Cropping raw data to match trajectory...\n');
    y = y(:,:,1:size(t,3));
end

disp('Size of trajectory t:')
disp(size(t))
disp('Size of raw data y:')
disp(size(y))

% We expect something like:
% y : [nRead x nCoils x nLines] or similar
% t : [3 x nRead x nLines]

%% 4) Load coil sensitivities
fprintf('\n=== Step 4: Loading coil sensitivities ===\n');

% or 'libre', whichever matches the current dataset
seqName = 'gre';  

coilDir = '/usr/src/app/head_motion_gre/C';

fprintf("coilDir = %s\n", coilDir);

disp(dir(coilDir))

% Load coil sensitivity maps C [Nx,Ny,Nz,nCoils], used for coil combination.
load(coilSenseFile, 'C');

% C is usually [Nx Ny Nz nCoils]
% Example: 48 x 48 x 48 x 52

%% 5) Reconstruction parameters
fprintf('\n=== Step 5: Defining reconstruction parameters ===\n');

% For now, keep these simple and explicit.
% Later we adapt them to your target resolution / FoV.
% match current coil sensitivity grid first
N_u  = [120 120 120];     
FoV  = 240;
dK_u = [1; 1; 1] / FoV;

fprintf('Reconstruction grid N_u = [%d %d %d]\n', N_u(1), N_u(2), N_u(3));
fprintf('FoV = %.1f\n', FoV);

C_size = size(C);
C_size = C_size(1:3);

% bmImResize: resample the coil maps onto the reconstruction grid N_u (e.g. 48^3 -> 120^3).
C = bmImResize(C, C_size, N_u);

%% 6) Sliding window parameters
fprintf('\n=== Step 6: Building sliding window mask ===\n');

% MonaLisa sequential binning uses timestamps to build temporal frames.
% Each frame contains all lines acquired during a fixed time window.
acquisitionParams = reader.acquisitionParams;

% Number of lines according to acquisition metadata
nLines_meta = acquisitionParams.nLine;

% Number of lines according to loaded raw data
% Depending on the reader, the line dimension may be the 3rd dimension.
nLines_y = size(y, 3);

fprintf('nLines from acquisitionParams = %d\n', nLines_meta);
fprintf('nLines from y                = %d\n', nLines_y);

% For safety, use the actual loaded data dimension.
nLines = nLines_y;

% Timestamps is still full length -> drop the same nShotOff shots worth
timeStamp = reader.acquisitionParams.timestamp;
timeStamp = reshape(timeStamp, [reader.acquisitionParams.nSeg, nShot + nShotOff]);
% drop the startup shots
timeStamp(:, 1:nShotOff) = [];
% write back trimmed
reader.acquisitionParams.timestamp = timeStamp(:).';

if numel(timeStamp) ~= nSeg * nShot
    error('Timestamp vector does not match nSeg*nShot.');
end

if removeSelfNav
    timeStamp2D = reshape(timeStamp, [nSeg, nShot]);
    timeStamp2D(1,:) = [];
    timeStamp = timeStamp2D(:).';
else
    timeStamp = timeStamp(:).';
end

% Shift so the first timestamp starts at zero
timeStamp = timeStamp - min(timeStamp);

% Convert to milliseconds.
% MonaLisa's sequential binning example uses a Siemens-specific factor of 2.5 ms. 
costTime = 2.5;
timestampMs = timeStamp * costTime;

% Make sure timestamp vector length matches nLines
timestampMs = timestampMs(:).';

if numel(timestampMs) < nLines
    error('Timestamp vector is shorter than the number of lines in y.');
elseif numel(timestampMs) > nLines
    warning('Timestamp vector is longer than y line count. Truncating to match y.');
    timestampMs = timestampMs(1:nLines);
end

% windowLengthSec and stride may be injected from the .sh; use a default if run standalone
if ~exist('windowLengthSec','var') || isempty(windowLengthSec)
    % default when run without an argument
    windowLengthSec = 3.5;   
end
if ~exist('strideRatio','var') || isempty(strideRatio)
    strideRatio = 0.7; 
end

% convert the stride into seconds
strideSec = strideRatio * windowLengthSec;

% window length and stride in miliseconds
windowLengthMs = windowLengthSec*1000;
strideMs = strideSec*1000;
% number of windows depends on the duration of the window and the stride
% overlap
notSSTime = 0;
totalDuration = max(timestampMs) - notSSTime;
nWindows = floor((totalDuration - windowLengthMs)/strideMs)+1;


% Preallocate sequential mask matrix
slidingMask = false(nWindows, nLines);
windowInfo.startMs  = zeros(1, nWindows);
windowInfo.endMs    = zeros(1, nWindows);
windowInfo.centerMs = zeros(1, nWindows);
windowInfo.nLines   = zeros(1, nWindows);

if nWindows < 1
    error('No temporal masks were created. Try increasing temporalWindowSec.');
end

motionAllowed = true(1, nLines);

for i = 1:nWindows
    windowStart = notSSTime + (i - 1) * strideMs;
    windowEnd   = windowStart + windowLengthMs;

    % Select all lines acquired during this temporal window
    windowMask = motionAllowed & (timestampMs >= windowStart) & (timestampMs < windowEnd);

    slidingMask(i, :) = windowMask;

    %Add metadata
    windowInfo.startMs(i) = windowStart;
    windowInfo.endMs(i) = windowEnd;
    windowInfo.centerMs(i) = (windowStart+windowEnd)/2;
    windowInfo.nLines(i) = sum(windowMask);
end

windowInfo.notSSTime       = notSSTime;
windowInfo.L_sec           = windowLengthSec;
windowInfo.S_sec           = strideSec;
windowInfo.overlap         = 1-strideSec/windowLengthSec;
windowInfo.nWindows        = nWindows;
windowInfo.totalDurationMs = totalDuration;

% Optional: report how many lines went into each frame
linesPerWindow = sum(slidingMask, 2);
disp('Lines per temporal frame:')
disp(linesPerWindow)

fprintf('Number of temporal windows created = %d\n', nWindows);

%% 7) Reconstruct each temporal frame
fprintf('\n=== Step 7: Reconstructing one 3D image per window ===\n');

delta = 2;
rho = 10 * delta;
nIter = 12;
nCGD = 8;
targetEnergy = 0.99;
nVirtualMax = 16;

% bmVolumeElement: per-sample density-compensation weights (Voronoi) for the full
% radial trajectory. Needed so gridding weights each k-space sample correctly.
ve_tot = bmVolumeElement(t, 'voronoi_full_radial3');

% 1) Normalization
% Full-acquisition reference reconstruction
% bmMathilda: fast gridded (non-iterative) reconstruction of the FULL acquisition.
% Used here only as a reference image for masking + normalization (not the final image).
x_tmp = bmMathilda(y, t, ve_tot, C, N_u, N_u, dK_u);

assert(exist('makeBrainMaskFromRecon', 'file') == 2, ...
    'makeBrainMaskFromRecon.mat was not found on the matlab path')
% Create a mask from this reconstruction
% makeBrainMaskFromRecon: Otsu-threshold + morphology on x_tmp to get a head-signal
% mask. 0.75 lowers the threshold (more inclusive). Used only for normalization below.
[maskBrain, maskInfo] = makeBrainMaskFromRecon(x_tmp, 0.75);

% Normalization
normalize_all  = median(abs(x_tmp(:)));
normalize_mask = median(abs(x_tmp(maskBrain)));
normalize_val  = normalize_mask;

fprintf('Normalization all voxels     : %.6g\n', normalize_all);
fprintf('Normalization inside mask    : %.6g\n', normalize_mask);
fprintf('Normalization ratio mask/all : %.4f\n', ...
    normalize_mask / normalize_all);

fprintf('Mask threshold: %.4f\n', maskInfo.threshold);

fprintf('Mask occupancy: %.2f %%\n', 100 * nnz(maskBrain) / numel(maskBrain));

% Normalize raw data by the in-mask median intensity, so 'delta' sits on a predictable
% O(1) scale (portable across datasets). Only computational use of the mask.
y = y / normalize_val;

% 2) Coil compression
% coilCompressSVD: SVD coil compression (nCoils -> nVirtual, keeping targetEnergy).
% Shrinks the per-window recon cost. Used for the motion-estimation CS pass only.
[y_svd, C_svd, Vcc, nVirtual, retainedEnergy] = coilCompressSVD( ...
    y, C, targetEnergy, nVirtualMax);

% 3) Reference image
% Reference image in the COMPRESSED coil basis (matches the solver's operators),
% used as the shared initial guess for every window.
x_ref = bmMathilda(y_svd, t, ve_tot, C_svd, N_u, N_u, dK_u);

% 4) Mitosis on compressed data
% bmMitosis: split the compressed data + trajectory into the temporal windows defined
% by slidingMask (one cell per window).
[y_windows, t_windows] = bmMitosis(y_svd, t, slidingMask);
ve_windows = bmVolumeElement(t_windows, 'voronoi_full_radial3');
% bmPermuteToCol: reshape each window's data to column layout [nPt,nCh] expected by the solver.
y_windows  = bmPermuteToCol(y_windows);

% 5) Sparse operators
% bmTraj2SparseMat: build the sparse gridding operator Gu and its adjoint Gut per window
% (maps non-Cartesian samples <-> Cartesian grid).
[Gu, Gut] = bmTraj2SparseMat(t_windows, ve_windows, N_u, dK_u);
ve_max = 10 * prod(dK_u);

% 6) Shared initialization
% Seed every window's iterate with the same reference image x_ref.
x_init = repmat({x_ref}, nWindows, 1);

temporalRes = strideSec;

fprintf('maskBrain voxels = %d / %d\n', nnz(maskBrain), numel(maskBrain));

outputFolder = fullfile(rootDir, subject, ...
    sprintf('sliding_L%.1f_S%.1f', windowLengthSec, strideSec));

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

%% Save a visualization of the mask
z = round(size(maskBrain,3) / 2);

fig = figure('Visible','off', 'Color','w');

imagesc(abs(x_tmp(:,:,z)));
axis image off
colormap gray
hold on

contour(double(maskBrain(:,:,z)), ...
    [0.5 0.5], ...
    'r', ...
    'LineWidth', 1);

title(sprintf( ...
    'Generated mask, threshold = %.4f', ...
    maskInfo.threshold));

maskQCFile = fullfile(outputFolder, 'generated_mask_QC.png');
exportgraphics(fig, maskQCFile, 'Resolution', 200);
close(fig);

fprintf('Mask QC image saved:\n%s\n', maskQCFile);

prepFile = fullfile(outputFolder, ...
    sprintf('prep_teva_tres%.1fs_Nu%dx%dx%d.mat', ...
    temporalRes, N_u(1), N_u(2), N_u(3)));

save(prepFile, ...
    'x_init', ...
    'y_windows', ...
    't_windows', ...
    've_windows', ...
    'windowLengthSec', ...
    'strideSec', ...
    'C_svd', ...
    'Gu', ...
    'Gut', ...
    'N_u', ...
    'FoV', ...
    'dK_u', ...
    'delta', ...
    'rho', ...
    'nCGD', ...
    've_max', ...
    'nIter', ...
    'windowInfo', ...
    'temporalRes', ...
    'outputFolder', ...
    'maskBrain', ...
    'maskInfo', ...
    'normalize_all', ...
    'normalize_mask', ...
    'normalize_val', ...
    '-v7.3');

%% ------------------------------------------------------------------------
%% Save experiment metadata
%% ------------------------------------------------------------------------

metadataFile = fullfile(rootDir, subject, ...
    sprintf('currentExperiment_L%.1f_S%.1f.mat', windowLengthSec, strideSec));

save(metadataFile, ...
    'prepFile', ...
    'outputFolder', ...
    'windowLengthSec', ...
    'strideSec', ...
    'temporalRes', ...
    'N_u',...
    'delta',...
    '-v7.3');

fprintf('\nCurrent experiment saved:\n%s\n', metadataFile);
%% ======================================================================
% Stage 2 (step3) runs in the SAME workspace as prep: x_init, y_windows,
% t_windows, ve_windows, C_svd, Gu, Gut, N_u, dK_u, delta, rho, nCGD,
% ve_max, nIter, windowInfo, temporalRes, outputFolder, maskBrain are all
% already live. No reload of currentExperiment / prepFile.

% 7) CS reconstruction
% STAGE: bmTevaMorphosia_chain — the temporally-COUPLED CS reconstruction. Solves all
% windows jointly with spatial + temporal regularization (delta = CS weight). Produces
% the 4D series x_cs used ONLY to estimate motion, not as the final image.
x_cs = bmTevaMorphosia_chain( ...
    x_init, [], [], ...
    y_windows, ve_windows, C_svd, ...
    Gu, Gut, N_u, ...
    [], [], ...
    delta, rho, 'normal', ...
    nCGD, ve_max, ...
    nIter, ...
    bmWitnessInfo(sprintf('motion_cs_delta%.4g', delta), []) );

temporalRes = strideSec;

deltaStr = strrep(num2str(delta), '.', 'p');

outputFolder = fullfile(rootDir, subject, sprintf('sliding_L%.1f_S%.1f', ...
        windowLengthSec, strideSec));

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

outputFile = fullfile(outputFolder,...
    sprintf('x_cs_tres%.1fs_Nu%dx%dx%d_delta%s_sliding.mat',...
    temporalRes, N_u(1), N_u(2), N_u(3), deltaStr));

save(outputFile,...
    'x_cs',...
    'FoV',...
    'N_u',...
    'temporalRes',...
    'windowInfo',...
    'delta',...
    'rho',...
    'maskBrain',...
    '-v7.3');


[~, fname] = fileparts(outputFile);
% Set the voxel size based on the current grid dimension size
voxelSize = FoV ./ N_u;
pixelDimensions = [voxelSize temporalRes];
outputPrefix = ['spm_' fname];

motionFile = fullfile(outputFolder, ['rp_spm_' fname '.txt']);

if exist(motionFile, 'file')
    fprintf('Motion file already exists, skipping SPM:\n%s\n', motionFile);
else
% estimate_rigid_motion_params_spm: run SPM rigid realignment on the x_cs series,
% writing rp_spm_*.txt = 6 motion params (Tx Ty Tz Rx Ry Rz) per window.
    estimate_rigid_motion_params_spm(outputFile, 'x_cs', pixelDimensions, outputPrefix);
end

clear x_cs
clear slidingMask
clear Gu Gut
clear y_windows
clear t_windows
clear ve_windows
clear y_svd
clear C_svd
clear x_init
clear x_tmp
clear Vcc
clear ve_b

fprintf('Large variables cleared.\n');
%% ======================================================================
%% Stage 3: motion-corrected final reconstruction (step5) 
% Use the live 'delta' written on the second stage, so step5's csFile name always matches.
selectedDelta = delta;
% STAGE 3: call step5 (below) to apply the motion trace to the FULL raw data and
% reconstruct the final image.
step5_recon_mathilda_moco(rootDir, subject, pathToPulseqSeq, ...
    pathToSiemensRd, coilSenseFile, selectedDelta, metadataFile);

% SLIDING WINDOW MATHILDA RECON
function step5_recon_mathilda_moco(rootDir, subject, pathToPulseqSeq,...
    pathToSiemensRd, coilSenseFile, selectedDelta, metadataFile)
% STEP5_RECON_MOCO
%
% Performs motion-corrected MRI reconstruction for all sequences and
% trajectories of a given subject. The function applies previously
% estimated rigid-body motion parameters (from SPM) to the raw k-space
% data before performing 1-bin (full acquisition) and 4-bin (temporal
% segmentation) reconstructions using the bmMathilda gridded reconstruction.
%
% Motion correction is applied directly in k-space by modifying the
% trajectory and phase of each readout line according to interpolated
% translations and rotations. The function maintains a fixed spatial
% reconstruction grid (N_u) but uses the Field-of-View (FoV) from the
% raw acquisition parameters.
%
% The function loops over all combinations of sequences and trajectories,
% skipping any missing data or missing motion parameters.
%
% OUTPUTS
%   Saves the following files in the corresponding subject/sequence/trajectory
%   motion_correction directory:
%     - normalization_moco.mat   : normalization factor after motion correction
%     - x0_1bin_moco.mat         : motion-corrected full acquisition reconstruction
%     - x0_4bins_moco.mat        : motion-corrected 4-bin reconstruction
%
% USAGE
%   step5_recon_mathilda_moco(rootDir, subject)
%
% INPUTS
%   rootDir : string
%       Path to the root dataset directory
%   subject : string
%       Subject identifier (e.g., 'sub-01')
%
% NOTES
%   1. Requires SPM motion parameter estimation to have been run beforehand
%      (produces rp_*.txt files in the motion_correction output folder).
%   2. The motion parameters are assumed to be in the form:
%          [Tx Ty Tz Rx Ry Rz]
%      where translations are in mm and rotations in radians.
%   3. Temporal resolution of the motion estimate is read from the step3
%      CS reconstruction if available, otherwise a default value is used.
%   4. Coil sensitivity maps and ROI masks are loaded from the derivatives
%      folder corresponding to the sequence.
%   5. Reconstruction grid (N_u) is fixed to [240 240 240], but FoV is
%      read from the raw data.
%
% EXAMPLE
%   step5_recon_mathilda_moco('/Volumes/sanDisk/5_december_2025', 'sub-01');
%
% See also: apply_rigid_motion_to_kspace, bmMathilda, bmMitosis

fprintf('\n=== Sliding-window motion-corrected reconstruction: %s ===\n', subject);

if nargin < 7 || isempty(metadataFile)
    metadataFile = fullfile(rootDir, subject, 'currentExperiment.mat');
end
load(metadataFile)

delta = selectedDelta;
deltaStr = strrep(num2str(delta), '.', 'p');

csFile = fullfile(outputFolder,...
    sprintf('x_cs_tres%.1fs_Nu%dx%dx%d_delta%s_sliding.mat',...
    temporalRes,N_u(1),N_u(2),N_u(3), deltaStr));

assert(exist(csFile,'file')==2,...
    'Missing reconstruction:\n%s',csFile);

[~,csBase] = fileparts(csFile);

motionFile = fullfile(outputFolder,...
    ['rp_spm_' csBase '.txt']);

assert(exist(motionFile,'file')==2,...
    'Missing motion file:\n%s',motionFile);

outputFile = fullfile(outputFolder,...
    sprintf('x_moco_mathilda_smoothedMotion_from_%s.mat',csBase));


fprintf('\nUsing reconstruction file:\n%s\n', csFile);
fprintf('Using motion file:\n%s\n', motionFile);

S = load(csFile, 'FoV', 'N_u', 'windowInfo');

FoV = S.FoV;
% DUAL-GRID BY DESIGN: motion was estimated on S.N_u (coarse, e.g. 120^3);
% the final motion-corrected image is delivered at N_u below (fine, 240^3).
% Motion params are physical (mm / rad), so they transfer across grids.
% grid the motion was estimated on
N_u_estim = S.N_u;           
% delivery grid
N_u = [240 240 240];         
windowInfo = S.windowInfo;

motionTimeMs = windowInfo.centerMs(:);

motionParams = load(motionFile);

smoothMotion = true;
if smoothMotion
    motionParams_raw = motionParams;
    smoothWin = 5;

    motionParams = movmedian(motionParams, smoothWin, 1);
    motionParams = movmean(motionParams, smoothWin, 1);
    % Keep the first valid estimated pose as the reference.
    % This prevents the clamped early k-space lines from receiving
    % a non-zero smoothed first motion estimate.
    motionParams = motionParams - motionParams(1,:);

    smoothFile = fullfile(outputFolder, ...
        sprintf('rp_smoothed_from_%s.mat', csBase));

    save(smoothFile, 'motionParams_raw', 'motionParams', ...
        'smoothWin', '-v7.3');
end
  
nFrames = size(motionParams, 1);
if numel(motionTimeMs) ~= nFrames
    error(['Mismatch between motion frames and window centers:\n'...
        'motionParams rows %d, windowInfo.centerMs =%d'],...
        nFrames, numel(motionTimeMs));
end
%% Reload raw data and trajectory
fprintf('Loading raw data and trajectory...\n');

reader = createRawDataReader(pathToSiemensRd, true);
reader.acquisitionParams.traj_type = 'pulseq';
reader.acquisitionParams.pulseqTrajFile_name = pathToPulseqSeq;

seqParams = extract_seq_params(pathToPulseqSeq);

if isfield(seqParams, 'n_seg')
    reader.acquisitionParams.nSeg = seqParams.n_seg;
elseif isfield(seqParams, 'nseg')
    reader.acquisitionParams.nSeg = seqParams.nseg;
end

if isfield(seqParams, 'n_shot')
    reader.acquisitionParams.nShot = seqParams.n_shot;
elseif isfield(seqParams, 'nshot')
    reader.acquisitionParams.nShot = seqParams.nshot;
end

reader.acquisitionParams.nShot_off = 15;
nShotOff = reader.acquisitionParams.nShot_off;

p = reader.acquisitionParams;

% step5 reloads the FULL raw data from scratch (independent of the windows).
y_tot = reader.readRawData(true,false);
nSeg = p.nSeg;
nShot = p.nShot - nShotOff;

[nCoils, nRead, nLinesRaw] = size(y_tot);

if nLinesRaw ~= nSeg * nShot
    error('Cannot reshape raw data: nLinesRas does not match nSeg*nShot.');
end

%% Remove self-navigation line
% or false depending on dataset
removeSelfNav = true; 

if removeSelfNav
    fprintf('Removing first segment of each shot from raw data...\n');
    y4 = reshape(y_tot, [nCoils, nRead, nSeg, nShot]);
    y4(:,:,1,:) = [];
    y_tot = reshape(y4, [nCoils, nRead, (nSeg-1)*nShot]);
else
    fprintf('Self-navigation removal disabled.\n');
end

t_tot = bmTraj(p);

fprintf('Raw data lines     : %d\n', size(y_tot,3));
fprintf('Trajectory lines   : %d\n', size(t_tot,3));

if size(y_tot,3) ~= size(t_tot,3)
    fprintf('Cropping raw data to match trajectory...\n');

    y_tot = y_tot(:,:,1:size(t_tot,3));

    fprintf('New raw data size:\n');
    disp(size(y_tot));
end

fprintf('Size y_tot after filtering:\n');
disp(size(y_tot));
fprintf('Size t_tot:\n');
disp(size(t_tot));

%% Build line timestamps
timeStamp = p.timestamp;
timeStamp = reshape(timeStamp, [nSeg, nShot + nShotOff]);
timeStamp(:, 1:nShotOff) = [];
timeStamp = timeStamp(:);

if numel(timeStamp) ~= nSeg*nShot
    error('Timestamp vector does not match nSeg*nShot.');
end

if removeSelfNav
    timeStamp2D = reshape(timeStamp, [nSeg, nShot]);
    timeStamp2D(1,:) = [];
    timeStamp = timeStamp2D(:).';
else
    timeStamp = timeStamp(:).';
end

timeStamp = timeStamp - min(timeStamp);

costTime = 2.5;
lineTimeMs = double(timeStamp) * costTime;
lineTimeMs = lineTimeMs(:);

nLines = size(y_tot, 3);

if numel(lineTimeMs) ~= nLines 
    error('lineTimeMs length (%d) does not match y_tot lines (%d).', ...
        numel(lineTimeMs), nLines);
end

%% Clamp raw line times to the valid motion-estimation interval
% This avoids extrapolating motion before the first SPM estimate
% and after the last SPM estimate.

lineTimeMsForCorrection = lineTimeMs;

lineTimeMsForCorrection(lineTimeMsForCorrection < motionTimeMs(1)) = ...
    motionTimeMs(1);

lineTimeMsForCorrection(lineTimeMsForCorrection > motionTimeMs(end)) = ...
    motionTimeMs(end);

fprintf('Motion time range: %.2f ms to %.2f ms\n', ...
    motionTimeMs(1), motionTimeMs(end));

fprintf('Raw line time range: %.2f ms to %.2f ms\n', ...
    min(lineTimeMs), max(lineTimeMs));

fprintf('Clamped line time range: %.2f ms to %.2f ms\n', ...
    min(lineTimeMsForCorrection), max(lineTimeMsForCorrection));

%% Apply motion correction to k-space
% apply_rigid_motion_to_kspace: the ACTUAL correction. Interpolates the SPM motion to
% each line's timestamp and applies it per line (rotate trajectory + translation phase).
% Output y_corr/t_corr = the full acquisition with motion removed.
[y_corr, t_corr, ~] = apply_rigid_motion_to_kspace(...
    y_tot,...
    t_tot,...
    lineTimeMsForCorrection,...
    motionParams,...
    motionTimeMs);

clear y_tot t_tot

%% Load coil sensitivities
load(coilSenseFile, 'C');
dK_u = [1 1 1]/FoV;
%% Recompute volume elements after trajectory correction
fprintf('Computing corrected volume elements...\n');
% Recompute density weights: the trajectory changed (rotated) after motion correction.
ve_corr = bmVolumeElement(t_corr, 'voronoi_full_radial3');
C_size = size(C);
C_size = C_size(1:3);

C = bmImResize(C, C_size, N_u);
%% Final full acquisition reconstruction
fprintf('Running final motion-corrected Mathilda reconstruction...\n');
% FINAL IMAGE (gridded): bmMathilda on the full motion-corrected data at N_u (240^3).
% Fast, unbiased, but streaky on undersampled data -> refined by Steva next.
x_moco = bmMathilda(...
    y_corr,...
    t_corr,...
    ve_corr,...
    C,...
    N_u,...
    N_u,...
    dK_u);
%% Save result    
save(outputFile,...
    'x_moco',...
    'motionParams',...
    'motionTimeMs',...
    'lineTimeMs',...
    'FoV',...
    'N_u',...
    '-v7.3');
fprintf('Saved motion-corrected reconstruction:\n%s\n', outputFile);


%% Refined optimized reconstruction with Steva
% fprintf('Running optimized motion-corrected Steva reconstruction...\n');
runSteva = true;

if runSteva
    %% Final single-frame regularized reconstruction (Steva / TV-ADMM)
    frSize = N_u;                         % 240^3 delivery grid (frSize == N_u)
    ve_max = 10 * prod(dK_u(:));
    nCGD   = 4;
    nIter  = 20;
    
    % Column layout: y [nPt, nCh], ve [1, nPt].  bmSteva col-reshapes C and x0 itself.
    y_col  = reshape(permute(y_corr, [2 3 1]), [], size(y_corr,1));
    ve_col = reshape(ve_corr, 1, []);
    
    % Sparse operators for the corrected trajectory at the delivery grid
    [Gu, Gut] = bmTraj2SparseMat({t_corr}, {ve_corr}, frSize, dK_u);
    
    % This build has NO adaptive mode: delta must be supplied and is weighed
    % against the data term, so normalize to ~O(1) first (step5 is raw-scale),
    % then rescale the result back. Makes delta interpretable and portable.
    sclN   = prctile(abs(x_moco(:)), 99);
    y_col  = y_col  / sclN;
    x0n    = x_moco / sclN;
    
    witnessInfo = bmWitnessInfo('steva_final', []);
    witnessInfo.save_witnessIm_flag = 0;
    % TV weight on the normalized scale (starting guess)
    delta = 0.1;   
    % ADMM penalty (rho = 10*delta convention)
    rho   = 10 * delta;     
    
    % Installed signature (15 args):
    % bmSteva(x0, z, u, y, ve, C, Gu, Gut, frSize, delta, rho, nCGD, ve_max, nIter, witnessInfo)
    % z,u are ADMM dual vars -> [] auto-inits them.
    % bmSteva: single-frame TV-regularized iterative recon. Takes x_moco as init and
    % de-streaks/denoises it. delta = TV weight (on the normalized scale). Final deliverable.
    x_refined_moco = bmSteva( ...
        x0n, [], [], ...
        y_col, ve_col, C, ...
        Gu{1}, Gut{1}, frSize, ...
        delta, rho, nCGD, ve_max, ...
        nIter, witnessInfo);
    
    x_refined_moco = x_refined_moco * sclN;   % restore original intensity scale
    
        % Persist the refined result + QC (bmSteva output was previously discarded) ---
        outputFileRefined = fullfile(outputFolder, ...
            sprintf('x_refined_moco_steva_delta%s_from_%s.mat', ...
                    strrep(num2str(delta),'.','p'), csBase));
        save(outputFileRefined, ...
            'x_refined_moco', 'x_moco', ...
            'motionParams', 'motionTimeMs', 'lineTimeMs', ...
            'FoV', 'N_u', 'delta', 'rho', 'nIter', 'sclN', '-v7.3');
        fprintf('Saved refined Steva reconstruction:\n%s\n', outputFileRefined);
    
        % side-by-side QC: Mathilda vs Steva, mid-slice
        zc  = round(size(x_refined_moco,3)/2);
        fig = figure('Visible','off','Color','w');
        tl  = tiledlayout(fig,1,2,'Padding','compact','TileSpacing','compact');
        nexttile(tl); imagesc(abs(x_moco(:,:,zc)));         axis image off; colormap gray; title('Mathilda (x\_moco)');
        nexttile(tl); imagesc(abs(x_refined_moco(:,:,zc))); axis image off; colormap gray; title(sprintf('Steva \\delta=%.3g', delta));
        exportgraphics(fig, fullfile(outputFolder, ...
            sprintf('steva_delta%s_QC_from_%s.png', strrep(num2str(delta),'.','p'), csBase)), 'Resolution', 200);
        close(fig);
    
        clear Gu Gut
else
    fprintf('Skipping Steva optimization.\n');
    fprintf('Using Mathilda motion-corrected reconstruction only.\n');
end
clear y_corr t_corr ve_corr C

fprintf('\nAll sliding-window motion-corrected reconstructions complete for subject %s\n', subject);
end