function [fn, zeta, A, Hhat, info] = mdof_frf_main(THM, ndof, iat, imiss, ibh, varargin)

% ================= OPTIONS =================
p = inputParser;
p.addParameter('Plot', true);
p.addParameter('Verbose', true);
p.addParameter('PrimaryChan', 1);
p.addParameter('WinFrac', 0.12);
p.addParameter('WeightFloorFrac', 1e-3);
p.addParameter('EstimateDelay', false);
p.addParameter('UseLM', true);
p.addParameter('LM_MaxIters', 100);
p.addParameter('LM_Lambda0', 1e-2);
p.parse(varargin{:});
opt = p.Results;

% ================= FRF TYPE =================
switch iat
    case 1, frf='';           model='multiplicative';
    case 2, frf='acc/force';  model='additive';
    case 3, frf='vel/force';  model='additive';
    case 4, frf='disp/force'; model='additive';
    otherwise, error('Unknown iat');
end

% ================= INPUT PARSE =================
[f, Hc] = parse_input_any(THM, opt.PrimaryChan);
f = f(:); Hc = Hc(:);
w = 2*pi*f;

if opt.Verbose
    fprintf('\n=== mdof_frf_main: ndof=%d  iat=%d  f=[%.3f..%.3f] Hz  N=%d ===\n',...
        ndof, iat, min(f), max(f), numel(f));
end

% ================= WEIGHTS =================
mag_floor = max(abs(Hc))*opt.WeightFloorFrac;
wts = 1 ./ max(abs(Hc), mag_floor);
wts = wts / max(wts);

% ================= IMPROVED PEAK SELECTION =================
[fn_seed, z_seed, mask] = robust_peak_selection(f, Hc, ndof, opt.WinFrac);

% ================= LINEAR COEFFICIENT SOLVE =================
tau0 = 0;
[A0, H0] = solve_coeffs(f,Hc,fn_seed,z_seed,tau0,frf,model,wts,mask);

% ================= LM REFINEMENT =================
fn_best = fn_seed; z_best = z_seed; tau_best = tau0;
A_best = A0; H_best = H0;

if opt.UseLM
    [fn_best,z_best,tau_best,A_best,H_best] = ...
        lm_refine(f,Hc,fn_seed,z_seed,tau0,frf,model,wts,mask,opt);
end

% ================= FINAL SOLVE =================
[fn,ord] = sort(fn_best);
zeta = z_best(ord);

[A, Hhat] = solve_coeffs(f,Hc,fn,zeta,tau_best,frf,model,wts,mask);

% ================= INFO =================
info.rmse_window = sqrt(mean(abs(Hc(mask)-Hhat(mask)).^2));
info.rmse_full   = sqrt(mean(abs(Hc-Hhat).^2));
info.tau = tau_best;

if opt.Verbose
    fprintf('Final window RMSE = %.6g   full RMSE = %.6g   tau=%.6g\n',...
        info.rmse_window, info.rmse_full, info.tau);
end

if opt.Plot
    plot_ri(f,Hc,Hhat,frf);
end

end

%% ========================================================================
% ROBUST PEAK SELECTION (guarantees weak modes are preserved)
%% ========================================================================
function [fn_seed,z_seed,mask] = robust_peak_selection(f,Hc,ndof,winFrac)

% ---- sort and remove duplicates safely ----
[f, idx] = sort(f(:));
Hc = Hc(idx);

[fu, iu] = unique(f,'stable');
f = fu;
Hc = Hc(iu);

N = numel(f);

% ---- peak metric ----
metric = abs(imag(Hc));
metric = smoothdata(metric,'movmean',max(5,round(0.002*N)));

% ---- find peaks WITHOUT passing frequency to findpeaks ----
[pks,locIdx] = findpeaks(metric,'MinPeakProminence',0.02*max(metric));

if isempty(locIdx)
    % fallback evenly spaced
    fn_seed = linspace(min(f),max(f),ndof).';
else
    locFreq = f(locIdx);

    % score by local energy
    scores = zeros(size(locFreq));
    band = max(f)-min(f);
    for k=1:numel(locFreq)
        m = abs(f-locFreq(k)) < 0.03*band;
        scores(k) = trapz(f(m),metric(m).^2);
    end

    [~,ord] = sort(scores,'descend');
    locFreq = locFreq(ord);

    fn_seed = [];
    minSep = 0.02*(max(f)-min(f));

    for k=1:numel(locFreq)
        if isempty(fn_seed) || all(abs(locFreq(k)-fn_seed) > minSep)
            fn_seed(end+1)=locFreq(k); %#ok<AGROW>
        end
        if numel(fn_seed)>=ndof
            break
        end
    end

    fn_seed = sort(fn_seed(:));
end

% ---- half-power damping ----
z_seed = zeros(numel(fn_seed),1);

for k=1:numel(fn_seed)
    [~,ik]=min(abs(f-fn_seed(k)));
    pk=metric(ik);
    hhalf=0.5*pk;

    iL=ik;
    while iL>1 && metric(iL)>hhalf
        iL=iL-1;
    end

    iR=ik;
    while iR<N && metric(iR)>hhalf
        iR=iR+1;
    end

    bw=f(iR)-f(iL);
    z_seed(k)=max(1e-4,min(0.2,bw/(2*fn_seed(k))));
end

% ---- force exactly ndof ----
if numel(fn_seed)<ndof
    extra = linspace(min(f),max(f),ndof);
    fn_seed = extra(:);
    z_seed = 0.02*ones(ndof,1);
end

% ---- build mask around ALL peaks ----
mask=false(size(f));
for k=1:numel(fn_seed)
    mask = mask | (f>fn_seed(k)*(1-winFrac) & f<fn_seed(k)*(1+winFrac));
end

if ~any(mask)
    mask(:)=true;
end

end



%% ========================================================================
% FRF KERNEL
%% ========================================================================
function Hk = frf_kernel(w,wn,z,frf)
den = (wn^2 - w.^2) + 1i*2*z*wn.*w;
switch frf
    case 'disp/force', Hk = 1./den;
    case 'vel/force',  Hk = (1i*w)./den;
    case 'acc/force',  Hk = (-w.^2)./den;
    otherwise,         Hk = (w.^2)./den;
end
end

%% ========================================================================
% LINEAR SOLVE
%% ========================================================================
function [A,Hhat] = solve_coeffs(f,Hc,fn,z,tau,frf,model,wts,mask)

w=2*pi*f;
Nm=numel(fn);

Phi=zeros(numel(f),Nm);
for k=1:Nm
    Phi(:,k)=frf_kernel(w,2*pi*fn(k),z(k),frf);
end

rhs = Hc;
Mf = Phi(mask,:);
bf = rhs(mask);

W = wts(mask);
M = [real(Mf).*W; imag(Mf).*W];
b = [real(bf).*W; imag(bf).*W];

A = M\b;

Hhat = Phi*A;

end

%% ========================================================================
% SIMPLE LM
%% ========================================================================
function [fn,z,tau,A,Hhat] = lm_refine(f,Hc,fn,z,tau,frf,model,wts,mask,opt)

for iter=1:opt.LM_MaxIters
    [A,Hhat]=solve_coeffs(f,Hc,fn,z,tau,frf,model,wts,mask);
    err = Hhat(mask)-Hc(mask);
    base = [real(err); imag(err)];

    J=zeros(numel(base),2*numel(fn));
    df=1e-4*fn;
    dz=1e-4*z;

    for k=1:numel(fn)
        fnp=fn; fnp(k)=fn(k)+df(k);
        [~,Hp]=solve_coeffs(f,Hc,fnp,z,tau,frf,model,wts,mask);
        ep=[real(Hp(mask)-Hc(mask)); imag(Hp(mask)-Hc(mask))];
        J(:,k)=(ep-base)/df(k);

        zp=z; zp(k)=z(k)+dz(k);
        [~,Hp]=solve_coeffs(f,Hc,fn,zp,tau,frf,model,wts,mask);
        ep=[real(Hp(mask)-Hc(mask)); imag(Hp(mask)-Hc(mask))];
        J(:,numel(fn)+k)=(ep-base)/dz(k);
    end

    dp = -(J'*J + opt.LM_Lambda0*eye(size(J,2))) \ (J'*base);

    fn = fn + dp(1:numel(fn));
    z  = min(0.2,max(1e-4,z+dp(numel(fn)+1:end)));

    if norm(dp)<1e-8
        break
    end
end

[A,Hhat]=solve_coeffs(f,Hc,fn,z,tau,frf,model,wts,mask);

end

%% ========================================================================
function plot_ri(f,Hc,Hhat,frf)
figure;
subplot(2,1,1)
plot(f,real(Hc),'k'); hold on;
plot(f,real(Hhat),'r');
grid on; ylabel('Real'); legend('Input','Fit');

subplot(2,1,2)
plot(f,imag(Hc),'k'); hold on;
plot(f,imag(Hhat),'r');
grid on; ylabel('Imag'); xlabel('Hz');
end

%% ========================================================================
function [f,Hc]=parse_input_any(THM,primary)
f=THM(:,1);
if size(THM,2)==3
    Hc=complex(THM(:,2),THM(:,3));
else
    Hc=THM(:,2);
end
end