#!/usr/bin/env python3
"""
sound-in-light-out-v1.py -- the numbers behind "Sound In, Light Out" on
solvetheuniverse.com, computed from stated inputs rather than quoted.

  1. Rayleigh collapse time and wall speed of an empty spherical cavity
  2. Adiabatic temperature of argon: the naive whole-collapse estimate
     (wrong, and why), and the last-stage estimate the article uses
  3. Energy and peak power of one flash; sound energy per water molecule

Run with no arguments to print the numbers used in the article.
Run with --selftest to check the script against results that do not
depend on it (Rayleigh's prefactor by direct integration, the integrator's
wall speed against the closed form, three closed-form controls on the
adiabatic step, and a hand value for the flash energy).

EVERY INPUT BELOW IS AN ASSUMPTION, chosen to be typical of published
single-bubble experiments. The one that matters most, R_MIN, is set to a
sixth of the rest radius, which is the compression ratio two published
estimates give (Storey & Szeri 2000; Bataller 2014 thesis); the script
prints that comparison, and also an excluded-volume floor computed from
argon's van der Waals constant (a soft floor, since b is four times the
atoms' own volume). The self-test checks the arithmetic; it does not
validate R_MIN. Change any input and the output changes.
"""
import argparse
import math
import sys

# ---- inputs, all assumed (see docstring) ----------------------------------
RHO_WATER = 1000.0        # kg/m^3
C_WATER = 1480.0          # m/s, speed of sound in water
P0 = 1.013e5              # Pa, ambient pressure
P_A = 1.3 * P0            # Pa, acoustic drive amplitude (typical SBSL ~1.2-1.5 atm)
R_MAX = 50e-6             # m, radius at the top of the expansion (typical)
R0 = 5e-6                 # m, rest radius (typical)
R_MIN = 0.85e-6           # m, minimum radius: a sixth of R0, the ratio of two published estimates (see docstring)
SIGMA_WATER = 0.0725      # N/m, surface tension, for the bubble's rest pressure P0 + 2 sigma / R0
# Published compression ratios R_min / R0 the chosen R_MIN is tested against:
PUBLISHED_RATIOS = {"Storey & Szeri 2000 (simulation, argon, 4.5 um; 0.14-0.18 across cases)": 0.156,
                    "Bataller 2014 thesis (RP fit to Mie scattering, xenon in water)": 0.547 / 3.20}
T0 = 293.0                # K, water and gas temperature before collapse
GAMMA_AR = 5.0 / 3.0      # monatomic ideal gas
B_ARGON = 0.03219e-3      # m^3/mol, van der Waals b for argon (CRC value)
N_A = 6.02214076e23
K_B = 1.380649e-23
PHOTONS = (1e5, 1e6)      # photons per flash: Barber & Putterman 1991 report "over 10^5"; 10^6 is a bright case
PHOTON_EV = 3.0           # eV, an average visible/near-UV photon (the spectrum is UV-weighted; this is a round value)
FLASH_S = 100e-12         # s, flash width; Gompf et al. 1997 measured 60 to >250 ps
DRIVE_HZ = 25.7e3         # Hz, one flash per cycle (the header photograph's drive frequency)
N_WATER = 3.34e28         # molecules per m^3 of water
EV = 1.602176634e-19      # J
RAYLEIGH_PREFACTOR = 0.914681   # exact value; 0.915 is the usual rounding


def rayleigh_time(r_max, dp, rho):
    """Rayleigh (1917): collapse time of an empty cavity from rest at r_max."""
    return RAYLEIGH_PREFACTOR * r_max * math.sqrt(rho / dp)


def rayleigh_wall_speed(r, r_max, dp, rho):
    """Rayleigh: U^2 = (2 dp / 3 rho) ((r_max/r)^3 - 1); inviscid, incompressible, empty cavity."""
    return math.sqrt(2.0 * dp / (3.0 * rho) * ((r_max / r) ** 3 - 1.0))


def radius_where_wall_speed_is(u, r_max, dp, rho):
    """Invert rayleigh_wall_speed for r."""
    return r_max / ((u * u * 3.0 * rho / (2.0 * dp) + 1.0) ** (1.0 / 3.0))


def sound_speed_argon(t):
    return math.sqrt(GAMMA_AR * 8.314462618 * t / 0.039948)


def excluded_volume_radius(r0, p, t, fraction=1.0):
    """Radius the gas in a bubble of rest radius r0 would occupy if packed to
    `fraction` of its van der Waals excluded volume: N atoms at (p, t), each
    excluding b/N_A. fraction=1 is the conventional loose floor (the RMP
    review's R0/8.86 at STP); fraction=0.25 is the atoms' own volume."""
    n_atoms = p * (4.0 / 3.0) * math.pi * r0 ** 3 / (K_B * t)
    v_excl = n_atoms * B_ARGON / N_A * fraction
    return (3.0 * v_excl / (4.0 * math.pi)) ** (1.0 / 3.0)


def adiabatic_temperature(t0, r_start, r_end, gamma, hard_core=0.0):
    """T = T0 ((V_start - b)/(V_end - b))^(gamma-1), optional excluded volume."""
    v_start = r_start ** 3 - hard_core ** 3
    v_end = r_end ** 3 - hard_core ** 3
    return t0 * (v_start / v_end) ** (gamma - 1.0)


def flash_energy(n_photons, photon_ev):
    return n_photons * photon_ev * EV


def acoustic_energy_density(p_a, rho, c):
    """Time-averaged energy density of a plane sound wave of pressure amplitude p_a."""
    return p_a ** 2 / (2.0 * rho * c ** 2)


def integrate_rayleigh_collapse(r_max, dp, rho, steps=200000, stop_frac=0.02, want_speed_at=None):
    """Direct integration of R R'' + 1.5 R'^2 = -dp/rho from rest at r_max,
    down to stop_frac * r_max. Independent of the closed forms above.
    Returns (time, speed_at_requested_radius), the speed linearly
    interpolated to the requested radius. Used only by --selftest."""
    r, v, t = r_max, 0.0, 0.0
    dt = rayleigh_time(r_max, dp, rho) / steps
    speed_at = None
    while r > stop_frac * r_max:
        a = (-dp / rho - 1.5 * v * v) / r
        v_new = v + a * dt
        r_new = r + v_new * dt
        if want_speed_at is not None and speed_at is None and r_new <= want_speed_at < r:
            f = (r - want_speed_at) / (r - r_new)
            speed_at = -(v + f * (v_new - v))
        v, r, t = v_new, r_new, t + dt
    return t, speed_at


def main():
    ap = argparse.ArgumentParser(description=__doc__, formatter_class=argparse.RawDescriptionHelpFormatter)
    ap.add_argument("--selftest", action="store_true")
    args = ap.parse_args()
    if args.selftest:
        return selftest()

    print("All inputs are assumptions; see the docstring.\n")
    tau = rayleigh_time(R_MAX, P0, RHO_WATER)
    c_ar = sound_speed_argon(T0)
    print(f"1. Rayleigh collapse of an empty cavity, Rmax={R_MAX*1e6:.0f} um, dp={P0/1e5:.2f} bar")
    print(f"   collapse time: {tau*1e6:.2f} us  (period at {DRIVE_HZ/1e3:.1f} kHz: {1e6/DRIVE_HZ:.1f} us)")
    for frac in (10, 20):
        u = rayleigh_wall_speed(R_MAX / frac, R_MAX, P0, RHO_WATER)
        print(f"   wall speed at Rmax/{frac}: {u:.0f} m/s")
    r_c = radius_where_wall_speed_is(c_ar, R_MAX, P0, RHO_WATER)
    print(f"   speed of sound in argon at {T0:.0f} K: {c_ar:.0f} m/s, reached at Rmax/{R_MAX/r_c:.1f}")
    print(f"   (the measured collapse is faster than Mach 4 in the gas: Weninger et al. 1997)")

    print(f"\n2. Adiabatic argon, gamma = 5/3")
    t_naive = adiabatic_temperature(T0, R_MAX, R_MIN, GAMMA_AR)
    print(f"   naive, whole collapse adiabatic, {R_MAX*1e6:.0f} um -> {R_MIN*1e6:.2f} um: {t_naive:,.0f} K  (WRONG: most of the collapse is slow enough to stay isothermal)")
    t_last = adiabatic_temperature(T0, R0, R_MIN, GAMMA_AR)
    h = excluded_volume_radius(R0, P0, T0)
    t_hc = adiabatic_temperature(T0, R0, R_MIN, GAMMA_AR, hard_core=h)
    print(f"   last stage only, {R0*1e6:.0f} um -> {R_MIN*1e6:.2f} um: {t_last:,.0f} K  (the article's estimate)")
    print(f"   R_MIN / R0 = {R_MIN/R0:.3f}; published estimates of that ratio:")
    for k, v in PUBLISHED_RATIOS.items():
        print(f"      {v:.3f}  {k}  -> would give {adiabatic_temperature(T0, R0, v*R0, GAMMA_AR):,.0f} K here")
    p_rest = P0 + 2 * SIGMA_WATER / R0
    h4 = excluded_volume_radius(R0, p_rest, T0, fraction=0.25)
    print(f"   excluded-volume floor from argon's van der Waals b: {h*1e6:.2f} um at 1 atm (R0/{R0/h:.1f}), "
          f"{excluded_volume_radius(R0, p_rest, T0)*1e6:.2f} um at the rest pressure {p_rest/P0:.2f} atm; "
          f"atoms' own volume (b/4): {h4*1e6:.2f} um. R_MIN is {R_MIN/h:.2f} x the 1 atm floor")
    print(f"   same last stage with the excluded volume subtracted: {t_hc:,.0f} K")
    print(f"   sensitivity: R_MIN 0.7 um -> {adiabatic_temperature(T0, R0, 0.7e-6, GAMMA_AR):,.0f} K; 1.0 um -> {adiabatic_temperature(T0, R0, 1.0e-6, GAMMA_AR):,.0f} K")
    print(f"   (water vapour trapped at collapse dissociates and absorbs energy, which lowers the real peak: Storey & Szeri 2000)")

    print("\n3. One flash, and the sound energy it came from")
    for n in PHOTONS:
        e = flash_energy(n, PHOTON_EV)
        print(f"   {n:.0e} photons x {PHOTON_EV} eV = {e:.1e} J; over {FLASH_S*1e12:.0f} ps = {e/FLASH_S*1e3:.1f} mW peak; "
              f"averaged at {DRIVE_HZ/1e3:.1f} kHz = {e*DRIVE_HZ:.1e} W")
    ed = acoustic_energy_density(P_A, RHO_WATER, C_WATER)
    per_mol_ev = ed / N_WATER / EV
    print(f"   acoustic energy density at p_a = {P_A/P0:.1f} atm: {ed:.2f} J/m^3 = {per_mol_ev:.1e} eV per water molecule")
    print(f"   one {PHOTON_EV} eV photon / that = {PHOTON_EV/per_mol_ev:.1e}, i.e. {math.log10(PHOTON_EV/per_mol_ev):.1f} orders of magnitude")
    return 0


def selftest():
    ok = True

    # (a) Rayleigh's prefactor by direct integration from rest. The integration
    # stops at 2% of Rmax; the time an already-collapsing cavity takes from a
    # small radius r to zero scales as r^(5/2), so the missing tail is about
    # 3e-5 of tau (the observed residual is smaller because the integrator's
    # truncation error partly offsets it). Compared against the exact prefactor
    # 0.914681, not the rounded 0.915, which would fail this tolerance.
    tau_formula = rayleigh_time(R_MAX, P0, RHO_WATER)
    tau_num, u_num = integrate_rayleigh_collapse(R_MAX, P0, RHO_WATER, want_speed_at=R_MAX / 10)
    err = abs(tau_num - tau_formula) / tau_formula
    print(f"[a] Rayleigh time: formula {tau_formula*1e6:.4f} us, integrated {tau_num*1e6:.4f} us, rel err {err:.1e}")
    ok &= err < 2e-4

    # (b) The integrator's wall speed at Rmax/10 against the closed form.
    u_formula = rayleigh_wall_speed(R_MAX / 10, R_MAX, P0, RHO_WATER)
    err_u = abs(u_num - u_formula) / u_formula
    # The interpolated semi-implicit Euler speed is first order in dt (about
    # 2e-3 at 200k steps, halving as steps double), so the tolerance is 5e-3.
    print(f"[b] wall speed at Rmax/10: formula {u_formula:.1f} m/s, integrated {u_num:.1f} m/s, rel err {err_u:.1e}")
    ok &= err_u < 5e-3

    # (c) Closed-form controls on the adiabatic step.
    same = adiabatic_temperature(T0, R0, R0, GAMMA_AR)
    print(f"[c1] no compression returns T0: {same:.1f} K (expect {T0})")
    ok &= abs(same - T0) < 1e-9
    iso = adiabatic_temperature(T0, R0, R_MIN, 1.0)
    print(f"[c2] gamma=1 (isothermal) returns T0: {iso:.1f} K (expect {T0})")
    ok &= abs(iso - T0) < 1e-9
    quad = adiabatic_temperature(T0, R0, R0 / 2, GAMMA_AR)
    print(f"[c3] halving the radius of a monatomic gas quadruples T: {quad/T0:.6f} (expect 4)")
    ok &= abs(quad / T0 - 4.0) < 1e-9

    # (d) Negative controls: a diatomic gas heats less than a monatomic one for
    # the same compression, and subtracting an excluded volume heats it more.
    diatomic = adiabatic_temperature(T0, R0, R_MIN, 1.4)
    mono = adiabatic_temperature(T0, R0, R_MIN, GAMMA_AR)
    hc = adiabatic_temperature(T0, R0, R_MIN, GAMMA_AR, hard_core=excluded_volume_radius(R0, P0, T0))
    print(f"[d] gamma=1.4 gives {diatomic:,.0f} K < gamma=5/3 gives {mono:,.0f} K < with excluded volume {hc:,.0f} K")
    ok &= diatomic < mono < hc

    # (e) Excluded-volume radius: for an ideal gas at 1 atm and 293 K the
    # excluded volume is b/(kT/p) = 0.03219e-3 / 0.02405 = 0.134 % of the
    # gas volume, so the radius ratio is 0.00134^(1/3) = 0.110, i.e. R0/9.1.
    h = excluded_volume_radius(R0, P0, T0)
    print(f"[e] excluded-volume radius R0/{R0/h:.2f} (hand: R0/9.1)")
    ok &= abs(R0 / h - 9.1) < 0.1

    # (f) Energy arithmetic against a hand value: 1e6 x 3 eV = 4.8065e-13 J.
    e = flash_energy(1e6, 3.0)
    print(f"[f] 1e6 x 3 eV = {e:.4e} J (hand: 4.8065e-13)")
    ok &= abs(e - 4.8065e-13) / 4.8065e-13 < 1e-3

    print("SELFTEST", "PASS" if ok else "FAIL")
    return 0 if ok else 1


if __name__ == "__main__":
    sys.exit(main())