# =====================================================================
# Companion code for "The Trap Number T: A Dimensionless Criterion for
# Boulder-Boulder Gravitational Dynamics on Near-Earth Rubble-Pile Asteroid
# Surfaces" (prepared for submission).
#
# Implements Equation 5 / Equation 8:
#     T = G * M_local / (g_surface * d_eff^2)
# This is the downslope force ratio  T = G*M_local/(g_surface*sin(alpha)*d_eff^2)
# (Eqs 3-5) evaluated at the STEEPEST-SLOPE LIMIT sin(alpha) = 1 (manuscript
# Section 2.2 / Section 3.1), a CONSERVATIVE LOWER BOUND: the slope-aware value
# is larger by 1/sin(alpha), and a trap (T > 1) occurs where the mapped value
# exceeds sin(alpha). Equivalent scaling (Eq 6): T = (M_local/M_ast)(R/d_eff)^2.
#
# Surface grid: 60 latitudes x 120 longitudes = 7,200 points (Section 3.1).
# Boulder mass: sphere, m = (4/3) pi rho_b (d/2)^3. Parameters per Section 3.1.
# THIS FILE: Eq 8 on the REAL Ryugu catalogue (Ray et al. 2026, 49,536 boulders). Headline
#   T_max ~= 0.5 (0.34 shape-corrected to 0.67 long-axis) at Otohime (Sec 3.2-3.3);
#   d_eff sweep 30/50/75/100 m -> 1.9/0.67/0.31/0.19 (Section 3.6).
# =====================================================================
"""
Trap number on the REAL Ryugu boulder catalogue (Ray et al. 2026).

Reads the published catalogue 'Ryugu_Boulder_Attributes.csv' (download from the
figshare deposit cited in the paper; columns Longitude_X, Latitude_Y, Length_m)
and computes the trap number directly from the observed boulders. It reproduces
every real-data number reported in the manuscript:
  - T_max at Otohime: 0.67 (long axis as diameter) / 0.34 (shape-corrected); best ~= 0.5
  - d_eff sensitivity (T_max): 1.9, 0.67, 0.31, 0.19 at d_eff = 30, 50, 75, 100 m
  - < 1% of the surface with T > 0.1; nowhere T > 1 at d_eff = 50 m
This is the authoritative real-data script for the paper.

Requires numpy and scipy (matplotlib only if SAVE_FIGURE = True).
"""
import csv, numpy as np
from scipy.spatial import cKDTree

CATALOGUE = 'Ryugu_Boulder_Attributes.csv'   # path to the downloaded Ray et al. (2026) file
SAVE_FIGURE = False                            # True also writes the Otohime trap-number map (Fig. 2)

# Ryugu parameters (identical to the manuscript and the supplementary workbook)
G=6.674e-11; M_AST=4.50e11; RHO_B=1600.0; D_EFF=50.0
A_ELL=502.0; C_ELL=437.744                     # Ryugu ellipsoid semi-axes (Ray et al. 2026)
g_surface=1.48e-4                              # = G*M/R^2 with R = 450 m

def load(path):
    lon=[];lat=[];L=[]
    with open(path,newline='') as f:
        for row in csv.DictReader(f):
            try: lon.append(float(row['Longitude_X']));lat.append(float(row['Latitude_Y']));L.append(float(row['Length_m']))
            except (ValueError,TypeError,KeyError): pass
    return np.array(lon)%360.0, np.array(lat), np.array(L)   # normalise longitudes to [0,360)

def to_xyz(lon,lat):
    lo=np.radians(lon); la=np.radians(lat)
    r=1.0/np.sqrt((np.cos(la)/A_ELL)**2+(np.sin(la)/C_ELL)**2)   # ellipsoid radius at each lat
    return np.column_stack([r*np.cos(la)*np.cos(lo), r*np.cos(la)*np.sin(lo), r*np.sin(la)])

def trap(pos,tree,mass,d_eff):
    return np.array([G*mass[nb].sum()/(g_surface*d_eff**2) for nb in tree.query_ball_point(pos,d_eff)])

def main():
    lon,lat,L=load(CATALOGUE)
    pos=to_xyz(lon,lat); tree=cKDTree(pos)
    print(f"boulders read: {len(L)}   largest: {L.max():.1f} m")

    print("\nT_max at Otohime (d_eff = 50 m):")
    for tag,sf in [("long-axis as diameter (upper)",1.0),("volume-equivalent sphere (lower)",0.8)]:
        m=RHO_B*(4/3)*np.pi*(sf*L/2)**3; T=trap(pos,tree,m,D_EFF); i=int(np.argmax(T))
        print(f"  {tag:34s}: T_max = {T.max():.3f}  at lon {lon[i]:.0f}, lat {lat[i]:.0f}   "
              f"(>0.1: {100*np.mean(T>0.1):.2f}%, >1: {100*np.mean(T>1):.2f}%)")

    print("\nd_eff sensitivity (T_max), long axis as diameter:")
    m=RHO_B*(4/3)*np.pi*(L/2)**3
    print("   " + "  ".join(f"{d}m={trap(pos,tree,m,d).max():.2f}" for d in (30,50,75,100)))
    print("   (shape-corrected values are about half of these)")
    print("\nBest estimate at Otohime: T ~= 0.5 (range 0.34 to 0.67); Ryugu is sub-threshold (T<1) at d_eff = 50 m.")

    if SAVE_FIGURE:
        save_otohime_map(lon,lat,L,pos,tree)

def save_otohime_map(lon,lat,L,pos,tree):
    import matplotlib; matplotlib.use('Agg'); import matplotlib.pyplot as plt
    m=RHO_B*(4/3)*np.pi*(L/2)**3
    j=int(np.argmax(L)); ctr=pos[j]; up=ctr/np.linalg.norm(ctr)
    east=np.cross([0,0,1.0],up); east/=np.linalg.norm(east); north=np.cross(up,east)
    near=np.where(np.linalg.norm(pos-ctr,axis=1)<260)[0]
    pe=(pos[near]-ctr)@east; pn=(pos[near]-ctr)@north
    gx=np.linspace(-160,160,140); GX,GY=np.meshgrid(gx,gx)
    gp=(ctr[None,None,:]+GX[...,None]*east+GY[...,None]*north).reshape(-1,3)
    Tg=np.array([G*m[nb].sum()/(g_surface*D_EFF**2) for nb in tree.query_ball_point(gp,D_EFF)]).reshape(GX.shape)
    fig,ax=plt.subplots(figsize=(6.2,5.2),dpi=300)
    pcm=ax.pcolormesh(GX,GY,Tg,shading='auto',cmap='magma',vmin=0,vmax=max(0.7,Tg.max()))
    fig.colorbar(pcm,ax=ax,fraction=0.046,pad=0.04).set_label(r'Trap number $\mathcal{T}$ (real boulders)')
    ax.scatter(pe,pn,s=np.clip(L[near]*2.2,4,260),facecolors='none',edgecolors='cyan',linewidths=0.7,alpha=0.8)
    ax.scatter([0],[0],s=L[j]*2.2,facecolors='none',edgecolors='white',linewidths=1.8)
    ax.set_xlabel('East offset from Otohime (m)'); ax.set_ylabel('North offset (m)'); ax.set_aspect('equal')
    fig.savefig('Fig2_Otohime_real.png',bbox_inches='tight'); plt.close()
    print("  saved Fig2_Otohime_real.png")

if __name__=='__main__':
    main()
