% order_tracking_function.m
% ver 2.0  by Tom Irvine
%
% General-purpose order tracking for automotive & rotating machinery
% Includes:
%   - Time histories
%   - Frequency map
%   - Computed order map (frequency → order)
%   - TRUE order map (angular resampling)
%
% Physics controlled by winLen
% Frequency grid controlled by nfft

function order_tracking_function(t,vib,rpm,winLen,nfft)

%% ================= INPUT CHECKS =================
if nargin < 5
    error('Usage: order_tracking_function(t,vib,rpm,winLen,nfft)')
end

t   = t(:);
vib = vib(:);
rpm = rpm(:);

%% ================= BASIC SETTINGS ===============
dt = mean(diff(t));
fs = 1/dt;

shaft_hz = max(rpm/60,1e-3);    % prevent divide-by-zero

%% ================= USER PARAMETERS ==============
% Spectrogram
win      = hann(winLen,'periodic');
noverlap = round(0.75*winLen);

% Display limits
maxOrderSignal = 3;      % highest physical order expected
orderMax       = 6;      % orders shown
freqMargin     = 2.0;

% TRUE order tracking (speed/accuracy trade)
angleStep_deg = 3;       % angular resampling step
revWindow     = 10;      % revolutions per FFT
revStep       = 3;       % FFT spacing (revolutions)

rng(1)  % repeatability if noise exists

%% ================= FIGURE 1: TIME HISTORIES =====
figure(1)
subplot(2,1,1)
plot(t,rpm,'r','LineWidth',1.5)
grid on
xlabel('Time (s)')
ylabel('RPM')
title('Rotational Speed')

subplot(2,1,2)
plot(t,vib,'b')
grid on
xlabel('Time (sec)')
ylabel('Accel (G)')
title('Vibration Signal')

%% ================= SPECTROGRAM ==================
[S,F,Tspec] = spectrogram(vib,win,noverlap,nfft,fs);
S_rms = abs(S)/sqrt(2);

%% ================= FREQUENCY MAP ================
f_max = ceil(maxOrderSignal*max(rpm)/60*freqMargin/10)*10;
maskF = F <= f_max;

figure(2)
imagesc(Tspec,F(maskF),20*log10(S_rms(maskF,:)))
axis xy
ylim([0 f_max])
xlabel('Time (s)')
ylabel('Frequency (Hz)')
title('Frequency Map')
colormap turbo
cb = colorbar;
ylabel(cb,'RMS Amplitude (dB)')

%% ================= COMPUTED ORDER MAP ===========
rpm_i = max(interp1(t,shaft_hz,Tspec,'linear','extrap'),1e-3);
order_matrix = F ./ rpm_i;

order_vec = linspace(0,orderMax,size(S,1));
S_order = zeros(numel(order_vec),numel(Tspec),'like',S_rms);

for k = 1:numel(Tspec)
    ok = order_matrix(:,k);
    [ok,ix] = sort(ok);
    S_order(:,k) = interp1(ok,S_rms(ix,k),order_vec,'linear',0);
end

figure(3)
imagesc(Tspec,order_vec,20*log10(S_order))
axis xy
ylim([0 orderMax])
xlabel('Time (s)')
ylabel('Order')
title('Computed Order Map')
colormap turbo
cb = colorbar;
ylabel(cb,'RMS Amplitude (dB)')
hold on
for n = 1:orderMax
    plot(Tspec,n*ones(size(Tspec)),'k--','LineWidth',0.5)
end
hold off

%% ================= TRUE ORDER TRACKING ==========
% Angular conversion
omega = 2*pi*shaft_hz;
theta = cumtrapz(t,omega)/(2*pi);     % revolutions

rev_step = angleStep_deg/360;
theta_u = theta(1):rev_step:theta(end);
vib_u   = interp1(theta,vib,theta_u,'linear','extrap');
t_u     = interp1(theta,t,theta_u,'linear','extrap');

win_len  = round(revWindow/rev_step);
step_len = round(revStep/rev_step);
idx0     = 1:step_len:(numel(theta_u)-win_len);

nfft_order = 2^nextpow2(win_len);
order_axis = (0:nfft_order/2-1)'/revWindow;

S_true = zeros(numel(order_axis),numel(idx0));
t_true = zeros(1,numel(idx0));

win_o = hann(win_len,'periodic');
U     = rms(win_o);

for k = 1:numel(idx0)
    idx = idx0(k):(idx0(k)+win_len-1);

    seg = vib_u(idx);
    seg = seg - mean(seg);
    seg = seg .* win_o;

    V = fft(seg,nfft_order);
    S_true(:,k) = abs(V(1:nfft_order/2)) / (sqrt(2)*U);

    t_true(k) = mean(t_u(idx));
end

%% ================= TRUE ORDER MAP ===============
maskO = order_axis <= orderMax;
S_true(S_true==0) = eps;
SdB = 20*log10(S_true(maskO,:));
Oplot = order_axis(maskO);

clim = [prctile(SdB(:),5) prctile(SdB(:),99)];

figure(4)
imagesc(t_true,Oplot,SdB)
axis xy
ylim([0 orderMax])
xlabel('Time (s)')
ylabel('Order')
title('TRUE Order Map (Order vs Time)')
caxis(clim)
colormap turbo
cb = colorbar;
ylabel(cb,'RMS Amplitude (dB)')
hold on
for n = 1:orderMax
    plot(t_true,n*ones(size(t_true)),'k--','LineWidth',0.5)
end
hold off

end
