Calculations/Data-Analyzer/+Plotter/plotSingleGaussianFitParameterPDF.m

function plotSingleGaussianFitParameterPDF(fitResults, scanValues, paramName, varargin)
%% plotSingleGaussianFitParameterPDF
% Plots 2D PDF (heatmap) of any parameter from single-Gaussian fit results
% for multiple scan parameters, with optional overlay of mean ± SEM.
%
% Inputs:
%   fitResults - struct array from single-Gaussian fits
%   scanValues - vector of scan parameter values
%   paramName  - string specifying the parameter to plot:
%                'A', 'mu', 'sigma', 'C', 'residual'
%
% Optional name-value pairs:
%   'OverlayMeanSEM' - logical, overlay mean ± SEM (default: true)
%   (other name-value pairs same as original function)

% --- Parse optional inputs ---
p = inputParser;
addParameter(p, 'Title', '', @(x) ischar(x) || isstring(x));
addParameter(p, 'XLabel', '', @(x) ischar(x) || isstring(x));
addParameter(p, 'YLabel', '', @(x) ischar(x) || isstring(x));
addParameter(p, 'FigNum', 1, @(x) isscalar(x));
addParameter(p, 'FontName', 'Arial', @ischar);
addParameter(p, 'FontSize', 14, @isnumeric);
addParameter(p, 'SkipSaveFigures', false, @islogical);
addParameter(p, 'SaveFileName', 'FitParameterPDF.fig', @ischar);
addParameter(p, 'SaveDirectory', pwd, @ischar);
addParameter(p, 'NumPoints', 200, @(x) isscalar(x));
addParameter(p, 'DataRange', [], @(x) isempty(x) || numel(x)==2);
addParameter(p, 'XLim', [], @(x) isempty(x) || numel(x)==2);
addParameter(p, 'YLim', [], @(x) isempty(x) || numel(x)==2);
addParameter(p, 'Colormap', @jet);
addParameter(p, 'PlotType', 'histogram', @(x) any(validatestring(x,{'kde','histogram'})));
addParameter(p, 'NumberOfBins', 50, @isscalar);
addParameter(p, 'NormalizeHistogram', true, @islogical);
addParameter(p, 'OverlayMeanSEM', true, @islogical);
parse(p, varargin{:});
opts = p.Results;

% --- Map paramName to index ---
paramMap = struct('A',1,'mu',2,'sigma',3,'C',4,'residual',5);
if ~isfield(paramMap,paramName)
    error('Invalid paramName. Must be one of: A, mu, sigma, C, residual');
end
paramIdx = paramMap.(paramName);

% --- Determine repetitions and scan parameters ---
N_params = numel(scanValues);
N_total  = numel(fitResults);
N_reps   = N_total / N_params;

% --- Extract chosen parameter values ---
paramValues = nan(N_reps, N_params);
for k = 1:N_total
    paramIdxScan = mod(k-1, N_params) + 1;
    repIdx       = floor((k-1)/N_params) + 1;

    if strcmpi(paramName, 'residual')
        paramValues(repIdx, paramIdxScan) = fitResults(k).residual;
    else
        paramValues(repIdx, paramIdxScan) = fitResults(k).pFit(paramIdx);
    end
end

% --- Mask out invalid zero fits ---
trueZeroMask = false(size(paramValues));
for i = 1:N_total
    paramIdxScan = mod(i-1, N_params) + 1;
    repIdx       = floor((i-1)/N_params) + 1;
    if ~fitResults(i).isValid && all(fitResults(i).pFit == 0)
        trueZeroMask(repIdx, paramIdxScan) = true;
    end
end

% --- Prepare data per scan parameter ---
dataCell = cell(N_params,1);
for i = 1:N_params
    dataCell{i} = paramValues(:,i);
end

% --- Determine y-range ---
if isempty(opts.DataRange)
    allData = cell2mat(dataCell(:));
    y_min   = min(allData);
    y_max   = max(allData);
else
    y_min   = opts.DataRange(1);
    y_max   = opts.DataRange(2);
end

% --- Prepare PDF grid/matrix ---
if strcmpi(opts.PlotType,'kde')
    y_grid     = linspace(y_min, y_max, opts.NumPoints);
    pdf_matrix = zeros(numel(y_grid), N_params);
else
    edges      = linspace(y_min, y_max, opts.NumberOfBins+1);
    binCenters = (edges(1:end-1) + edges(2:end))/2;
    pdf_matrix = zeros(numel(binCenters), N_params);
end

% --- Compute PDFs ---
for i = 1:N_params
    data = dataCell{i};
    data = data(~isnan(data));
    if isempty(data), continue; end
    if strcmpi(opts.PlotType,'kde')
        f = ksdensity(data, y_grid);
        pdf_matrix(:,i) = f;
    else
        counts = histcounts(data, edges);
        if opts.NormalizeHistogram
            binWidth = edges(2) - edges(1);
            counts   = counts / (sum(counts) * binWidth);
        end
        pdf_matrix(:,i) = counts(:);
    end
end

% --- Mask out scans without valid data ---
dataMask = ~isnan(paramValues);
maskPerScan = any(dataMask,1);
pdf_matrix(:, ~maskPerScan) = NaN;

% --- Plot heatmap ---
fig = figure(opts.FigNum); clf(fig);
set(fig, 'Color', 'w', 'Position', [100 100 950 750]);

if strcmpi(opts.PlotType,'kde')
    h = imagesc(scanValues, y_grid, pdf_matrix);
    set(h, 'AlphaData', ~isnan(pdf_matrix));
else
    h = imagesc(scanValues, binCenters, pdf_matrix);
    set(h, 'AlphaData', ~isnan(pdf_matrix));
end

set(gca, 'YDir', 'normal', 'FontName', opts.FontName, 'FontSize', opts.FontSize);
xlabel(opts.XLabel, 'FontSize', opts.FontSize, 'FontName', opts.FontName);
ylabel(opts.YLabel, 'FontSize', opts.FontSize, 'FontName', opts.FontName);
title(opts.Title, 'FontSize', opts.FontSize+2, 'FontWeight', 'bold');

cmap = feval(opts.Colormap);
colormap([1 1 1; cmap]);  % white for NaN
c = colorbar;
ylabel(c, 'Probability Density', 'Rotation', -90, 'FontName', opts.FontName, 'FontSize', opts.FontSize);

if ~isempty(opts.XLim), xlim(opts.XLim); end
if ~isempty(opts.YLim), ylim(opts.YLim); end

% --- Overlay mean ± SEM ---
if opts.OverlayMeanSEM
    meanParam = nanmean(paramValues,1);
    semParam  = nanstd(paramValues,0,1) ./ sqrt(sum(~isnan(paramValues),1));
    hold on;
    xVec = reshape(scanValues, 1, []);  
    fill([xVec, fliplr(xVec)], [meanParam - semParam, fliplr(meanParam + semParam)], ...
         [0.2 0.4 0.8], 'FaceAlpha',0.2, 'EdgeColor','none');
    plot(scanValues, meanParam, 'k-', 'LineWidth', 2);
end

end