function M = f_fold_k(M, sz, k)
% f_fold_k  Fold a matrix back into a tensor along mode-k
%
%   This function reverses the operation of f_unfold_k, converting a
%   matrix back into a K-way tensor of size `sz` by folding along mode-k.
%
%   Inputs:
%       M  - Matrix of size [sz(k), prod(sz(~k))] (mode-k unfolded tensor)
%       sz - Original size of the K-mode tensor
%       k  - Mode along which the unfolding was done (1-based indexing)
%
%   Output:
%       M  - Folded tensor of shape sz, with original mode order restored

K = length(sz);         % Total number of modes
R = k;                  % Row mode
C = [k+1:K, 1:k-1];     % Column modes (remaining modes after mode-k)

% Step 1: Reshape the matrix into size [sz(k), sz(k+1), ..., sz(K), sz(1), ..., sz(k-1)]
M = reshape(M, [sz(R), sz(C)]);

% Step 2: Permute to restore the original mode order
% Construct the inverse of [R, C], i.e., move mode-k back to position k
order_dst = [K - k + 2 : K, 1 : K - k + 1];
M = permute(M, order_dst);
end
