% Driver for Chapters 4 and 5: the square membrane.
%
% The editable script supplies
%
%     [y, U] = drum2d(u, uold, lam, M, scheme, ir, jr, K)
%
% which advances a scheme M steps on the unit square, rim held at zero.
%
% The exact solution is a sum of the modes sin(m pi x) sin(n pi y), each with
% frequency (c/2) sqrt(m^2 + n^2).  Those same sine vectors, sampled on the
% grid, are exact eigenvectors of every stencil in these chapters, so a
% scheme started from a single mode keeps it and changes only its frequency,
% which the driver measures from the recording and nothing else.
%
% mode 'strike'  a struck drum, with any scheme: the membrane at rest, given a
%                Gaussian blow of velocity at one point, simulated and heard
%                at another.  The reference is the modal sum, evaluated
%                exactly at the receiver.  The time step is lam h / c, so the
%                sound comes out at its own sample rate, 1/dt.
% mode 'cents'   the pitch error of modes in every direction, at three
%                wavelengths, for several Courant numbers (c = 1 here).
% mode 'orders'  the pitch error of three fixed modes against N, for each
%                scheme at its own Courant number, and the time each run took.
p = wv_params();
out = struct();

if strcmp(p.mode, 'strike')
    c = p.c;
    N = p.N;
    h = 1 / N;
    lam = p.lam;
    dt = lam * h / c;
    fs = 1 / dt;
    M = round(p.T / dt);
    K = p.kmax;
    [nn, mm] = meshgrid(1:K, 1:K);
    k2 = (mm .^ 2 + nn .^ 2) * pi ^ 2;
    w = c * sqrt(k2);
    sw = p.width;
    % sine coefficients of the initial velocity exp(-r^2 / (2 sw^2)); the
    % Gaussian is far enough from the rim that its truncation there is below
    % rounding, so these are exact
    A = 4 * 2 * pi * sw ^ 2 * sin(mm * pi * p.xs) .* sin(nn * pi * p.ys) .* exp(-k2 * sw ^ 2 / 2);
    x = (0:N)' * h;
    S = sin(pi * x * (1:K));
    u0 = zeros(N + 1, N + 1);
    um = S * (A .* sin(-w * dt) ./ w) * S';
    ir = round(p.xr * N) + 1;
    jr = round(p.yr * N) + 1;
    t0 = tic;
    [y, ~] = drum2d(u0, um, lam, M, p.scheme, ir, jr, 0);
    out.ms = toc(t0) * 1000;
    % the exact signal at the same node
    amp = A .* sin(mm * pi * x(ir)) .* sin(nn * pi * x(jr)) ./ w;
    keep = find(abs(amp) > 1e-10 * max(abs(amp(:))));
    t = (1:M)' * dt;
    ye = zeros(M, 1);
    for q = keep'
        ye = ye + amp(q) * sin(w(q) * t);
    end
    e = y - ye;
    out.N = N;
    out.lam = lam;
    out.fs = fs;
    out.rel = sqrt(sum(e .^ 2) / sum(ye .^ 2));
    out.nmodes = numel(keep);
    sc = max(abs(ye));
    out.clipExact = wv_audio(ye, fs, sc);
    out.clipFd = wv_audio(y, fs, sc);
    out.clipDiff = wv_audio(e, fs, sc);
    nw = min(M, round(0.4 * fs));
    [fr, dbe] = wv_spectrum(ye(1:nw), fs, p.fmax, 800);
    [~, dbf] = wv_spectrum(y(1:nw), fs, p.fmax, 800);
    out.fr = fr;
    out.dbe = dbe;
    out.dbf = dbf;
    % the lowest few mode frequencies that the blow excites, for the plot
    fm = w(:) / (2 * pi);
    am = abs(amp(:));
    [fm, ord] = sort(fm);
    am = am(ord);
    sel = am > 1e-3 * max(am) & fm < p.fmax;
    fsel = fm(sel);
    out.fmodes = fsel(1:min(end, 40))';
    out.f11 = c * sqrt(2) / 2;
    w0 = find(t >= 0.1 & t < 0.1 + 4 / out.f11);
    out.tw = t(w0)' * 1000;
    out.yw = y(w0)';
    out.yew = ye(w0)';
elseif strcmp(p.mode, 'cents')
    N = p.N;
    h = 1 / N;
    x = (0:N)' * h;
    lams = p.lams(:)';
    Rs = p.rs(:)';
    th = linspace(4, 86, p.nang) * pi / 180;
    MN = zeros(0, 3);
    for r = 1:numel(Rs)
        R = Rs(r);
        for q = 1:numel(th)
            m = round(R * cos(th(q)));
            n = round(R * sin(th(q)));
            if m >= 1 && n >= 1 && m < N && n < N
                MN(end + 1, :) = [m, n, r];
            end
        end
    end
    % no mode twice
    [~, iu] = unique(MN(:, 1) * 10000 + MN(:, 2));
    MN = MN(sort(iu), :);
    nq = size(MN, 1);
    Phi = zeros((N + 1) ^ 2, nq);
    for q = 1:nq
        f = sin(MN(q, 1) * pi * x) * sin(MN(q, 2) * pi * x)';
        Phi(:, q) = f(:);
    end
    wq = pi * sqrt(MN(:, 1) .^ 2 + MN(:, 2) .^ 2);     % exact, with c = 1
    fex = wq / (2 * pi);
    MEAS = NaN(numel(lams), nq);
    PRED = NaN(numel(lams), nq);
    for k = 1:numel(lams)
        lam = lams(k);
        dt = lam * h;
        Kst = ceil(12 / (min(fex) * dt));
        u0 = reshape(sum(Phi, 2), N + 1, N + 1);
        um = reshape(Phi * cos(wq * dt), N + 1, N + 1);
        [~, U] = drum2d(u0, um, lam, Kst, 'fd2', 1, 1, Kst);
        A = Phi' * U;
        for q = 1:nq
            fmq = wv_peakfreq(A(q, :), dt, 0, 0.5 / dt);
            MEAS(k, q) = wv_cents(fmq, fex(q));
        end
        s2 = sin(MN(:, 1) * pi * h / 2) .^ 2 + sin(MN(:, 2) * pi * h / 2) .^ 2;
        wp = 2 * asin(lam * sqrt(s2)) / dt;
        PRED(k, :) = wv_cents(wp / (2 * pi), fex)';
    end
    out.N = N;
    out.lams = wv_vec(lams);
    out.m = MN(:, 1)';
    out.n = MN(:, 2)';
    out.ring = MN(:, 3)';
    out.rs = Rs;
    out.angle = atan2(MN(:, 2), MN(:, 1))' * 180 / pi;
    out.ppw = (2 * N ./ sqrt(MN(:, 1) .^ 2 + MN(:, 2) .^ 2))';
    out.meas = wv_rows(MEAS);
    out.pred = wv_rows(PRED);
    out.agree = max(abs(MEAS(:) - PRED(:)));
else
    Ns = p.Ns(:)';
    schemes = p.schemes;
    lams = p.lams(:)';
    modes = reshape(p.modes, 2, [])';
    nmode = size(modes, 1);
    ns = numel(schemes);
    CE = NaN(ns * nmode, numel(Ns));
    TM = NaN(ns, numel(Ns));
    WORST = NaN(ns, numel(Ns));
    for s = 1:ns
        drum2d(zeros(12), zeros(12), 0.5, 20, schemes{s}, 3, 3, 0);    % warm-up
        for j = 1:numel(Ns)
            N = Ns(j);
            h = 1 / N;
            x = (0:N)' * h;
            lam = lams(s);
            dt = lam * h;
            % the first interior node: sin(m pi / N) is never zero for m < N,
            % so no mode has a nodal line through it
            ir = 2;
            jr = 2;
            wst = 0;
            for q = 1:nmode
                m = modes(q, 1);
                n = modes(q, 2);
                if m >= N || n >= N
                    continue;
                end
                wq = pi * sqrt(m ^ 2 + n ^ 2);
                f = sin(m * pi * x) * sin(n * pi * x)';
                M = ceil(12 * 2 * pi / (wq * dt));
                if q == 1
                    % the lowest mode's run is also the one that is timed:
                    % twelve of its periods, whatever the scheme, fastest of
                    % at least three runs while they are cheap
                    best = Inf;
                    tot = 0;
                    rep = 0;
                    while rep == 0 || (rep < 20 && tot < 0.3 && (rep < 3 || tot < 0.05))
                        t0 = tic;
                        y = drum2d(f, f * cos(wq * dt), lam, M, schemes{s}, ir, jr, 0);
                        el = toc(t0);
                        best = min(best, el);
                        tot = tot + el;
                        rep = rep + 1;
                    end
                    TM(s, j) = best * 1000;
                else
                    y = drum2d(f, f * cos(wq * dt), lam, M, schemes{s}, ir, jr, 0);
                end
                fm = wv_peakfreq(y, dt, 0, 0.5 / dt);
                cq = wv_cents(fm, wq / (2 * pi));
                CE((s - 1) * nmode + q, j) = cq;
                wst = max(wst, abs(cq));
            end
            WORST(s, j) = wst;
        end
    end
    out.Ns = Ns;
    out.schemes = schemes;
    out.lams = wv_vec(lams);
    out.modes = wv_rows(modes);
    out.cents = wv_rows(CE);
    out.worst = wv_rows(WORST);
    out.ms = wv_rows(TM);
    Q = NaN(ns, numel(Ns));
    for s = 1:ns
        Q(s, :) = wv_orders(WORST(s, :), Ns);
    end
    out.orders = wv_rows(Q);
end

wv_write(out);