% Driver for Chapter 2: the leapfrog scheme on a string and in a tube.
%
% The editable script supplies
%
%     [y, U] = leapfrog(u, uold, lam, M, bc, jr, K)
%
% which advances the scheme M steps from two given time levels.
%
% The driver supplies exact data at both starting levels, so that the only
% error left is the scheme's, and it measures that error against solutions it
% computes independently:
%
% mode 'run'    a plucked string at the audio rate.  The time step is one
%               sample, 1/fs, and the grid is N intervals, so the Courant number
%               is lam = 2 f0 N / fs and choosing the grid chooses lam.  The
%               reference is d'Alembert's solution at the receiver.
% mode 'cents'  every partial's frequency, measured.  The field is projected
%               onto the discrete sine vectors sin(n pi j / N), which are exact
%               eigenvectors of the three-point difference, so each projection
%               is a single sinusoid whose frequency wv_peakfreq reads off.
%               The exact frequency of partial n is n f0.
% mode 'click'  a short pulse circulating in a periodic tube, recorded at one
%               point, against the exact pulse train.
p = wv_params();
out = struct();

if strcmp(p.mode, 'run')
    fs = p.fs;
    f0 = p.f0;
    N = p.N;
    c = 2 * f0;                     % a string of length 1
    h = 1 / N;
    dt = 1 / fs;
    lam = c * dt / h;
    M = round(p.T * fs);
    x = (0:N)' * h;
    u0 = drv_phi(x, p.beta);
    um = 0.5 * (drv_phi(x + c * dt, p.beta) + drv_phi(x - c * dt, p.beta));
    if strcmp(p.where, 'bridge')
        jr = 2;                     % the first node off the bridge
    else
        jr = round(p.xp * N) + 1;
    end
    t0 = tic;
    [y, ~] = leapfrog(u0, um, lam, M, 'fixed', jr, 0);
    out.ms = toc(t0) * 1000;
    t = (1:M)' * dt;
    ye = 0.5 * (drv_phi(x(jr) - c * t, p.beta) + drv_phi(x(jr) + c * t, p.beta));

    sc = max(abs(ye));
    bad = find(~isfinite(y) | abs(y) > 50 * sc, 1);
    out.unstable = ~isempty(bad);
    if out.unstable
        out.blowup = bad * dt * 1000;
        y(bad:end) = 0;
    else
        out.blowup = NaN;
    end
    out.lam = lam;
    out.N = N;
    e = y - ye;
    % relative rms error, period by period
    P = fs / f0;
    np = floor(M / P);
    er = NaN(1, np);
    for k = 1:np
        i0 = floor((k - 1) * P) + 1;
        i1 = floor(k * P);
        er(k) = sqrt(sum(e(i0:i1) .^ 2) / sum(ye(i0:i1) .^ 2));
    end
    out.tp = ((1:np) - 0.5) / f0 * 1000;
    out.err = er;
    last = t > p.T - 1;
    out.relerr = sqrt(sum(e(last) .^ 2) / sum(ye(last) .^ 2));
    % three periods starting at one second, finely enough to see the drift
    w = find(t >= 1 & t < 1 + 3 / f0);
    out.tw = t(w)' * 1000;
    out.yw = y(w)';
    out.yew = ye(w)';
    % predicted pitch of partials 1 and 10, for the status strip
    out.pred1 = drv_pred(1, N, lam, dt, f0);
    out.pred10 = drv_pred(10, N, lam, dt, f0);
    out.clipExact = wv_audio(ye, fs, sc);
    out.clipFd = wv_audio(max(min(y, 2 * sc), -2 * sc), fs, sc);
    out.clipDiff = wv_audio(max(min(e, 2 * sc), -2 * sc), fs, sc);
elseif strcmp(p.mode, 'cents')
    f0 = p.f0;
    N = p.N;
    c = 2 * f0;
    h = 1 / N;
    x = (0:N)' * h;
    lams = p.lams(:)';
    nl = numel(lams);
    % the partials to measure: every one of the first ten, and twenty more
    % spread over the rest of the grid's range
    sel = unique([1:min(10, N - 1), round(linspace(1, N - 1, 20))]);
    S = sin(pi * sel(:) * (1:N - 1) / N);          % discrete sine vectors
    nall = 1:N - 1;
    MEAS = NaN(nl, numel(sel));
    PRED = NaN(nl, N - 1);
    MS = NaN(1, nl);
    for k = 1:nl
        lam = lams(k);
        dt = lam * h / c;
        K = ceil(12 / (f0 * dt));                  % twelve periods of the fundamental
        % every partial present, with amplitude 1/n, released from rest
        u0 = zeros(N + 1, 1);
        um = zeros(N + 1, 1);
        for n = nall
            u0 = u0 + sin(n * pi * x) / n;
            um = um + sin(n * pi * x) * cos(2 * pi * n * f0 * dt) / n;
        end
        t0 = tic;
        [~, U] = leapfrog(u0, um, lam, K, 'fixed', 1, K);
        MS(k) = toc(t0) * 1000;
        A = S * U(2:N, :);
        for i = 1:numel(sel)
            fm = wv_peakfreq(A(i, :), dt, 0, 0.5 / dt);
            MEAS(k, i) = wv_cents(fm, sel(i) * f0);
        end
        PRED(k, :) = drv_pred(nall, N, lam, dt, f0);
    end
    out.N = N;
    out.f0 = f0;
    out.lams = wv_vec(lams);
    out.sel = sel;
    out.nall = nall;
    out.meas = wv_rows(MEAS);
    out.pred = wv_rows(PRED);
    out.ms = wv_vec(MS);
    % the largest disagreement between measurement and prediction, in cents
    D = NaN(nl, numel(sel));
    for k = 1:nl
        D(k, :) = MEAS(k, :) - PRED(k, sel);
    end
    out.agree = max(abs(D(:)));
else
    % the click in a periodic tube
    c = p.c;
    fs = p.fs;
    dt = 1 / fs;
    L = p.L;
    N = round(L * p.lam / (c * dt));
    h = L / N;
    lam = c * dt / h;
    M = round(p.T * fs);
    x = (0:N - 1)' * h;
    sx = c * p.sigt;
    u0 = drv_train(x, 0, L, sx);
    um = drv_train(x + c * dt, 0, L, sx);          % moving right: u(x, -dt) = g(x + c dt)
    jr = round(N / 2) + 1;
    t0 = tic;
    [y, ~] = leapfrog(u0, um, lam, M, 'periodic', jr, 0);
    out.ms = toc(t0) * 1000;
    t = (1:M)' * dt;
    ye = drv_train(x(jr) - c * t, 0, L, sx);
    out.N = N;
    out.lam = lam;
    out.fcut = asin(min(lam, 1)) / (pi * dt);
    sc = max(abs(ye));
    out.clipExact = wv_audio(ye, fs, sc);
    out.clipFd = wv_audio(y, fs, sc);
    % windows around the first and the twentieth arrival
    tp = L / c;
    ta = x(jr) / c;
    for k = [1, 20]
        tk = ta + (k - 1) * tp;
        w = find(t >= tk - 0.002 & t < tk + 0.020);
        out.(sprintf('tw%d', k)) = (t(w)' - tk) * 1000;
        out.(sprintf('y%d', k)) = y(w)' / sc;
        out.(sprintf('ye%d', k)) = ye(w)' / sc;
    end
    out.rel = sqrt(sum((y - ye) .^ 2) / sum(ye .^ 2));
    % spectrogram of the simulated recording
    nw = 512;
    hop = 256;
    nfr = floor((M - nw) / hop) + 1;
    win = wv_window(nw);
    nb = nw / 2 + 1;
    G = zeros(nb, nfr);
    for k = 1:nfr
        seg = y((k - 1) * hop + (1:nw));
        F = abs(fft(win .* seg));
        G(:, k) = F(1:nb);
    end
    G = max(20 * log10(max(G / max(G(:)), 1e-6)), -80);
    keep = 1:2:nb;                                  % every other bin is plenty to draw
    out.spec = wv_rows(G(keep, :));
    out.specF = [0, (keep(end) - 1) * fs / nw];
    out.specT = [0, ((nfr - 1) * hop + nw / 2) / fs];
end

wv_write(out);

function v = drv_phi(x, beta)
% the plucked shape, extended oddly about 0 and periodically with period 2
r = mod(x, 2);
sg = ones(size(r));
sg(r > 1) = -1;
r(r > 1) = 2 - r(r > 1);
v = min(r / beta, (1 - r) / (1 - beta));
v = sg .* v;
end

function g = drv_train(x, x0, L, sx)
% a Gaussian pulse of width sx at x0, repeated with period L
d = mod(x - x0 + L / 2, L) - L / 2;
g = exp(-0.5 * (d / sx) .^ 2);
end

function c = drv_pred(n, N, lam, dt, f0)
% the leapfrog dispersion relation sin(w dt / 2) = lam sin(n pi / (2 N)),
% as a pitch error in cents against the exact n f0
w = 2 * asin(lam * sin(n * pi / (2 * N))) / dt;
c = wv_cents(w / (2 * pi), n * f0);
end