function plotFitParameterPDF(fitResults, scanValues, paramName, varargin)
%% plotFitParameterPDF
% Author:       Karthik
% Date:         2025-10-06
% Version:      1.2
%
% Description:
%   Plots 2D PDF (heatmap) of any parameter from two-Gaussian fit results
%   for multiple scan parameters, with optional overlay of mean ± SEM.
%
% Inputs:
%   fitResults - struct array from fitTwoGaussianCurves
%   scanValues - vector of scan parameter values
%   paramName  - string specifying the parameter to plot:
%                'A1', 'mu1', 'sigma1', 'A2', 'mu2', 'sigma2'
%
% Optional name-value pairs (same as before, plus OverlayMeanSEM):
%   'OverlayMeanSEM' - logical, overlay mean ± SEM (default: true)

% --- 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('A1',1,'mu1',2,'sigma1',3,'A2',4,'mu2',5,'sigma2',6);
if ~isfield(paramMap,paramName)
    error('Invalid paramName. Must be one of: A1, mu1, sigma1, A2, mu2, sigma2');
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;
    paramValues(repIdx, paramIdxScan) = fitResults(k).pFit(paramIdx);
end

% --- Create mask of true zeros (from fillZeroStruct) ---
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 non-data regions (make them white) ---
dataMask = ~isnan(paramValues);          % where valid fits exist
maskPerScan = any(dataMask,1);           % scans that have any valid data
pdf_matrix(:, ~maskPerScan) = NaN;       % make empty scans white

% --- 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));  % make NaNs transparent
else
    h = imagesc(scanValues, binCenters, pdf_matrix);
    set(h, 'AlphaData', ~isnan(pdf_matrix));  % make NaNs transparent
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');

% --- Colormap and colorbar ---
cmap = feval(opts.Colormap);
colormap([1 1 1; cmap]);  % prepend white for NaN regions
c = colorbar;
if strcmpi(opts.PlotType,'kde') || opts.NormalizeHistogram
    ylabel(c, 'Probability Density', 'Rotation', -90, 'FontName', opts.FontName, 'FontSize', opts.FontSize);
else
    ylabel(c, 'Counts', 'Rotation', -90, 'FontName', opts.FontName, 'FontSize', opts.FontSize);
end

if ~isempty(opts.XLim)
    xlim(opts.XLim);
end

if ~isempty(opts.YLim)
    ylim(opts.YLim);
end

% --- Overlay mean ± SEM if requested ---
if opts.OverlayMeanSEM
    meanParam = nanmean(paramValues,1);
    semParam  = nanstd(paramValues,0,1) ./ sqrt(sum(~isnan(paramValues),1));
    hold on;
    xVec = reshape(scanValues, 1, []);  % ensures 1 × N_params
    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