"""
2D shirt-face view of entry-side back spatter.

Geometry: shirt is 5 mm in front of the entry wound. Bullet passes through
the shirt at the origin (entry hole), then through the victim. Back-spatter
drops are ejected back toward the shooter through the entry wound and
deposit on the shirt within 1-2 ms.

Physics: Comiskey/Yarin/Kim/Attinger 2016 (Phys Rev Fluids 1, 043201)
slender-bullet back-spatter model.
  - Drop sizes: 27 - 84 um (RT instability, Eq. 22)
  - Drop velocities: 19 - 183 m/s (Eq. 23)
  - Ejection cone half-angle: ~57 deg (high-speed video studies)

Stain size on fabric: Eq. (29) spread factor xi = 0.61 (We/Oh)^0.166
  - Conservative for fabric (real fabric spreads more than smooth paper)
  - Final stain diameter = xi * l_drop / sin(impact_angle)

Visualization: laid-flat shirt as seen by an investigator looking at it.
Origin = bullet hole. Off-white fabric, dark red stains.
"""
import numpy as np
import matplotlib
matplotlib.use('Agg')
import matplotlib.pyplot as plt
from matplotlib.patches import Ellipse, Circle
from matplotlib.collections import PatchCollection

np.random.seed(42)

# ==================== PARAMETERS ====================
Va = 800.0
m_bullet = 9.7e-3
a = 0.00381                  # bullet radius
rho = 1060.0
sigma = 0.06045
rho_air = 1.2
g = 9.81
mu_air = 1.81e-5

theta_tip = 14 * np.pi / 180     # spitzer half-angle
h_penetration = 0.02
c_slender = 0.001
w_slender = 0.9
theta_backward = 57 * np.pi / 180  # back-spatter cone half-angle

Z_shirt = -0.005                 # 5 mm in front of skin
total_back_vol_mL = 0.3

WAKE_FACTOR = 0.2
V = Va / np.sqrt(1.0 + 4.0 * rho * a**3 / (3.0 * m_bullet))

# Display
MAX_VISIBLE_STAINS = 3000        # cap for rendering performance
SHIRT_VIEW_HALF = 0.08           # show +/- 8 cm of shirt around hole

# ==================== 2016 BACK-SPATTER SUBFAMILIES ====================
def back_subfamilies(n_rings=80):
    h = h_penetration
    r_min = theta_tip * h * 1.05
    r_max = 0.05
    r_edges = np.geomspace(r_min, r_max, n_rings + 1)
    r_mid = np.sqrt(r_edges[:-1] * r_edges[1:])
    dr = np.diff(r_edges)

    v_z_mag = theta_tip**2 * V * h / r_mid           # Eq. 23
    A_mag = theta_tip**2 * V**2 / r_mid              # Eq. 21
    l_i = 2*np.pi * w_slender / np.sqrt(rho * A_mag / (3 * sigma))  # Eq. 22

    db = dr / h
    M_i = 2 * np.pi * rho * (h * theta_tip)**2 * db * c_slender    # Eq. 24
    M_i *= (total_back_vol_mL * 1e-6 * rho) / M_i.sum()

    n_i = M_i / (rho * (4/3) * np.pi * (l_i/2)**3)
    return r_mid, l_i, v_z_mag, n_i, M_i

r_b, l_b, v_b, n_b, M_b = back_subfamilies()
total_drops_predicted = n_b.sum()
print(f"V0 = {V:.1f} m/s")
print(f"Total back-spatter drops predicted by model: {total_drops_predicted:,.0f}")
print(f"Total back-spatter mass: {M_b.sum()*1e3:.3f} g ({M_b.sum()/rho*1e6:.2f} mL)")
print(f"Drop diameter range: {l_b.min()*1e6:.0f} - {l_b.max()*1e6:.0f} um")

# ==================== SAMPLE DROPS ====================
N_SAMPLE = MAX_VISIBLE_STAINS
idx = np.random.choice(len(r_b), size=N_SAMPLE, p=n_b/n_b.sum())
l_s = l_b[idx]
u0_s = v_b[idx]

# Wide cone (theta_backward, with some scatter) into back half-space (-z direction)
theta_s = theta_backward + np.random.normal(0, 0.10, size=N_SAMPLE)
theta_s = np.clip(theta_s, 0.01, np.pi/2 - 0.01)
phi_s = np.random.uniform(0, 2*np.pi, size=N_SAMPLE)

# ==================== TRAJECTORY TO SHIRT (5 mm) ====================
def integrate_to_shirt(l_drops, u0s, thetas, phis, Z_stop=Z_shirt):
    """Integrate drops launched back (-z) until they hit the shirt at z=Z_stop.
    Returns: X, Y position on shirt; impact speed; impact angle from normal."""
    n = l_drops.size
    diam = l_drops
    m = rho * (np.pi/6.0) * diam**3
    A_cs = np.pi * (diam/2.0)**2

    vx = u0s * np.sin(thetas) * np.cos(phis)
    vy = u0s * np.sin(thetas) * np.sin(phis)
    vz = -u0s * np.cos(thetas)               # back toward shooter
    x = np.zeros(n); y = np.zeros(n); z = np.zeros(n)

    sp0 = np.sqrt(vx*vx + vy*vy + vz*vz)
    Re0 = rho_air * sp0 * diam / mu_air
    Cd0 = (0.28 + 6/np.sqrt(np.maximum(Re0,1e-3)) + 21/np.maximum(Re0,1e-3)) * WAKE_FACTOR
    stop_t = m / (0.5 * Cd0 * rho_air * A_cs * np.maximum(sp0, 1.0))
    dt_per = np.clip(0.02 * stop_t, 1e-8, 5e-4)

    alive = np.ones(n, dtype=bool)
    hit = np.zeros(n, dtype=bool)
    X = np.full(n, np.nan); Y = np.full(n, np.nan)
    impact_angle = np.full(n, np.nan)    # angle from surface normal (0 = head-on)
    impact_speed = np.full(n, np.nan)

    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/np.sqrt(Re_c) + 21/Re_c) * WAKE_FACTOR
        Fm = 0.5 * Cd * rho_air * A_cs * sp * sp / m
        return (-Fm*vxv/sp_safe, -g - Fm*vyv/sp_safe, -Fm*vzv/sp_safe)

    max_iters = 50000
    for _ 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)
        x_new = x + dt*vxm; y_new = y + dt*vym; z_new = z + dt*vzm
        vx_new = vx + dt*ax2; vy_new = vy + dt*ay2; vz_new = vz + dt*az2

        # back drops moving in -z, crossing z = Z_shirt (negative)
        crossed = alive & (z_new <= Z_stop) & (z > Z_stop)
        if crossed.any():
            frac = (z[crossed] - Z_stop) / np.maximum(z[crossed] - z_new[crossed], 1e-12)
            X[crossed] = x[crossed] + frac * (x_new[crossed] - x[crossed])
            Y[crossed] = y[crossed] + frac * (y_new[crossed] - y[crossed])
            sp = np.sqrt(vx_new[crossed]**2 + vy_new[crossed]**2 + vz_new[crossed]**2)
            impact_speed[crossed] = sp
            # angle from surface normal (z-axis): cos(angle) = |vz|/speed
            impact_angle[crossed] = np.arccos(np.minimum(np.abs(vz_new[crossed]) / np.maximum(sp, 1e-9), 1.0))
            hit[crossed] = True
            alive[crossed] = False

        # Drops that stalled before reaching shirt - mark as missed
        stalled = alive & (vz_new > -0.01)
        alive[stalled] = False

        x, y, z = x_new, y_new, z_new
        vx, vy, vz = vx_new, vy_new, vz_new

    return X, Y, hit, impact_speed, impact_angle

print("\nIntegrating back-spatter trajectories to shirt...")
X, Y, hit_mask, V_imp, ang_imp = integrate_to_shirt(l_s, u0_s, theta_s, phi_s)
print(f"  hits: {hit_mask.sum()} / {N_SAMPLE} ({100*hit_mask.mean():.1f}%)")

# Keep only hits
X = X[hit_mask]; Y = Y[hit_mask]
l_hit = l_s[hit_mask]
V_hit = V_imp[hit_mask]
ang_hit = ang_imp[hit_mask]
N_hit = len(X)

# ==================== STAIN SIZE & SHAPE ====================
# Eq. (29-30) of 2017 paper:
#   We = rho l v^2 / sigma     (impact Weber number)
#   Oh = mu / sqrt(rho l sigma) (Ohnesorge)
#   xi = 0.61 (We/Oh)^0.166     (spread factor)
# stain diameter = xi * l_drop
# For oblique impact, stain elongates by 1/sin(alpha) where alpha is angle from surface.
# Our impact_angle is measured from surface normal, so the elongation factor is
# 1/cos(impact_angle)  (more oblique = longer stain).
mu_blood = 3.7e-3   # Pa.s (from paper Table II ~3.7 cP)
We = rho * l_hit * V_hit**2 / sigma
Oh = mu_blood / np.sqrt(rho * l_hit * sigma)
xi = 0.61 * (We / Oh)**0.166
stain_dia = xi * l_hit                          # along-surface diameter
elongation = 1.0 / np.maximum(np.cos(ang_hit), 0.1)
stain_major = stain_dia * np.sqrt(elongation)   # major axis
stain_minor = stain_dia / np.sqrt(elongation)   # minor axis

# Stain orientation: aligned with the in-plane velocity direction at impact
# (For oblique impact, the drop streaks in the direction it was moving along surface.)
# In-plane direction from origin to impact point gives the radial direction;
# elongation is along that radial direction (drops moving outward radially).
stain_angle = np.arctan2(Y, X)

print(f"\nStain statistics:")
print(f"  Stain diameter range: {stain_dia.min()*1e6:.0f} - {stain_dia.max()*1e6:.0f} um")
print(f"  Median stain diameter: {np.median(stain_dia)*1e6:.0f} um")
print(f"  Impact speeds: {V_hit.min():.1f} - {V_hit.max():.1f} m/s")
r_imp = np.sqrt(X**2 + Y**2)
print(f"  Radial spread on shirt: median={np.median(r_imp)*1000:.1f} mm, max={r_imp.max()*1000:.1f} mm")

# ==================== RENDER SHIRT FACE ====================
fig, ax = plt.subplots(1, 1, figsize=(10, 10))

# Off-white shirt fabric background
shirt_color = '#f4ede1'   # cream / off-white
ax.set_facecolor(shirt_color)
fig.patch.set_facecolor('#1a1a1a')

# Subtle fabric texture: random small light gray dots
np.random.seed(7)
fab_x = np.random.uniform(-SHIRT_VIEW_HALF, SHIRT_VIEW_HALF, 2000)
fab_y = np.random.uniform(-SHIRT_VIEW_HALF, SHIRT_VIEW_HALF, 2000)
ax.scatter(fab_x, fab_y, s=0.4, c='#dccfb6', alpha=0.4, marker='.', linewidths=0)

# Build stain patches as ellipses
patches = []
# Color blood stains with slight variation in red/dark-red based on speed
blood_colors = []
for i in range(N_hit):
    # Convert to cm for display
    e = Ellipse((X[i]*100, Y[i]*100),
                width=stain_major[i]*100*2,   # diameter, not radius
                height=stain_minor[i]*100*2,
                angle=np.degrees(stain_angle[i]))
    patches.append(e)
    # Faster impact = darker, more saturated red (more blood deposited)
    v_norm = np.clip((V_hit[i] - 5) / 100, 0, 1)
    # blood red gradient: dark crimson to bright red
    r_c = 0.55 + 0.30 * v_norm
    g_c = 0.05 + 0.05 * v_norm
    b_c = 0.08 + 0.05 * v_norm
    blood_colors.append((r_c, g_c, b_c, 0.85))

pc = PatchCollection(patches, facecolors=blood_colors, edgecolors='none', linewidths=0)
ax.add_collection(pc)

# Central bullet hole - black circle at origin (radius ~ a = 3.8 mm)
hole = Circle((0, 0), a*100, facecolor='#0a0a0a', edgecolor='#000', linewidth=0.5, zorder=10)
ax.add_patch(hole)
# Dark scorching/halo around hole
for ring_r, alpha in [(a*100*1.8, 0.30), (a*100*1.4, 0.50)]:
    ring = Circle((0, 0), ring_r, facecolor='#2a1010', edgecolor='none', alpha=alpha, zorder=9)
    ax.add_patch(ring)

# Scale bar lower right
sb_x0, sb_y0 = SHIRT_VIEW_HALF*100 - 2.5, -SHIRT_VIEW_HALF*100 + 0.8
ax.plot([sb_x0, sb_x0 + 2.0], [sb_y0, sb_y0], color='#222', lw=3, zorder=20)
ax.text(sb_x0 + 1.0, sb_y0 + 0.3, '2 cm', ha='center', va='bottom',
        fontsize=11, color='#222', weight='bold', zorder=20)

# Title and labels
ax.set_xlim(-SHIRT_VIEW_HALF*100, SHIRT_VIEW_HALF*100)
ax.set_ylim(-SHIRT_VIEW_HALF*100, SHIRT_VIEW_HALF*100)
ax.set_aspect('equal')
ax.set_xticks([])
ax.set_yticks([])
for spine in ax.spines.values():
    spine.set_color('#444'); spine.set_linewidth(2)

fig.text(0.5, 0.96, 'Entry-side shirt face — predicted back-spatter pattern',
         ha='center', fontsize=15, color='white', weight='bold')
fig.text(0.5, 0.93,
         f'.30 cal @ 800 m/s, lateral neck/jugular impact, shirt 5 mm in front of entry wound',
         ha='center', fontsize=11, color='#bbb')
fig.text(0.5, 0.025,
         f'{N_hit:,} stains rendered (model predicts ~{total_drops_predicted/1e6:.1f}M drops total). '
         f'Stain Ø range: {stain_dia.min()*1e6:.0f}–{stain_dia.max()*1e6:.0f} μm. '
         f'Central hole: bullet perforation (Ø {2*a*1000:.1f} mm).\n'
         f'Comiskey/Yarin/Kim/Attinger 2016 (slender-bullet back spatter, Eqs. 21–24) + 2017 stain spread (Eq. 29). '
         f'Muzzle gases NOT modeled.',
         ha='center', fontsize=9, color='#aaa', style='italic')

plt.subplots_adjust(left=0.05, right=0.95, top=0.91, bottom=0.10)
out_png = '/home/claude/shirt_face_2d.png'
plt.savefig(out_png, dpi=180, bbox_inches='tight', facecolor='#1a1a1a')
print(f"\nSaved: {out_png}")

# Also save a higher-resolution version for forensic record
out_png_hi = '/home/claude/shirt_face_2d_hires.png'
plt.savefig(out_png_hi, dpi=300, bbox_inches='tight', facecolor='#1a1a1a')
print(f"Hi-res: {out_png_hi}")