checkSignalSaturation.m

function results = checkSignalSaturation(rootFolder, tag, varargin)
% checkSignalSaturation  Quality-control check for calcium-indicator (GECI)
% saturation, using an empirical per-cell "headroom" metric.
%
% OPTIONAL NeuroSpikeX add-on. Elevated basal fluorescence (e.g. after a
% stimulus) can push cells toward the top of the indicator's dynamic range,
% which could suppress detectable transients and be mistaken for a biological
% change. This tool quantifies, per cell, how much usable dynamic range
% ("headroom") remains and flags cells operating near the ceiling.
%
%   Headroom_i = 1 - F0_i / Fsat
%       F0_i  = 10th percentile of cell i's RAW fluorescence trace (its basal level)
%       Fsat  = 95th percentile of ALL raw fluorescence values in the recording
%               (an empirical estimate of the indicator's usable ceiling)
%
% Headroom near 1 = far from saturation; near 0 = close to the ceiling.
%
% IMPORTANT: headroom must be computed from RAW fluorescence (before dF/F
% normalization). NeuroCa saves raw traces as 'result.mat' (cells x frames);
% this is the default source. Pass 'RawTraces' if your raw data live elsewhere.
%
% USAGE
%   results = checkSignalSaturation(rootFolder, '0plus');
%   % Derive the "low-headroom" cut from a baseline/sham recording:
%   results = checkSignalSaturation(rootFolder, '0plus', ...
%                 'ReferenceRaw', fullfile(shamRoot,'0plus','result.mat'));
%
% INPUTS
%   rootFolder - sample root folder (with per-timepoint subfolders and the
%                NeuroSpikeX outputs).
%   tag        - timepoint tag to evaluate (e.g. '0plus').
%
% NAME-VALUE OPTIONS
%   'RawTraces'      - cells x frames raw matrix, OR a path to a .mat holding
%                      it. Default: <rootFolder>/<tag>/result.mat.
%   'LowHeadroomCut' - numeric headroom threshold below which a cell is
%                      "low-headroom". Default [] -> 5th percentile of THIS
%                      recording's headroom (within-sample fallback).
%   'ReferenceRaw'   - raw traces (matrix or path) from a reference recording
%                      (e.g. sham/baseline); the cut becomes the 5th percentile
%                      of the reference headroom. Overrides 'LowHeadroomCut'.
%   'OutputDir'      - default <rootFolder>/AdditionalAnalyses/SignalSaturation
%   'SaveFigures'    - save figures + CSV (default true)
%   'ShowFigures'    - display figures (default true)
%
% OUTPUT (summary CSV + figures saved to OutputDir)
%   results - struct with per-cell headroom, F0, the low-headroom mask/cut,
%             overall and active-only low-headroom fractions, and the
%             headroom-vs-spike-rate association among active cells.

    p = inputParser;
    p.addParameter('RawTraces', fullfile(rootFolder, tag, 'result.mat'));
    p.addParameter('LowHeadroomCut', []);
    p.addParameter('ReferenceRaw', []);
    p.addParameter('OutputDir', fullfile(rootFolder,'AdditionalAnalyses','SignalSaturation'));
    p.addParameter('SaveFigures', true);
    p.addParameter('ShowFigures', true);
    p.parse(varargin{:});
    opt = p.Results;
    if ~exist(opt.OutputDir,'dir'), mkdir(opt.OutputDir); end

    fprintf('---------------------------------------------------------\n');
    fprintf('SIGNAL SATURATION / HEADROOM QC: %s\n', tag);
    fprintf('---------------------------------------------------------\n');

    %% ---------- Raw fluorescence ----------
    raw = resolveRaw(opt.RawTraces, 'RawTraces');
    if isempty(raw)
        error(['Could not load raw fluorescence for "%s".\n' ...
               'Expected NeuroCa raw traces (cells x frames) at:\n  %s\n' ...
               'Pass ''RawTraces'' with the correct matrix or path.'], tag, string(opt.RawTraces));
    end

    [headroom, F0, Fsat] = computeHeadroom(raw);

    %% ---------- Low-headroom threshold ----------
    if ~isempty(opt.ReferenceRaw)
        refRaw = resolveRaw(opt.ReferenceRaw, 'ReferenceRaw');
        refHead = computeHeadroom(refRaw);
        cut = prctile(refHead, 5);
        cutSource = 'reference recording (5th pct)';
    elseif ~isempty(opt.LowHeadroomCut)
        cut = opt.LowHeadroomCut;
        cutSource = 'user-specified';
    else
        cut = prctile(headroom, 5);
        cutSource = 'within-sample (5th pct)';
    end

    lowMask = headroom < cut;

    %% ---------- Map spike rate onto cells ----------
    [cellRate, active] = perCellSpikeRate(rootFolder, tag, size(raw,1));
    activeMask = false(size(raw,1),1);
    activeMask(active(active>=1 & active<=numel(activeMask))) = true;

    fracLowAll    = mean(lowMask);
    fracLowActive = NaN;
    if any(activeMask), fracLowActive = mean(lowMask(activeMask)); end

    fprintf('Cells: %d | Fsat (95th pct raw) = %.2f\n', size(raw,1), Fsat);
    fprintf('Low-headroom cut (%s): %.4f\n', cutSource, cut);
    fprintf('Low-headroom cells: %d / %d (%.1f%% overall', sum(lowMask), numel(lowMask), 100*fracLowAll);
    if ~isnan(fracLowActive)
        fprintf('; %.1f%% of active cells)\n', 100*fracLowActive);
    else
        fprintf(')\n');
    end

    %% ---------- Headroom vs activity (active cells) ----------
    rho = NaN; pval = NaN; slope = NaN;
    xg = headroom(activeMask); yg = cellRate(activeMask);
    good = isfinite(xg) & isfinite(yg);
    if sum(good) >= 5
        try
            mdl = fitlm(xg(good), yg(good), 'RobustOpts','on');
            slope = mdl.Coefficients.Estimate(2);
        catch, mdl = []; end
        [rho, pval] = corr(xg(good), yg(good), 'Type','Spearman','Rows','complete');
        fprintf('Headroom vs spike rate (active cells): Spearman rho=%.3f (p=%.3g), slope=%.4g\n', rho, pval, slope);
    else
        mdl = [];
        fprintf('Not enough active cells for headroom-vs-rate correlation.\n');
    end

    %% ---------- Package ----------
    results = struct();
    results.tag              = tag;
    results.headroom         = headroom;
    results.F0               = F0;
    results.Fsat             = Fsat;
    results.lowHeadroomCut   = cut;
    results.cutSource        = cutSource;
    results.lowMask          = lowMask;
    results.fracLowAll       = fracLowAll;
    results.fracLowActive    = fracLowActive;
    results.spikeRatePerCell = cellRate;
    results.activeMask       = activeMask;
    results.rho_headroom_rate = rho;
    results.p_headroom_rate   = pval;
    results.slope_headroom_rate = slope;

    %% ---------- Figures + CSV ----------
    if opt.SaveFigures || opt.ShowFigures
        vis = 'on'; if ~opt.ShowFigures, vis = 'off'; end

        fH = figure('Color','w','Visible',vis,'Position',[200 200 620 420]);
        histogram(headroom, 'Normalization','count','FaceColor',[0.3 0.5 0.8]); hold on;
        xline(cut,'--r','LineWidth',1.5,'Label','Low-headroom cut');
        xlabel('Headroom (0 = at ceiling, 1 = far from ceiling)'); ylabel('Number of cells');
        title(sprintf('%s: headroom distribution (%.1f%% low)', tag, 100*fracLowAll),'Interpreter','none');
        box off;

        fS = figure('Color','w','Visible',vis,'Position',[200 200 620 460]);
        scatter(xg(good), yg(good), 14, 'filled','MarkerFaceAlpha',0.6); hold on;
        if ~isempty(mdl)
            xx = linspace(min(xg(good)), max(xg(good)), 100)';
            plot(xx, predict(mdl, xx), 'k','LineWidth',1.5);
        end
        xline(cut,'--r','LineWidth',1.2,'Label','Low-headroom cut');
        xlabel('Headroom at this timepoint'); ylabel('Spike rate (Hz)');
        title(sprintf('%s: spike rate vs headroom (active cells)', tag),'Interpreter','none'); box off;
        if isfinite(rho)
            txt = sprintf('Spearman \\rho = %.3f (p = %.3g)', rho, pval);
            xl = xlim; yl = ylim;
            text(xl(1)+0.05*range(xl), yl(2)-0.08*range(yl), txt, 'BackgroundColor','w', ...
                'EdgeColor',[.85 .85 .85],'Margin',4);
        end

        if opt.SaveFigures
            saveFig(fH, fullfile(opt.OutputDir, sprintf('Headroom_distribution_%s', tag)));
            saveFig(fS, fullfile(opt.OutputDir, sprintf('Headroom_vs_spikeRate_%s', tag)));
            Tsum = table(string(tag), size(raw,1), Fsat, cut, string(cutSource), ...
                sum(lowMask), fracLowAll, fracLowActive, rho, pval, slope, ...
                'VariableNames', {'timepoint','n_cells','Fsat_p95','low_headroom_cut','cut_source', ...
                'n_low_headroom','frac_low_all','frac_low_active','rho_headroom_rate', ...
                'p_headroom_rate','slope_headroom_rate'});
            csvPath = fullfile(opt.OutputDir, sprintf('SignalSaturation_summary_%s.csv', tag));
            writetable(Tsum, csvPath);
            fprintf('Saved figures + summary to:\n  %s\n', opt.OutputDir);
        end
    end
    fprintf('---------------------------------------------------------\n');
end

%% ======================= local helpers =======================

function [H, F0, Fsat] = computeHeadroom(raw)
    F0   = prctile(raw, 10, 2);          % basal per cell
    Fsat = prctile(raw(:), 95);          % empirical ceiling
    H    = 1 - F0 ./ max(Fsat, eps);
    H    = max(0, min(1, H));
end

function raw = resolveRaw(src, name)
    raw = [];
    if isnumeric(src) && ~isempty(src)
        raw = double(src); return;
    end
    if (ischar(src) || isstring(src)) && isfile(src)
        S = load(src);
        if isfield(S,'result'), raw = double(S.result); return; end
        fn = fieldnames(S);
        % pick the largest 2-D numeric variable
        best = []; bestN = 0;
        for k = 1:numel(fn)
            v = S.(fn{k});
            if isnumeric(v) && ismatrix(v) && numel(v) > bestN
                best = v; bestN = numel(v);
            end
        end
        if ~isempty(best), raw = double(best); return; end
    end
    if isempty(raw)
        warning('%s could not be resolved to a raw trace matrix.', name);
    end
end

function [cellRate, active] = perCellSpikeRate(rootFolder, tag, nCells)
% Build a per-cell spike-rate vector (length nCells); inactive cells = 0.
    cellRate = zeros(nCells,1);
    active = [];
    fCS = fullfile(rootFolder, ['analyzed_fdata_' tag '.mat']);
    fQ  = fullfile(rootFolder, ['all_quantitative_analyses_' tag '.mat']);
    if ~isfile(fCS) || ~isfile(fQ)
        warning('Missing NeuroSpikeX outputs for "%s"; spike-rate overlay unavailable.', tag);
        return;
    end
    Scs = load(fCS);
    cs  = pickField(Scs, {['calcium_spikes_' tag], ['calcium_transients_' tag]}, {'calcium_spikes','calcium_transients'});
    Sq  = load(fQ, ['calcium_spike_calculations_' tag]);
    csc = pickField(Sq, {['calcium_spike_calculations_' tag]}, {'calcium_spike_calculations'});
    if isempty(cs) || isempty(csc), return; end
    % active_cells index the CLEANED fdata (ProcessNeuroCa drops empty cells),
    % while `raw` here still has one row per NeuroCa cell. Translate so the
    % spike rates land on the right raw traces.
    si = pickField(Scs, {['signal_info_' tag]}, {'signal_info'});
    active = neuroCaCellIDs(si, cs.active_cells(:), sprintf('%s active cells', tag));
    rate   = csc.cell_spike_rate(:);
    for i = 1:numel(active)
        id = active(i);
        if isfinite(id) && id>=1 && id<=nCells && i<=numel(rate)
            cellRate(id) = rate(i);
        end
    end
    active = active(isfinite(active));
end

function v = pickField(S, preferredNames, tokens)
    v = [];
    for k = 1:numel(preferredNames)
        if isfield(S, preferredNames{k}), v = S.(preferredNames{k}); return; end
    end
    fn = fieldnames(S);
    for t = 1:numel(tokens)
        hit = fn(contains(fn, tokens{t}, 'IgnoreCase', true));
        if ~isempty(hit), v = S.(hit{1}); return; end
    end
end

function saveFig(figHandle, basePath)
    set(figHandle,'Renderer','painters');
    try, print(figHandle, [basePath '.pdf'], '-dpdf','-vector','-bestfit'); catch, end
    try, saveas(figHandle, [basePath '.png']); catch, end
    try, savefig(figHandle, [basePath '.fig']); catch, end
end