# =====================================================================
# 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: Generates manuscript Figures 1-6 (saved into this same folder).
# =====================================================================
"""
Regenerate figures for the revised manuscript.
Generates the manuscript figures. Figs 1, 3, 4, 5, 6 use the synthetic model
(synthetic Ryugu value 0.44; the headline Ryugu value is the real ~0.5). Fig 2
is the real-catalogue 3-panel trap map and is built by the dedicated module
code_fig2_trap_map.py (imported below), so this all-in-one run and the paper
figure are identical.
"""
import numpy as np
import matplotlib
matplotlib.use('Agg')
import matplotlib.pyplot as plt
from matplotlib.colors import LinearSegmentedColormap
import matplotlib.patheffects as pe

np.random.seed(42)
G = 6.674e-11
import os
outdir = os.path.join(os.path.dirname(os.path.abspath(__file__)), '')
print('Figures will be saved in this folder:', outdir)

plt.rcParams.update({
    'font.family':'sans-serif','font.size':10,'axes.linewidth':0.6,
    'figure.dpi':200,'savefig.dpi':300,'savefig.bbox':'tight',
    'savefig.pad_inches':0.18,'axes.spines.top':False,'axes.spines.right':False,
    'mathtext.fontset':'cm'
})

# Symbol forms:
#   TSYM  : the bare symbol (use alone or before a subscript)
#   TW    : symbol followed by a word -> thin space baked inside the math so it never glues
TSYM = r'$\mathcal{T}$'
TW   = r'$\mathcal{T}\,$' + ' '       # symbol immediately before a word
def TEQ(v):                          # render "𝒯 = value" as a single math group (proper spacing around =)
    return r'$\mathcal{T}\;=\;%s$' % v

C_BLUE='#2563EB'; C_PURPLE='#7C3AED'; C_CORAL='#DC2626'; C_TEAL='#059669'; C_AMBER='#D97706'; C_GRAY='#6B7280'

asteroids = {
    'Ryugu':{'R':450,'M':4.5e11,'g':1.48e-4,'rho_b':1600,'N':4400,'d_min':5,'d_max':160,'alpha':-2.665,
             'lat_bias':lambda lat:0.6 if abs(lat)<15 else(1.4 if abs(lat)<50 else 1.0),
             'named':[(-80,180,140,'Otohime'),(4,296,50,'Cendrillon'),(-30,50,60,'Urashima'),(20,100,40,'Kintaro')]},
    'Bennu':{'R':245,'M':7.329e10,'g':8.2e-5,'rho_b':1600,'N':1500,'d_min':2,'d_max':60,'alpha':-2.9,
             'lat_bias':lambda lat:1.0+0.8*abs(lat)/90,
             'named':[(-40,156,51,'Roc'),(23,42,21,'BenBen'),(-55,310,43,'Simurgh')]},
    'Itokawa':{'R':165,'M':3.51e10,'g':8.6e-5,'rho_b':2500,'N':800,'d_min':1,'d_max':50,'alpha':-3.05,
               'lat_bias':lambda lat:1.5 if abs(lat)>30 else 0.3,
               'named':[(10,250,50,'Yoshinodai')]},
}

def sample_pl(N,dmin,dmax,a):
    b=a-1; u=np.random.uniform(0,1,N)
    d=((dmax**(b+1)-dmin**(b+1))*u+dmin**(b+1))**(1/(b+1))
    return np.clip(d,dmin,dmax)

def gen(p):
    N=p['N']; d=sample_pl(N,p['d_min'],p['d_max'],p['alpha'])
    lats,lons=[],[]
    while len(lats)<N:
        la=np.degrees(np.arcsin(np.random.uniform(-1,1)))
        if np.random.uniform()<p['lat_bias'](la)/1.5:
            lats.append(la); lons.append(np.random.uniform(0,360))
    lats=np.array(lats[:N]); lons=np.array(lons[:N])
    for i,(la,lo,di,nm) in enumerate(p['named']):
        if i<N: lats[i]=la; lons[i]=lo; d[i]=di
    m=(4/3)*np.pi*p['rho_b']*(d/2)**3
    return d,m,lats,lons

def Tmap(p,blat,blon,m,d,latg,long,d_eff=50):
    R,g=p['R'],p['g']
    out=np.zeros((len(latg),len(long)))
    for i,la in enumerate(latg):
        for j,lo in enumerate(long):
            cd=(np.sin(np.radians(la))*np.sin(np.radians(blat))+
                np.cos(np.radians(la))*np.cos(np.radians(blat))*np.cos(np.radians(lo-blon)))
            dist=R*np.arccos(np.clip(cd,-1,1))
            Ml=m[dist<d_eff].sum()
            out[i,j]=G*Ml/(g*d_eff**2)
    return out

latg=np.linspace(-85,85,60); long=np.linspace(2.5,357.5,120)
data={}
for nm,p in asteroids.items():
    d,m,la,lo=gen(p)
    data[nm]={'d':d,'m':m,'la':la,'lo':lo,'T':Tmap(p,la,lo,m,d,latg,long),'p':p}
print("Computed T maps")

cmap=LinearSegmentedColormap.from_list('g',['#F0F9FF','#93C5FD','#3B82F6','#1D4ED8','#FCD34D','#F59E0B','#DC2626','#7F1D1D'])

# ---------- FIGURE 1: surface maps ----------
fig,axes=plt.subplots(3,1,figsize=(7.5,10.8),constrained_layout=True)
titles={'Ryugu':r'(162173) Ryugu:  $\mathcal{T}_{max}\;=\;0.44 \pm 0.05$',
        'Bennu':r'(101955) Bennu:  $\mathcal{T}_{max}\;=\;0.038 \pm 0.006$',
        'Itokawa':r'(25143) Itokawa:  $\mathcal{T}_{max}\;=\;0.051 \pm 0.001$'}
for idx,nm in enumerate(['Ryugu','Bennu','Itokawa']):
    ax=axes[idx]; Tm=data[nm]['T']; LON,LAT=np.meshgrid(long,latg)
    vmax=max(0.1,Tm.max()*1.05)
    c=ax.pcolormesh(LON,LAT,Tm,cmap=cmap,shading='gouraud',vmin=0,vmax=vmax)
    fig.colorbar(c,ax=ax,shrink=0.85,label=TSYM,aspect=28)
    for la,lo,di,nmb in data[nm]['p']['named']:
        ms=max(5,min(10,di/14))
        ax.plot(lo,la,'^',color='white',markersize=ms,markeredgecolor='black',markeredgewidth=0.6,zorder=5)
        # keep label inside the panel and offset clearly from the marker
        dx,dy = 16,11
        if lo>300: dx=-16
        ha='left' if dx>0 else 'right'
        ax.annotate(nmb,xy=(lo,la),xytext=(lo+dx,la+dy),textcoords='data',fontsize=7.5,color='white',ha=ha,va='center',
                    path_effects=[pe.withStroke(linewidth=1.8,foreground='black')],zorder=6)
    if Tm.max()>0.1:
        try:
            cs=ax.contour(LON,LAT,Tm,levels=[0.1],colors=['white'],linewidths=0.9,linestyles='--')
            ax.clabel(cs,fmt=r'$\mathcal{T}\;=\;0.1$',fontsize=7,colors='white')
        except: pass
    ax.set_title(titles[nm],fontsize=10.5,pad=6)
    ax.set_xlabel('Longitude (deg)'); ax.set_ylabel('Latitude (deg)')
    ax.set_xlim(0,360); ax.set_ylim(-85,85)
fig.suptitle('Trap number '+TW+r'across three near-Earth asteroids ($d_{\mathrm{eff}}$ = 50 m)',fontsize=12.5)
fig.savefig(outdir+'Fig1_surface_maps.png'); plt.close()
print("Fig 1 done")

# ---------- FIGURE 2: real-catalogue trap map (3-panel) ----------
# Generated by the dedicated module code_fig2_trap_map.py so the all-in-one
# run and the manuscript figure are identical (field map + zoom + 1/r^2 panel).
from code_fig2_trap_map import build_fig2
_f2 = build_fig2(os.path.join(outdir, 'Fig2.png'))
print("Fig 2 done ->", _f2)

# ---------- FIGURE 3: T vs Bond number classification ----------
fig,ax=plt.subplots(figsize=(7.2,5.4))
ax.set_xlim(0.001,100); ax.set_ylim(0.001,100); ax.set_xscale('log'); ax.set_yscale('log')
ax.fill_between([0.001,1],0.001,1,alpha=0.06,color=C_GRAY)
ax.fill_between([1,100],0.001,1,alpha=0.08,color=C_PURPLE)
ax.fill_between([0.001,1],1,100,alpha=0.08,color=C_BLUE)
ax.fill_between([1,100],1,100,alpha=0.08,color=C_TEAL)
ax.axhline(1,color=C_GRAY,lw=0.8,ls='--'); ax.axvline(1,color=C_GRAY,lw=0.8,ls='--')
# quadrant labels placed in clear corners, away from data points
ax.text(0.0055,40,'Gravity-dominated\ntrapping',fontsize=10,ha='center',color=C_BLUE,fontweight='bold')
ax.text(45,40,'Both mechanisms\nactive',fontsize=10,ha='center',color=C_TEAL,fontweight='bold')
ax.text(55,0.0022,'Cohesion-dominated\ntrapping',fontsize=10,ha='center',color=C_PURPLE,fontweight='bold')
ax.text(0.0055,0.0022,'Standard\ngeopotential\nsufficient',fontsize=9,ha='center',color=C_GRAY,fontweight='bold')
# data points with labels spaced so they do not collide
ax.plot(0.3,0.44,'s',color=C_TEAL,markersize=11,markeredgecolor='white',markeredgewidth=1.5,zorder=5)
ax.annotate('Ryugu\n(Otohime)',xy=(0.3,0.44),xytext=(0.085,0.62),fontsize=8,color=C_TEAL,fontweight='bold',
            ha='center',arrowprops=dict(arrowstyle='->',color=C_TEAL,lw=0.7))
ax.plot(5,0.051,'s',color=C_CORAL,markersize=10,markeredgecolor='white',markeredgewidth=1.5,zorder=5)
ax.annotate('Itokawa',xy=(5,0.051),xytext=(2.7,0.16),fontsize=8,color=C_CORAL,fontweight='bold',
            ha='center',arrowprops=dict(arrowstyle='->',color=C_CORAL,lw=0.7))
ax.plot(10,0.038,'s',color=C_BLUE,markersize=11,markeredgecolor='white',markeredgewidth=1.5,zorder=5)
ax.annotate('Bennu',xy=(10,0.038),xytext=(22,0.022),fontsize=8,color=C_BLUE,fontweight='bold',
            ha='center',arrowprops=dict(arrowstyle='->',color=C_BLUE,lw=0.7))
ax.plot(8,2.3,'*',color=C_AMBER,markersize=15,markeredgecolor='white',markeredgewidth=1,zorder=5)
ax.annotate('Dimorphos\n(predicted)',xy=(8,2.3),xytext=(3.2,3.6),fontsize=8,color=C_AMBER,fontweight='bold',
            ha='center',arrowprops=dict(arrowstyle='->',color=C_AMBER,lw=0.7))
ax.plot(15,7,'*',color='#047857',markersize=15,markeredgecolor='white',markeredgewidth=1,zorder=5)
ax.annotate('50 m NEA\n(predicted)',xy=(15,7),xytext=(34,9.5),fontsize=8,color='#047857',fontweight='bold',
            ha='center',arrowprops=dict(arrowstyle='->',color='#047857',lw=0.7))
ax.set_xlabel(r'Cohesive Bond number $B_c$',fontsize=10)
ax.set_ylabel(r'Trap number $\mathcal{T}$',fontsize=10)
ax.set_title('Two-parameter trapping classification:  '+TW+r'versus $B_c$',fontsize=10)
fig.savefig(outdir+'Fig3_classification.png'); plt.close()
print("Fig 3 done")

# ---------- FIGURE 4: slope vs T_max ----------
fig,ax=plt.subplots(figsize=(7.2,4.7))
slopes=np.linspace(-4.0,-2.0,50); Tvals=[]
for s in slopes:
    dt=sample_pl(2000,5,160,s); mt=(4/3)*np.pi*1600*(dt/2)**3
    Ml=np.sort(mt)[::-1][:10].sum()
    Tvals.append(G*Ml/(1.48e-4*50**2))
ax.semilogy(slopes,Tvals,color=C_PURPLE,lw=2.5)
# explicit, non-overlapping label anchors for each marker
lab={'Ryugu':(-2.78,2.6),'Bennu':(-2.45,0.026),'Itokawa':(-3.72,0.0058)}
pts=[('Ryugu',-2.665,0.44,C_TEAL),('Bennu',-2.9,0.038,C_BLUE),('Itokawa',-3.05,0.051,C_CORAL)]
for nm,s,tm,col in pts:
    ax.plot(s,tm,'o',color=col,markersize=11,zorder=5,markeredgecolor='white',markeredgewidth=1.5)
    txt=f'{nm}\nslope = {s}\n'+TEQ(tm)
    ax.annotate(txt,xy=(s,tm),xytext=lab[nm],fontsize=8,color=col,fontweight='bold',
                ha='center',arrowprops=dict(arrowstyle='->',color=col,lw=0.9))
ax.axhline(1.0,color=C_CORAL,lw=1.2,ls='--',alpha=0.6); ax.text(-3.97,1.35,TEQ('1')+' (trap formation)',fontsize=8,color=C_CORAL)
ax.axhline(0.1,color=C_AMBER,lw=1,ls=':',alpha=0.6); ax.text(-3.97,0.135,TEQ('0.1')+' (detectable)',fontsize=8,color=C_AMBER)
ax.axvline(-2.7,color=C_GRAY,lw=0.8,ls=':',alpha=0.5); ax.text(-2.69,0.0016,'slope = -2.7\n(threshold)',fontsize=8,color=C_GRAY)
ax.set_xlabel('Cumulative power-law slope of boulder SFD',fontsize=11)
ax.set_ylabel(r'$\mathcal{T}_{max}$ (Ryugu-like body)',fontsize=11)
ax.set_xlim(-4,-2); ax.set_ylim(0.0008,5)
ax.set_title(r'Power-law slope controls the trap number: shallower slopes raise $\mathcal{T}$',fontsize=10)
ax.grid(True,alpha=0.12)
fig.savefig(outdir+'Fig4_slope_vs_T.png'); plt.close()
print("Fig 4 done")

# ---------- FIGURE 5: bar chart ----------
fig,ax=plt.subplots(figsize=(6.8,4))
bodies=['Itokawa','Bennu','Ryugu','Dimorphos\n(pred.)','50 m NEA\n(pred.)']
Ts=[0.051,0.038,0.44,2.3,6.8]; errs=[0.001,0.006,0.05,0,0]
colors=[C_CORAL,C_BLUE,C_TEAL,C_AMBER,'#047857']
bars=ax.bar(bodies,Ts,yerr=errs,color=colors,edgecolor='black',linewidth=0.6,width=0.6,capsize=4)
for bar,v in zip(bars,Ts):
    ax.text(bar.get_x()+bar.get_width()/2,v+0.15,f'{v:.2f}' if v<1 else f'{v:.1f}',ha='center',fontsize=9,fontweight='bold')
ax.axhline(1.0,color=C_CORAL,lw=1.5,ls='--',alpha=0.5); ax.text(4.5,1.18,TEQ('1'),fontsize=8,color=C_CORAL,ha='right')
ax.axhline(0.1,color=C_AMBER,lw=1,ls=':',alpha=0.5); ax.text(4.5,0.22,TEQ('0.1'),fontsize=8,color=C_AMBER,ha='right')
ax.axvspan(-0.5,2.5,alpha=0.04,color=C_BLUE); ax.axvspan(2.5,4.5,alpha=0.04,color=C_AMBER)
ax.text(1,7.7,'Visited',fontsize=9,ha='center',color=C_GRAY,style='italic')
ax.text(3.5,7.7,'Predicted',fontsize=9,ha='center',color=C_AMBER,style='italic')
ax.set_ylabel(r'$\mathcal{T}_{max}$',fontsize=11); ax.set_ylim(0,8.7)
ax.set_title('Maximum '+TW+'from measured to predicted NEOs',fontsize=10)
fig.savefig(outdir+'Fig5_comparison_bar.png'); plt.close()
print("Fig 5 done")

# ---------- FIGURE 6: three-panel stability diagram ----------
# minimum bowl (left) | saddle / Otohime (middle) | real point-mass eigenvalues,
# always a saddle, T=1 nominal only (right). Generated by the dedicated module
# so the manuscript figure and the code match.
from code_stability_diagram import build_figure as _build_stability
fig=_build_stability()
fig.savefig(outdir+'Fig6_stability.png',dpi=200,bbox_inches='tight'); plt.close()
print("Fig 6 (stability, 3-panel) done")
print("All 6 figures regenerated.")

import glob as _glob
print("\nSaved figure files:")
for _p in sorted(_glob.glob(os.path.join(outdir,'Fig*.png'))):
    print("  ", _p)
