function R = tdcs_models(S, meta, cfg)
%TDCS_MODELS Fit the tDCS GLMMs, per-animal tests, and (unmerged) classification.
%   R = TDCS_MODELS(S, META, CFG) consumes one scenario table S and its
%   META (from TDCS_SCENARIO_DATA) plus CFG (from TDCS_CONFIG), and
%   returns a struct R:
%
%     R.countGLME   GeneralizedLinearMixedModel, Poisson count-level model:
%                     success ~ group + day_c + day_c2 + (1|subject)
%     R.rateGLME    GeneralizedLinearMixedModel, Binomial rate-level model
%                   (total==0 sessions dropped first):
%                     success ~ group + day_c + day_c2 + (1|subject)
%     R.learnGLME   GeneralizedLinearMixedModel, Poisson learning-rate
%                   model with a group x day_c interaction:
%                     success ~ group*day_c + day_c2 + (1|subject)
%     R.learn       learning-rate comparison, Box-B2 vs Box-A2. The reported
%                   result is a PER-ANIMAL slope test (cluster-honest, like
%                   R.perAnimal): each subject's OLS slope of success on day,
%                   then Welch t / two-sided Mann-Whitney across groups:
%                   .slopeWelchP, .slopeMWUp, .nB, .nA, .meanSlopeB,
%                   .meanSlopeA. The GLMM group x day_c joint F-test is also
%                   kept for reference (.glmeJointP, .glmeJointF, .glmeJointDF1,
%                   .glmeJointDF2) but is ANTICONSERVATIVE here: fitglme's
%                   coefTest offers only observation-level DF, which overstates
%                   significance for this few-subject design (it reads p~=0
%                   even when the per-animal slopes are clearly parallel), so
%                   it is NOT the basis for the conclusion.
%     R.perAnimal   per-subject Box-B2 vs Box-A2 comparison (pure stats,
%                   no GLME involved): .countWelchP, .countMWUp,
%                   .rateWelchP, .rateMWUp, .nB, .nA. "count" = per-subject
%                   mean success; "rate" = per-subject pooled
%                   sum(success)/sum(total). Welch = TTEST2(...,
%                   'Vartype','unequal') (two-sided); MWU = RANKSUM(...,
%                   'tail','right') (one-sided, B2 > A2).
%     R.h2          anchor check (Box-A2 vs Box-B2) read off the
%                   count/rate GLMEs' group_<anchorLow> fixed effect
%                   (coded vs the cfg.ref = anchorHigh reference level):
%                   .countIRR, .countOneSidedP, .rateOR, .rateOneSidedP.
%                   IRR/OR = 1/exp(coef) is the Box-B2-over-Box-A2
%                   advantage (so > 1 supports H2); one-sided p =
%                   NORMCDF(coef/se), testing coef < 0 (B2 above A2).
%     R.classify    (unmerged scenarios only, i.e. meta.isUnmerged) struct
%                   with one field per unknown in cfg.unknowns (field name
%                   via matlab.lang.makeValidName), each holding .count and
%                   .rate sub-structs with .vsA2 (.ratio, .p), .vsB2
%                   (.ratio, .p), and .label (one of 'B2-like', 'A2-like',
%                   'AMBIGUOUS (nearer <anchor>)', 'UNLIKE BOTH'), per
%                   METHODS.md's classification logic. For merged
%                   scenarios (meta.isUnmerged == false), R.classify is an
%                   empty (no-field) struct.
%
%   See analysis/tdcs_glm.py (the Python this ports) and
%   docs/superpowers/plans/2026-07-20-matlab-tdcs-analysis.md (Task 3).

countFormula = 'success ~ group + day_c + day_c2 + (1|subject)';
R.countGLME = fitglme(S, countFormula, 'Distribution', 'Poisson', 'Link', 'log');

Sr = S(S.total > 0, :);
R.rateGLME = fitglme(Sr, countFormula, 'Distribution', 'Binomial', ...
    'BinomialSize', Sr.total, 'Link', 'logit');

learnFormula = 'success ~ group*day_c + day_c2 + (1|subject)';
R.learnGLME = fitglme(S, learnFormula, 'Distribution', 'Poisson');

R.perAnimal = localPerAnimal(S, cfg);
R.h2 = localH2(R.countGLME, R.rateGLME, cfg);
R.learn = localLearn(R.learnGLME, S, cfg);

if meta.isUnmerged
    R.classify = localClassifyAll(R.countGLME, R.rateGLME, cfg);
else
    R.classify = struct();
end

end

% --------------------------------------------------------------- per-animal
function out = localPerAnimal(S, cfg)
%LOCALPERANIMAL Per-subject Box-B2 vs Box-A2 Welch t-test / MWU (pure stats).
b = S(S.group == cfg.anchorHigh, :);
a = S(S.group == cfg.anchorLow, :);

[bCount, bRate, nB] = localSubjectSummary(b);
[aCount, aRate, nA] = localSubjectSummary(a);

[~, countWelchP] = ttest2(bCount, aCount, 'Vartype', 'unequal');
countMWUp = ranksum(bCount, aCount, 'tail', 'right');
[~, rateWelchP] = ttest2(bRate, aRate, 'Vartype', 'unequal');
rateMWUp = ranksum(bRate, aRate, 'tail', 'right');

out = struct( ...
    'countWelchP', countWelchP, 'countMWUp', countMWUp, ...
    'rateWelchP', rateWelchP, 'rateMWUp', rateMWUp, ...
    'nB', nB, 'nA', nA);
end

function [meanSuccess, pooledRate, nSubj] = localSubjectSummary(T)
%LOCALSUBJECTSUMMARY Per-subject mean success and pooled success/total.
subj = unique(T.subject);
nSubj = numel(subj);
meanSuccess = zeros(nSubj, 1);
pooledRate = zeros(nSubj, 1);
for i = 1:nSubj
    rows = T.subject == subj(i);
    meanSuccess(i) = mean(T.success(rows));
    pooledRate(i) = sum(T.success(rows)) / sum(T.total(rows));
end
end

% --------------------------------------------------------------------- H2
function h2 = localH2(countGLME, rateGLME, cfg)
%LOCALH2 Anchor check: Box-A2 vs Box-B2, read off each GLME's group term.
[countCoef, countSE] = localCoefAndSE(countGLME, cfg.anchorLow, cfg.ref);
[rateCoef, rateSE] = localCoefAndSE(rateGLME, cfg.anchorLow, cfg.ref);

h2 = struct( ...
    'countIRR', 1 / exp(countCoef), ...
    'countOneSidedP', normcdf(countCoef / countSE), ...
    'rateOR', 1 / exp(rateCoef), ...
    'rateOneSidedP', normcdf(rateCoef / rateSE));
end

% -------------------------------------------------------------- classify
function out = localClassifyAll(countGLME, rateGLME, cfg)
%LOCALCLASSIFYALL Classify every unknown in cfg.unknowns vs both anchors.
out = struct();
for i = 1:numel(cfg.unknowns)
    unknown = cfg.unknowns{i};
    fieldName = matlab.lang.makeValidName(unknown);
    out.(fieldName) = struct( ...
        'count', localClassifyOne(countGLME, unknown, cfg), ...
        'rate', localClassifyOne(rateGLME, unknown, cfg));
end
end

function verdict = localClassifyOne(model, unknown, cfg)
%LOCALCLASSIFYONE Compare UNKNOWN vs both anchors on one GLME; label it.
%   Mirrors classify_one() in analysis/tdcs_glm.py / METHODS.md's
%   "Classification logic (unknowns)": indistinguishable (p >= 0.05) from
%   one anchor and different from the other => that anchor's label;
%   indistinguishable from both => AMBIGUOUS (report the numerically
%   nearer one); different from both => UNLIKE BOTH.
[coefLo, pLo] = localContrast(model, unknown, cfg.anchorLow, cfg.ref);
[coefHi, pHi] = localContrast(model, unknown, cfg.anchorHigh, cfg.ref);

vsA2 = struct('ratio', exp(coefLo), 'p', pLo);
vsB2 = struct('ratio', exp(coefHi), 'p', pHi);

likeLo = pLo >= 0.05;
likeHi = pHi >= 0.05;
if likeHi && ~likeLo
    label = 'B2-like';
elseif likeLo && ~likeHi
    label = 'A2-like';
elseif likeLo && likeHi
    if abs(coefHi) < abs(coefLo)
        nearer = cfg.anchorHigh;
    else
        nearer = cfg.anchorLow;
    end
    label = sprintf('AMBIGUOUS (nearer %s)', nearer);
else
    label = 'UNLIKE BOTH';
end

verdict = struct('vsA2', vsA2, 'vsB2', vsB2, 'label', label);
end

% ------------------------------------------------------------- contrasts
function [coef, se] = localCoefAndSE(model, aName, refName)
%LOCALCOEFANDSE Coefficient + SE for group level A_NAME vs the reference.
v = localContrastVector(model, aName, refName, refName);
beta = model.Coefficients.Estimate;
Cov = model.CoefficientCovariance;
coef = v * beta;
se = sqrt(v * Cov * v');
end

function [coef, p] = localContrast(model, aName, bName, refName)
%LOCALCONTRAST Linear contrast A_NAME vs B_NAME (both coded vs REF_NAME).
%   Two-sided p via COEFTEST (an F-test on one row = the two-sided
%   Wald/t-test); works for any pair, including B_NAME == REF_NAME.
v = localContrastVector(model, aName, bName, refName);
beta = model.Coefficients.Estimate;
coef = v * beta;
p = coefTest(model, v);
end

function v = localContrastVector(model, aName, bName, refName)
%LOCALCONTRASTVECTOR +1 at group_A_NAME (if not the reference), -1 at
%   group_B_NAME (if not the reference), over MODEL.CoefficientNames.
cn = model.CoefficientNames;
v = zeros(1, numel(cn));
if ~strcmp(aName, refName)
    v(localCoefIndex(cn, aName)) = v(localCoefIndex(cn, aName)) + 1;
end
if ~strcmp(bName, refName)
    v(localCoefIndex(cn, bName)) = v(localCoefIndex(cn, bName)) - 1;
end
end

function idx = localCoefIndex(coefNames, groupName)
%LOCALCOEFINDEX Index into CoefficientNames for a "group_<groupName>" term.
idx = find(strcmp(coefNames, ['group_' groupName]), 1);
if isempty(idx)
    error('tdcs_models:missingGroupTerm', ...
        'No "group_%s" coefficient in this GLME (CoefficientNames: %s).', ...
        groupName, strjoin(coefNames, ', '));
end
end

% -------------------------------------------------------------- learning
function out = localLearn(learnGLME, S, cfg)
%LOCALLEARN Learning-rate result: per-animal slope test (primary, cluster-
%   honest) plus the GLMM interaction joint F-test (reference, anticonservative).
bSlopes = localSubjectSlopes(S(S.group == cfg.anchorHigh, :));
aSlopes = localSubjectSlopes(S(S.group == cfg.anchorLow, :));
[~, slopeWelchP] = ttest2(bSlopes, aSlopes, 'Vartype', 'unequal');
slopeMWUp = ranksum(bSlopes, aSlopes);   % two-sided: slope-diff direction unknown

j = localJointInteractionTest(learnGLME);
out = struct( ...
    'slopeWelchP', slopeWelchP, 'slopeMWUp', slopeMWUp, ...
    'nB', numel(bSlopes), 'nA', numel(aSlopes), ...
    'meanSlopeB', mean(bSlopes), 'meanSlopeA', mean(aSlopes), ...
    'glmeJointP', j.jointP, 'glmeJointF', j.jointF, ...
    'glmeJointDF1', j.jointDF1, 'glmeJointDF2', j.jointDF2);
end

function slopes = localSubjectSlopes(T)
%LOCALSUBJECTSLOPES Per-subject OLS slope of success on day.
subj = unique(T.subject);
slopes = zeros(numel(subj), 1);
for i = 1:numel(subj)
    rows = T.subject == subj(i);
    coeffs = polyfit(T.day(rows), T.success(rows), 1);
    slopes(i) = coeffs(1);
end
end

function out = localJointInteractionTest(learnGLME)
%LOCALJOINTINTERACTIONTEST Joint Wald test that all group x day_c
%   interaction fixed effects are zero (do groups learn at different
%   rates?), mirroring the GEE wald_test in analysis/tdcs_glm.py.
cn = learnGLME.CoefficientNames;
isInteraction = contains(cn, ':day_c') & ~contains(cn, 'day_c2');
idx = find(isInteraction);
H = zeros(numel(idx), numel(cn));
for i = 1:numel(idx)
    H(i, idx(i)) = 1;
end
[p, F, df1, df2] = coefTest(learnGLME, H);
out = struct('jointP', p, 'jointF', F, 'jointDF1', df1, 'jointDF2', df2);
end