function [params, Tfit, rep] = trans_mag_pickfit(TH, nstages, varargin)
% TRANS_MAG_PICKFIT (robust)
% Magnitude-only modal pick with defensive peak handling.
%
% Returns:
%   params = [fn  zeta  A]
%   Tfit   = | sum_i A_i /((1 - r_i^2) + j*2*zeta_i*r_i) |
%   rep    = struct with RMSE, seeds, Qtable, StageLog

    % ---------- Parse inputs ----------
    p = inputParser;
    p.FunctionName    = 'trans_mag_pickfit';
    p.CaseSensitive   = false;
    p.PartialMatching = true;
    p.KeepUnmatched   = true;

    addRequired(p,'TH',      @(x)isnumeric(x)&&size(x,2)>=2);
    addRequired(p,'nstages', @(x)isnumeric(x)&&isscalar(x)&&x>=1);

    addParameter(p,'NumModes',           [],     @(x)isnumeric(x)&&isscalar(x)&&x>=1);
    addParameter(p,'DoParabolicPolish',  true,   @(x)islogical(x)||isnumeric(x));
    addParameter(p,'DoHalfPowerTouchup', true,   @(x)islogical(x)||isnumeric(x));
    addParameter(p,'DoSKPolish',         false,  @(x)islogical(x)||isnumeric(x)); %#ok<INUSD>
    addParameter(p,'Plot',               false,  @(x)islogical(x)||isnumeric(x));

    parse(p, TH, nstages, varargin{:});
    opt = p.Results;

    if isempty(opt.NumModes)
        nm_req = nstages;
    else
        nm_req = opt.NumModes;
    end

    % ---------- Prep ----------
    f    = TH(:,1);
    Tmag = TH(:,2);
    [f, Tmag] = sanitize(f, Tmag);

    % ---------- Seed peaks ----------
    [fpks, apks, ipks] = seed_peaks(f, Tmag, nm_req);

    nm_found = numel(ipks);
    if nm_found == 0
        error('trans_mag_pickfit:NoPeaks','No valid peaks found.');
    end

    if nm_found < nm_req
        nm = nm_found;
        StageLog = {sprintf('Seeds found: %d (requested %d) — reduced.', nm_found, nm_req)};
    else
        nm = nm_req;
        StageLog = {sprintf('Seeds found: %d', nm_found)};
    end

    fpks = fpks(1:nm);
    apks = apks(1:nm);
    ipks = ipks(1:nm);

    % ---------- Parabolic peak polish ----------
    if opt.DoParabolicPolish
        for k = 1:nm
            [fpk, Apk] = parab_peak(f, Tmag, ipks(k));
            if isfinite(fpk) && fpk > f(1) && fpk < f(end)
                fpks(k) = fpk;
                apks(k) = Apk;
            end
        end
        StageLog{end+1} = 'Parabolic peak polish applied.';
    else
        StageLog{end+1} = 'Parabolic peak polish skipped.';
    end

    % ---------- Half-power damping ----------
    z = zeros(nm,1);
    for k = 1:nm
        [~, idx] = min(abs(f - fpks(k)));
        z(k) = halfpower_zeta(f, Tmag, idx);
        if ~isfinite(z(k)) || z(k) <= 0
            z(k) = 0.02;   % conservative fallback
        end
    end

    if opt.DoHalfPowerTouchup
        StageLog{end+1} = 'Half-power damping computed.';
    else
        StageLog{end+1} = 'Half-power damping defaulted.';
    end

    % ---------- Amplitude estimate ----------
    A = max(2*z.*apks(:), 1e-12);

    params = [fpks(:), z(:), A(:)];

    % ---------- Outputs ----------
    Tfit = modal_sum_mag(params, f);

    rep = struct();
    rep.RMSE     = sqrt(mean((Tfit - Tmag).^2,'omitnan'));
    rep.seeds    = struct('fpks',fpks,'apks',apks,'ipks',ipks);
    rep.StageLog = StageLog;
    rep.Qtable   = table(params(:,1), 1./(2*params(:,2)), ...
                         'VariableNames', {'fn_Hz','Q'});

    if logical(opt.Plot)
        figure('Name','trans\_mag\_pickfit QA','Color','w');
        plot(f, Tmag, 'k-', 'LineWidth', 1); hold on;
        plot(f, Tfit, 'r--', 'LineWidth', 1);
        for i = 1:nm
            plot(f, abs(modal_term_c(params(i,1), params(i,2), params(i,3), f)), ':');
        end
        grid on; box on;
        xlabel('Frequency (Hz)');
        ylabel('|T| (linear)');
        title(sprintf('Modal pickfit (RMSE = %.3g)', rep.RMSE));
        hold off;
    end
end

% ======================= Helpers =======================

function [f, Tmag] = sanitize(f, Tmag)
    good = isfinite(f) & isfinite(Tmag) & Tmag > 0;
    f = f(good);
    Tmag = Tmag(good);
    [f, ord] = sort(f);
    Tmag = Tmag(ord);
end

function [fpks, apks, ipks] = seed_peaks(f, Tmag, nm)
    y  = Tmag(:);
    ys = movmedian(y, max(3, round(numel(f)/400)), 'omitnan');

    isMax = islocalmax(ys);
    cand  = find(isMax);
    cand  = cand(:);

    if numel(cand) < nm
        [~,cand] = maxk(y, min(nm, numel(y)));
        cand = cand(:);
    end

    xlog = log10(f(:));
    minSep = 0.03;

    chosen = zeros(0,1);   % FORCE column
    for ii = 1:numel(cand)
        i = cand(ii);
        if isempty(chosen) || min(abs(xlog(chosen) - xlog(i))) > minSep
            chosen(end+1,1) = i; %#ok<AGROW>
            if numel(chosen) == nm
                break;
            end
        end
    end

    if numel(chosen) < nm
        remain = setdiff(cand, chosen, 'stable');
        k = min(nm - numel(chosen), numel(remain));
        if k > 0
            chosen = [chosen; remain(1:k)];
        end
    end

    ipks = sort(chosen);
    apks = y(ipks);
    fpks = f(ipks);
end

function [f_pk, T_pk] = parab_peak(f, Tmag, k)
    k = max(2, min(numel(f)-1, k));
    ii = k-1:k+1;
    x = f(ii);
    y = Tmag(ii);

    p = polyfit(x, y, 2);
    if p(1) >= 0
        [T_pk, j] = max(y);
        f_pk = x(j);
    else
        f_pk = -p(2)/(2*p(1));
        T_pk = polyval(p, f_pk);
    end
end

function z = halfpower_zeta(f, Tmag, k)
    k = max(2, min(numel(f)-1, k));
    Tpk = Tmag(k);
    if ~isfinite(Tpk) || Tpk <= 0
        z = NaN;
        return;
    end

    Thp = Tpk/sqrt(2);
    iL = find(Tmag(1:k) <= Thp, 1, 'last');
    iR = find(Tmag(k:end) <= Thp, 1, 'first');

    if isempty(iL) || isempty(iR)
        z = NaN;
        return;
    end

    iR = iR + (k-1);
    fL = interp1(Tmag(iL:iL+1), f(iL:iL+1), Thp, 'linear', 'extrap');
    fR = interp1(Tmag(iR-1:iR), f(iR-1:iR), Thp, 'linear', 'extrap');

    if ~(isfinite(fL) && isfinite(fR) && fR > fL)
        z = NaN;
        return;
    end

    fn_hat = f(k);
    z = (fR - fL) / (2*fn_hat);
end

function H = modal_term_c(fn, zeta, A, f)
    r = f ./ fn;
    H = A ./ ((1 - r.^2) + 1j*(2*zeta.*r));
end

function Tfit = modal_sum_mag(params, f)
    H = zeros(size(f));
    for i = 1:size(params,1)
        H = H + modal_term_c(params(i,1), params(i,2), params(i,3), f);
    end
    Tfit = abs(H);
end
