"""
Comiskey-Yarin-Attinger 2017 back spatter -- corrected pipeline.

Implements paper equations (18, 19, 21, 23, 24, 25) with these regularizations:
  * v_z capped at V (potential-flow singularity at r=a is unphysical).
  * Total ejected mass rescaled to a chosen value. The paper's c=786.6 was
    calibrated for one experimental geometry; the DISTRIBUTION SHAPE is what
    the equations predict, while absolute mass must be re-pinned per scenario.
  * Optional wake drag-reduction factor (paper Sec III, refs [7,8]) since
    trailing drops in a dense spatter cone experience reduced drag in the
    aerodynamic wake of leading drops. Without this effect, single-drop
    trajectories of 20-100 um drops at 800 m/s have stopping distances
    20-100 cm and the model under-predicts long-range deposition.

Set WAKE_FACTOR < 1 to reduce drag; the paper's [4] used a 2-phase jet
treatment which is approximately wake_factor ~ 0.1-0.3 for drops behind
the leading edge. 1.0 = no wake effect (standalone single-drop physics).
"""
import numpy as np
import matplotlib
matplotlib.use('Agg')
import matplotlib.pyplot as plt

np.random.seed(42)

# ==================== PARAMETERS ====================
SCENARIO = "jon"            # "paper" or "jon"
WAKE_FACTOR = 0.2           # 1 = no wake effect, <1 = reduced drag for spray drops

if SCENARIO == "paper":
    Va = 1000.0; m_bullet = 2.916e-3; a = 0.285e-2
    Z_target = 0.50; total_ejected_vol_mL = 5.0
elif SCENARIO == "jon":
    Va = 800.0;  m_bullet = 9.7e-3;  a = 0.00381
    Z_target = 0.61; total_ejected_vol_mL = 1.5

theta_spread = 15 * np.pi / 180
rho   = 1060.0
sigma = 0.06045
c     = 786.6
w     = 0.112
rho_air = 1.2
g = 9.81
mu_air = 1.81e-5

V = Va / np.sqrt(1.0 + 4.0 * rho * a**3 / (3.0 * m_bullet))     # Eq. (20)

print(f"=== Scenario: {SCENARIO} ===")
print(f"Va = {Va} m/s   m = {m_bullet*1e3:.2f} g   a = {a*1e3:.3f} mm")
print(f"V0 = {V:.2f} m/s   r_max = 5a = {5*a*1e3:.2f} mm   Z = {Z_target*100:.0f} cm")
print(f"Target ejected volume (rescaled): {total_ejected_vol_mL} mL\n")

# ==================== EQUATIONS ====================
def vz_at(r):
    s = np.sqrt(np.maximum(r*r - a*a, 1e-30))
    return (2.0 * V / np.pi) * (a / s - np.arcsin(np.clip(a/r, 0.0, 1.0)))

def A_at(r):
    s = np.sqrt(np.maximum(r*r - a*a, 1e-30))
    return (2.0 * V*V / (np.pi * c * a)) * (a/s - np.arcsin(np.clip(a/r, 0.0, 1.0)))

def lstar(r):
    return 2.0 * np.pi * w / np.sqrt(rho * A_at(r) / (3.0 * sigma))

def M_bin_closed(b_lo, b_hi):
    """Eq. (23)."""
    s_hi = np.sqrt(b_hi*b_hi - a*a); s_lo = np.sqrt(b_lo*b_lo - a*a)
    return 4.0*rho*c*a*(
        2.0*a*(s_hi - s_lo)
        - 0.5*(b_hi*b_hi*np.arcsin(a/b_hi) - b_lo*b_lo*np.arcsin(a/b_lo))
    )

# ==================== SUBFAMILIES ====================
n_rings = 100
u = np.linspace(np.log(1.05 - 1.0), np.log(5.0 - 1.0), n_rings + 1)
r_edges = a * (1.0 + np.exp(u))
r_mid = np.sqrt(r_edges[:-1] * r_edges[1:])

M_i_raw = np.array([M_bin_closed(r_edges[k], r_edges[k+1]) for k in range(n_rings)])
target_mass = (total_ejected_vol_mL * 1e-6) * rho
M_i = M_i_raw * (target_mass / M_i_raw.sum())   # rescale to chosen volume
l_i = lstar(r_mid)
v_i_raw = vz_at(r_mid)
v_i = np.clip(v_i_raw, 0.0, V)                  # physical cap

vol_drop = (4.0/3.0) * np.pi * (l_i/2.0)**3
n_i = M_i / (rho * vol_drop)

print("--- Subfamily summary (rescaled) ---")
print(f"Rings: {n_rings},  r/a range: {r_mid[0]/a:.3f} - {r_mid[-1]/a:.3f}")
print(f"Total mass: {M_i.sum()*1e3:.3f} g  ({M_i.sum()/rho*1e6:.2f} mL)")
print(f"Total drops: {n_i.sum():,.0f}")
print(f"v_z range: {v_i.min():.1f} - {v_i.max():.1f} m/s (cap = V = {V:.0f})")
print(f"l* range: {l_i.min()*1e6:.1f} - {l_i.max()*1e6:.0f} um")
print(f"Number-weighted mean l*: {np.average(l_i, weights=n_i)*1e6:.1f} um")
print(f"Mass-weighted mean l*:   {np.average(l_i, weights=M_i)*1e6:.0f} um")
print(f"Drops per ring (first/last): {n_i[0]:.2e} / {n_i[-1]:.2e}\n")

# ==================== VECTORIZED TRAJECTORY ====================
def trajectories_vec(l_drops, u0s, psis, Z_tgt, t_max=2.0):
    """Vectorized RK2 with per-drop adaptive substeps."""
    n = l_drops.size
    diam = l_drops
    m = rho * (np.pi/6.0) * diam**3
    A_cs = np.pi * (diam/2.0)**2

    x  = np.zeros(n); z = np.zeros(n); y = np.zeros(n)
    vx = u0s * np.sin(psis)
    vz = u0s * np.cos(psis)
    vy = np.zeros(n)
    alive = np.ones(n, dtype=bool)
    hit   = np.zeros(n, dtype=bool)
    X_hit = np.full(n, np.nan); Y_hit = np.full(n, np.nan); Sp_hit = np.full(n, np.nan)

    # Per-drop time step: stopping time tau_s ~ m / (0.5 Cd_typ rho_air A v).
    # Use dt = 0.02 * stopping_time, floored at 1e-6 s.
    sp0 = np.sqrt(vx*vx + vy*vy + vz*vz)
    Re0 = rho_air * sp0 * diam / mu_air
    Cd0 = 0.28 + 6.0/np.sqrt(np.maximum(Re0,1e-3)) + 21.0/np.maximum(Re0,1e-3)
    stopping_time = m / (0.5 * Cd0 * rho_air * A_cs * np.maximum(sp0, 1.0))
    dt_per = np.clip(0.02 * stopping_time, 1e-7, 5e-4)

    def accel(vxv, vyv, vzv):
        sp = np.sqrt(vxv*vxv + vyv*vyv + vzv*vzv)
        sp_safe = np.maximum(sp, 1e-9)
        Re = rho_air * sp * diam / mu_air
        Re_c = np.maximum(Re, 1e-3)
        Cd = (0.28 + 6.0/np.sqrt(Re_c) + 21.0/Re_c) * WAKE_FACTOR
        F_over_m = 0.5 * Cd * rho_air * A_cs * sp * sp / m
        ax = -F_over_m * vxv / sp_safe
        ay = -g - F_over_m * vyv / sp_safe
        az = -F_over_m * vzv / sp_safe
        return ax, ay, az

    # March in global time; each drop takes its own dt
    t = 0.0
    # Cap loop iterations
    max_iters = int(t_max / dt_per.min()) + 1
    max_iters = min(max_iters, 200000)
    for it in range(max_iters):
        if not alive.any():
            break
        dt = np.where(alive, dt_per, 0.0)
        ax1, ay1, az1 = accel(vx, vy, vz)
        vxm = vx + 0.5*dt*ax1; vym = vy + 0.5*dt*ay1; vzm = vz + 0.5*dt*az1
        ax2, ay2, az2 = accel(vxm, vym, vzm)
        z_new = z + dt * vzm
        x_new = x + dt * vxm
        y_new = y + dt * vym
        vx_new = vx + dt*ax2; vy_new = vy + dt*ay2; vz_new = vz + dt*az2

        crossed = alive & (z_new >= Z_tgt) & (z < Z_tgt) & ~hit
        if crossed.any():
            frac = (Z_tgt - z[crossed]) / np.maximum(z_new[crossed] - z[crossed], 1e-12)
            X_hit[crossed] = x[crossed] + frac * (x_new[crossed] - x[crossed])
            Y_hit[crossed] = y[crossed] + frac * (y_new[crossed] - y[crossed])
            Sp_hit[crossed] = np.sqrt(vx_new[crossed]**2 + vy_new[crossed]**2 + vz_new[crossed]**2)
            hit[crossed] = True
            alive[crossed] = False

        # kill drops that have stopped going forward or fallen
        dead = alive & ((vz_new < 0.05) | (y_new < -1.0))
        alive[dead] = False

        x, y, z = x_new, y_new, z_new
        vx, vy, vz = vx_new, vy_new, vz_new
        t += dt_per.max()
        if t > t_max:
            break

    return X_hit, Y_hit, hit, Sp_hit

# ==================== SIMULATE ====================
print("Sampling and integrating trajectories...")
n_samples = 20000
p = n_i / n_i.sum()
idx = np.random.choice(len(r_mid), size=n_samples, p=p)
phi = np.random.uniform(0, 2*np.pi, size=n_samples)

l_s   = l_i[idx]
u0_s  = v_i[idx]
# For normal impact (delta=0) every drop launches at psi=theta from bullet axis.
# Radial spread on target comes from gravity acting differently on drops with
# different sizes and launch speeds (slower drops fall more, sweeping out a fan).
psi_s = np.full(n_samples, theta_spread)

X_loc, Y_loc, hit_mask, Sp = trajectories_vec(l_s, u0_s, psi_s, Z_target)

X_target = X_loc * np.cos(phi) + 0.0          # rotate launch x to lateral on target
Y_target = X_loc * np.sin(phi) + Y_loc        # vertical includes gravity drop

X = X_target[hit_mask]; Y = Y_target[hit_mask]
sizes = l_s[hit_mask]; v_imp = Sp[hit_mask]
radial = np.sqrt(X*X + Y*Y)

print(f"Hits at Z = {Z_target*100:.0f} cm: {hit_mask.sum()} / {n_samples} ({100*hit_mask.mean():.1f}%)")
if len(X):
    print(f"Radial:  min={radial.min()*100:.2f}  median={np.median(radial)*100:.2f}  "
          f"95th={np.percentile(radial,95)*100:.2f}  max={radial.max()*100:.2f} cm")
    print(f"Sizes at impact: {sizes.min()*1e6:.0f}-{sizes.max()*1e6:.0f} um "
          f"(median {np.median(sizes)*1e6:.0f}, mass-weighted {np.average(sizes, weights=sizes**3)*1e6:.0f})")
    print(f"Impact speed: {v_imp.min():.1f}-{v_imp.max():.1f} m/s (median {np.median(v_imp):.1f})")
    hist, edges = np.histogram(radial*100, bins=40, range=(0, max(radial.max()*100, 30)))
    peak = 0.5*(edges[np.argmax(hist)] + edges[np.argmax(hist)+1])
    print(f"Peak stain density at radial ~ {peak:.1f} cm")
print()

# ==================== PLOTS ====================
fig, axs = plt.subplots(2, 2, figsize=(14, 10), facecolor='#0a0f1a')
for ax in axs.flat:
    ax.set_facecolor('#0a0f1a')
    for s in ax.spines.values(): s.set_color('#888')
    ax.tick_params(colors='#ccc')
    ax.xaxis.label.set_color('#ccc'); ax.yaxis.label.set_color('#ccc')
    ax.title.set_color('#fff')

ax = axs[0, 0]
ax.plot(r_mid/a, l_i*1e6, color='#00d4ff', lw=2)
ax.set_xlabel('r / a'); ax.set_yscale('log')
ax.set_ylabel('l* [um]', color='#00d4ff')
ax.tick_params(axis='y', labelcolor='#00d4ff')
ax2 = ax.twinx(); ax2.set_facecolor('#0a0f1a')
ax2.plot(r_mid/a, v_i, color='#ff8800', lw=2, label='v_z (capped)')
ax2.plot(r_mid/a, v_i_raw, color='#ff8800', lw=1, ls='--', alpha=0.5, label='v_z raw (singular)')
ax2.axhline(V, color='#ffaa00', lw=0.8, ls=':', alpha=0.6)
ax2.set_ylim(0, V*1.3)
ax2.set_ylabel('v_z [m/s]', color='#ff8800')
ax2.tick_params(axis='y', labelcolor='#ff8800')
for s in ax2.spines.values(): s.set_color('#888')
ax2.legend(loc='upper right', facecolor='#0a0f1a', labelcolor='#ccc', edgecolor='#444')
ax.set_title(f'Eqs. (18), (21): free-surface velocity and droplet size\nV0={V:.0f} m/s, a={a*1e3:.2f} mm')

ax = axs[0, 1]
ax.hist(l_i*1e6, weights=n_i, bins=40, color='#ff3355', alpha=0.85)
ax.set_xlabel('Droplet diameter [um]'); ax.set_ylabel('Drops per bin')
ax.set_yscale('log')
ax.set_title(f'Size distribution\n{n_i.sum():,.0f} drops, {M_i.sum()*1e3:.2f} g rescaled')

ax = axs[1, 0]
if len(radial):
    ax.hist(radial*100, bins=40, color='#ff3355', alpha=0.85,
            range=(0, max(radial.max()*100, 30)))
ax.set_xlabel(f'Radial distance at Z = {Z_target*100:.0f} cm [cm]')
ax.set_ylabel('Stain count')
ax.set_title(f'Radial stain distribution ({len(X)} hits)')

ax = axs[1, 1]
if len(X):
    sc = ax.scatter(X*100, Y*100, s=np.clip(sizes*8e5, 1, 120),
                    c=v_imp, cmap='plasma', alpha=0.6, edgecolors='none')
    cb = plt.colorbar(sc, ax=ax, label='Impact speed [m/s]')
    cb.ax.yaxis.label.set_color('#ccc'); cb.ax.tick_params(colors='#ccc')
ax.set_xlabel('X [cm]'); ax.set_ylabel('Y [cm] (gravity-affected)')
ax.set_title('2D spatter pattern on target')
ax.set_aspect('equal'); ax.axhline(0, color='#444', lw=0.5); ax.axvline(0, color='#444', lw=0.5)

plt.tight_layout()
plt.savefig('/home/claude/spatter_final.png', dpi=140, bbox_inches='tight', facecolor='#0a0f1a')
print("Figure: /home/claude/spatter_final.png")