%% analyze_visualsearch.m
% Plots median RT as a function of set size for pop-out vs conjunction.
% Robust to single file or directory of files; adds optional bootstrap CIs.

close all; clearvars;

%% ---------------- User settings ----------------
dataInput   = fullfile('..','..','data');  % can be a folder OR a single CSV file
filePattern = 'visualsearch_sub-*_*.csv';  % used if dataInput is a folder

rtUnit      = 's';      % 's' if your task saved seconds; 'ms' if already ms
trimMethod  = 'iqr';    % 'none' | 'iqr' | 'range'
iqrK        = 3;        % RT trimming: keep within [Q1-k*IQR, Q3+k*IQR] per condition
rtRange_ms  = [150 3000]; % only used if trimMethod='range'

doBootstrap = true;
nBoot       = 2000;     % bootstrap samples per condition
rng(1);                 % reproducibility for CIs

saveFig     = true;
outDir      = fullfile('..','..','images');
outBaseName = 'visualsearch_medianRT_by_setsize';

%% ---------------- Load data ----------------
T = readVisualSearchData(dataInput, filePattern);

% ---- Standardize expected variable names ----
% (If your CSV uses different names, add aliases here.)
T = standardizeVars(T);

% ---- RT units ----
RT = T.rt;
switch lower(rtUnit)
    case 's'
        RT = 1000 * RT; % convert to ms
    case 'ms'
        % ok
    otherwise
        error('rtUnit must be ''s'' or ''ms''.');
end

% Variables
respKey       = T.respKey;
targetPresent = logical(T.targetPresent);
setSize       = T.setSize;
searchType    = parseSearchType(T.searchType); % returns categorical with labels

% ---- Correctness ----
% Expectation: respKey is 0/1 for absent/present. If stored as char keys, handle that.
resp = parseResponse(respKey);
isCorrect = (resp == targetPresent);

fprintf('Accuracy: %d/%d correct (%.1f%% errors)\n', sum(isCorrect), numel(isCorrect), 100*mean(~isCorrect));

% ---- Keep only correct trials ----
Tcorr = table(RT(isCorrect), setSize(isCorrect), searchType(isCorrect), ...
    'VariableNames', {'rt_ms','setSize','searchType'});

% ---- Optional RT trimming (on correct trials) ----
Tcorr = trimRT(Tcorr, trimMethod, iqrK, rtRange_ms);

%% ---------------- Aggregate: median per condition ----------------
% Condition = searchType x setSize
[G, st, ss] = findgroups(Tcorr.searchType, Tcorr.setSize);
medRT = splitapply(@median, Tcorr.rt_ms, G);
nTrials = splitapply(@numel, Tcorr.rt_ms, G);

summary = table(st, ss, nTrials, medRT, 'VariableNames', ...
    {'searchType','setSize','n','medianRT_ms'});

% Sort nicely for plotting
summary = sortrows(summary, {'searchType','setSize'});

%% ---------------- Bootstrap CIs (optional) ----------------
if doBootstrap
    ciLo = nan(height(summary),1);
    ciHi = nan(height(summary),1);

    for i = 1:height(summary)
        sel = (Tcorr.searchType==summary.searchType(i)) & (Tcorr.setSize==summary.setSize(i));
        x = Tcorr.rt_ms(sel);

        if numel(x) < 5
            ciLo(i) = NaN; ciHi(i) = NaN;
            continue
        end

        bootMed = bootstrp(nBoot, @median, x);
        ci = prctile(bootMed, [2.5 97.5]);
        ciLo(i) = ci(1);
        ciHi(i) = ci(2);
    end

    summary.ciLo = ciLo;
    summary.ciHi = ciHi;
end

%% ---------------- Plot ----------------
figure(1); clf; hold on;

cats = categories(summary.searchType);

h = gobjects(numel(cats),1);
for k = 1:numel(cats)
    sel = summary.searchType == cats{k};
    x = summary.setSize(sel);
    y = summary.medianRT_ms(sel);

    % plot line + markers
    h(k) = plot(x, y, 'o-', 'LineWidth', 2, 'MarkerSize', 9, 'MarkerFaceColor', 'w');

    % error bars (bootstrap CI)
    if doBootstrap && all(ismember({'ciLo','ciHi'}, summary.Properties.VariableNames))
        lo = summary.ciLo(sel);
        hi = summary.ciHi(sel);
        % only plot where finite
        ok = isfinite(lo) & isfinite(hi);
        if any(ok)
            errorbar(x(ok), y(ok), y(ok)-lo(ok), hi(ok)-y(ok), ...
                'LineStyle','none', 'LineWidth', 1);
        end
    end
end

grid on;
box off;

xlabel('Set size (items)');
ylabel('Median reaction time (ms)');
title('Visual search: Pop-out vs Conjunction');

xticks(sort(unique(summary.setSize)));
xlim([min(summary.setSize)-1, max(summary.setSize)+1]);

legend(h, cats, 'Location','northwest');

%% ---------------- Save outputs ----------------
if saveFig
    if ~exist(outDir,'dir'); mkdir(outDir); end
    exportgraphics(gcf, fullfile(outDir, [outBaseName '.png']), 'Resolution', 300);
    exportgraphics(gcf, fullfile(outDir, [outBaseName '.pdf']));
end

% Also save summary table next to figure (helps students)
writetable(summary, fullfile(outDir, [outBaseName '_summary.csv']));

disp(summary);

%% =================================================================
%% Local functions
function T = readVisualSearchData(dataInput, pattern)
    if isfolder(dataInput)
        files = dir(fullfile(dataInput, pattern));
        if isempty(files)
            error('No files found in folder "%s" with pattern "%s".', dataInput, pattern);
        end
        Tall = cell(numel(files),1);
        for i = 1:numel(files)
            f = fullfile(files(i).folder, files(i).name);
            Tall{i} = readtable(f);
            Tall{i}.sourceFile = repmat(string(files(i).name), height(Tall{i}), 1);
        end
        T = vertcat(Tall{:});
    else
        T = readtable(dataInput);
        T.sourceFile = repmat(string(dataInput), height(T), 1);
    end
end

function T = standardizeVars(T)
    % Make variable name matching more forgiving (case-insensitive)
    vn = lower(string(T.Properties.VariableNames));

    % Helper to rename if present
    function renameIf(oldNames, newName)
        idx = find(ismember(vn, lower(string(oldNames))), 1);
        if ~isempty(idx)
            T.Properties.VariableNames{idx} = newName;
            vn(idx) = lower(newName);
        end
    end

    renameIf({'rt','reactiontime','reaction_time'}, 'rt');
    renameIf({'respkey','response','resp'}, 'respKey');
    renameIf({'targetpresent','target_present','present'}, 'targetPresent');
    renameIf({'setsize','set_size','nitems','n'}, 'setSize');
    renameIf({'searchtype','condition','mode'}, 'searchType');

    required = {'rt','respKey','targetPresent','setSize','searchType'};
    missing = required(~ismember(required, string(T.Properties.VariableNames)));
    if ~isempty(missing)
        error('Missing required columns: %s', strjoin(missing, ', '));
    end
end

function st = parseSearchType(x)
    % Returns categorical with labels {'pop-out','conjunction'} when possible.
    if iscategorical(x)
        st = x;
    elseif isstring(x) || ischar(x) || iscellstr(x)
        st = categorical(string(x));
    else
        % numeric: assume 0=pop-out, 1=conjunction (Wallisch-style)
        st = categorical(x, [0 1], {'pop-out','conjunction'});
        return
    end

    % Normalize common spellings
    s = lower(string(st));
    s = replace(s, "popout", "pop-out");
    s = replace(s, "pop out", "pop-out");
    s = replace(s, "conj", "conjunction");
    s = replace(s, "conjunctive", "conjunction");
    st = categorical(s);

    % If it already came as '0'/'1' strings, map them:
    if all(ismember(categories(st), ["0","1"]))
        st = categorical(double(st)-double("0"), [0 1], {'pop-out','conjunction'});
    end
end

function resp = parseResponse(respKey)
    % Convert response keys to logical/0-1.
    % If numeric already (0/1), keep.
    if isnumeric(respKey)
        resp = logical(respKey);
        return
    end
    rk = string(respKey);
    rk = strip(lower(rk));
    rk = replace(rk, "numpad", ""); % tolerate "numpad1" etc if present
    % Accept '1'/'0', 'y'/'n', 'present'/'absent'
    resp = nan(numel(rk),1);
    resp(rk=="1" | rk=="y" | rk=="yes" | rk=="present") = 1;
    resp(rk=="0" | rk=="n" | rk=="no"  | rk=="absent")  = 0;

    if any(isnan(resp))
        warning('Some responses could not be parsed; marking as incorrect.');
        resp(isnan(resp)) = -999; % will fail correctness check
    end
    resp = logical(resp);
end

function T = trimRT(T, method, iqrK, rtRange_ms)
    method = lower(method);
    switch method
        case 'none'
            return
        case 'range'
            sel = T.rt_ms >= rtRange_ms(1) & T.rt_ms <= rtRange_ms(2);
            T = T(sel,:);
        case 'iqr'
            % Trim within each condition (searchType x setSize)
            [G,~,~] = findgroups(T.searchType, T.setSize);
            keep = false(height(T),1);

            for g = 1:max(G)
                idx = (G==g);
                x = T.rt_ms(idx);
                if numel(x) < 8
                    keep(idx) = true; % don't trim tiny cells
                    continue
                end
                q1 = prctile(x,25);
                q3 = prctile(x,75);
                iqr = q3-q1;
                lo = q1 - iqrK*iqr;
                hi = q3 + iqrK*iqr;
                keep(idx) = (x>=lo & x<=hi);
            end
            T = T(keep,:);
        otherwise
            error('Unknown trimMethod "%s".', method);
    end
end