function [Qdd, Xdd, force, helmholtzHz, contactMs, peakForceN, maxRealLam] = woodblock_solve(P, basis, oversample)
% WOODBLOCK_SOLVE Integrate the coupled mallet / body / cavity system.
%
% The system is nonlinear only while the mallet is in contact (a fraction of
% a millisecond), so the solve is split in two:
%   contact phase: RK4 on the full nonlinear system, until the mallet has
%     rebounded far enough that no future contact is possible;
%   free phase: the remaining coupled body+cavity system is linear, so it is
%     evolved in closed form from one eigendecomposition of the (2N+2)
%     system matrix, evaluated vectorized over the whole time axis.
% This departs from the reference woodblock.py, which runs RK4 throughout;
% the free-phase eigen solution is exact in time, so the two agree to the
% reference's own integration tolerance.
%
% Returns modal accelerations Qdd (N x T), slit-plug acceleration Xdd (1 x T),
% contact force (1 x T), and scalar diagnostics. maxRealLam should be
% strictly negative; anything else means the coupling was assembled wrongly.

N = basis.n;
fsHi = P.sample_rate * oversample;
dt = 1 / fsHi;
nSteps = floor(P.duration_s * fsHi);

om = basis.omega;
om2 = om.^2;
tzo = 2 * basis.zeta .* om;
phiC = basis.phiC;
A = basis.A;

rho0 = P.air_density; c = P.speed_of_sound;
V = P.cavity_volume_m3; S = P.slit_area_m2;
% End correction on the slit: ~0.85 * equivalent radius per open end.
Leff = P.slit_length_m + 2 * 0.85 * sqrt(S / pi);
omH = c * sqrt(S / (V * Leff));
tzoH = 2 * P.cavity_damping * omH;
pScale = rho0 * c^2 / V;
cavDrive = c^2 / (V * Leff);
K = P.mallet_stiffness;
alpha = P.contact_exponent;
mu = P.contact_damping;
invMh = 1 / P.mallet_mass;

% ---- contact phase: RK4 on [q; qd; xi; xid; y; yd] ----
maxContact = min(nSteps, ceil(0.02 * fsHi));   % 20 ms hard cap
z = zeros(2 * N + 4, 1);
z(2 * N + 4) = P.mallet_velocity;
Qc = zeros(N, maxContact);
Xc = zeros(1, maxContact);
forceC = zeros(1, maxContact);
Tc = 0;
for k = 1:maxContact
  [d1, qdd1, xidd1, fc1] = wb_deriv(z, N, om2, tzo, phiC, A, S, pScale, omH, tzoH, cavDrive, K, alpha, mu, invMh);
  Qc(:, k) = qdd1;
  Xc(k) = xidd1;
  forceC(k) = fc1;
  Tc = k;
  [d2, ~, ~, ~] = wb_deriv(z + 0.5 * dt * d1, N, om2, tzo, phiC, A, S, pScale, omH, tzoH, cavDrive, K, alpha, mu, invMh);
  [d3, ~, ~, ~] = wb_deriv(z + 0.5 * dt * d2, N, om2, tzo, phiC, A, S, pScale, omH, tzoH, cavDrive, K, alpha, mu, invMh);
  [d4, ~, ~, ~] = wb_deriv(z + dt * d3, N, om2, tzo, phiC, A, S, pScale, omH, tzoH, cavDrive, K, alpha, mu, invMh);
  z = z + (dt / 6) * (d1 + 2 * d2 + 2 * d3 + d4);

  % No-recontact exit: after separation the mallet coasts at constant
  % velocity while the surface stays within B (a conservative bound from
  % the current modal energy), so once the mallet is receding and further
  % than 1.5*B below the surface rest plane, contact can never resume.
  y = z(2 * N + 3);
  yd = z(2 * N + 4);
  if fc1 == 0 && yd <= 0
    q = z(1:N);
    qd = z(N + 1:2 * N);
    B = sum(abs(phiC) .* sqrt(q.^2 + (qd ./ om).^2));
    if y <= -1.5 * B
      break;
    end
  end
end

% ---- free phase: eigendecomposition of the linear body+cavity system ----
% state x = [q; qd; xi; xid], xdot = MM x
T2 = nSteps - Tc;
MM = zeros(2 * N + 2);
MM(1:N, N + 1:2 * N) = eye(N);
MM(N + 1:2 * N, 1:N) = -(diag(om2) + pScale * (A * A.'));
MM(N + 1:2 * N, N + 1:2 * N) = -diag(tzo);
MM(N + 1:2 * N, 2 * N + 1) = -pScale * S * A;
MM(2 * N + 1, 2 * N + 2) = 1;
MM(2 * N + 2, 1:N) = -cavDrive * A.';
MM(2 * N + 2, 2 * N + 1) = -omH^2;
MM(2 * N + 2, 2 * N + 2) = -tzoH;

[VV, DD] = eig(MM);
lam = diag(DD);
maxRealLam = max(real(lam));
cvec = VV \ z(1:2 * N + 2);
w = lam .* cvec;                       % acceleration weights per eigenmode
t2 = (0:T2 - 1) * dt;                  % t2 = 0 is the state right after contact
E = exp(lam * t2);                     % (2N+2) x T2, complex
Qf = real((VV(N + 1:2 * N, :) * diag(w)) * E);
Xf = real((VV(2 * N + 2, :) * diag(w)) * E);

Qdd = [Qc(:, 1:Tc) Qf];
Xdd = [Xc(1:Tc) Xf];
force = [forceC(1:Tc) zeros(1, T2)];

helmholtzHz = omH / (2 * pi);
contactMs = sum(force > 0) * dt * 1e3;
peakForceN = max(force);
end

function [d, qdd, xidd, fc] = wb_deriv(z, N, om2, tzo, phiC, A, S, pScale, omH, tzoH, cavDrive, K, alpha, mu, invMh)
% Derivative of the coupled system with Hunt-Crossley mallet contact.
q = z(1:N);
qd = z(N + 1:2 * N);
xi = z(2 * N + 1);
xid = z(2 * N + 2);
y = z(2 * N + 3);
yd = z(2 * N + 4);

wsurf = phiC.' * q;
wdsurf = phiC.' * qd;
delta = y - wsurf;
if delta > 0
  % f = K d^a (1 + lambda * ddot); lambda in s/m keeps the correction O(1)
  fc = K * delta^alpha * (1 + mu * (yd - wdsurf));
  if fc < 0
    fc = 0;
  end
else
  fc = 0;
end

Aq = A.' * q;
pc = pScale * (S * xi + Aq);

qdd = -tzo .* qd - om2 .* q + fc * phiC - A * pc;
xidd = -tzoH * xid - omH^2 * xi - cavDrive * Aq;

d = zeros(2 * N + 4, 1);
d(1:N) = qd;
d(N + 1:2 * N) = qdd;
d(2 * N + 1) = xid;
d(2 * N + 2) = xidd;
d(2 * N + 3) = yd;
d(2 * N + 4) = -fc * invMh;
end