# =====================================================================
# 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: Generic driver applying Equation 8 to any catalogue (reads ryugu_boulders.csv).
# =====================================================================
"""
Compute the trap number T on a REAL boulder catalogue (observed boulders),
using the same definition as the paper:

        T = G * M_local / (g_surface * d_eff^2)

where M_local is the total mass of catalogued boulders lying within a distance
d_eff of a given surface point. T is evaluated at every catalogued boulder
location, and the maximum is reported (this is where traps are strongest,
i.e. next to the biggest blocks such as Otohime).

INPUT: a CSV file with one row per boulder and columns for
   - longitude  (degrees East)
   - latitude   (degrees)
   - size       (long-axis length OR diameter, in metres)
Edit the CONFIG block below so the three column names match your file.

Only numpy is required. If scipy is installed it runs much faster.

NOTE: For Ryugu specifically, code_real_ryugu_trap_number.py is the authoritative
script: it uses the Ryugu ellipsoid and reproduces the manuscript numbers exactly.
This generic adapter uses a sphere of mean radius and works for any catalogue.
"""

import numpy as np, csv, sys

# ----------------------------- CONFIG ----------------------------------
CATALOGUE_CSV = 'ryugu_boulders.csv'   # path to the real catalogue file
COL_LON  = 'longitude'                 # column name: longitude (deg E)
COL_LAT  = 'latitude'                  # column name: latitude  (deg)
COL_SIZE = 'length_m'                  # column name: long-axis length (m)

# Asteroid parameters (defaults below are for Ryugu)
G        = 6.674e-11    # gravitational constant (SI)
M_AST    = 4.50e11      # asteroid mass            (kg)
R_AST    = 450.0        # mean radius              (m)  [Ryugu]
RHO_B    = 1600.0       # boulder density          (kg/m^3)
D_EFF    = 50.0         # interaction scale        (m)
SHAPE_FACTOR = 1.0      # diameter = SHAPE_FACTOR * size  (use ~0.8 if 'size'
                        # is a long axis and you want an equivalent diameter)
# -----------------------------------------------------------------------

g_surface = G * M_AST / R_AST**2

def load(path):
    with open(path, newline='') as f:
        rdr = csv.DictReader(f)
        cols = rdr.fieldnames
        for need in (COL_LON, COL_LAT, COL_SIZE):
            if need not in cols:
                sys.exit(f"Column '{need}' not found. Columns in your file are:\n  {cols}\n"
                         f"Edit the CONFIG block to use the correct names.")
        lon, lat, size = [], [], []
        for row in rdr:
            try:
                lon.append(float(row[COL_LON])); lat.append(float(row[COL_LAT]))
                size.append(float(row[COL_SIZE]))
            except (ValueError, TypeError):
                continue   # skip blank/non-numeric rows
    return np.array(lon), np.array(lat), np.array(size)

def to_xyz(lon_deg, lat_deg, R):
    lon = np.radians(lon_deg); lat = np.radians(lat_deg)
    return np.column_stack([R*np.cos(lat)*np.cos(lon),
                            R*np.cos(lat)*np.sin(lon),
                            R*np.sin(lat)])

def main():
    lon, lat, size = load(CATALOGUE_CSV)
    diam = SHAPE_FACTOR * size
    mass = RHO_B * (4.0/3.0)*np.pi*(diam/2.0)**3      # sphere of that diameter
    pos  = to_xyz(lon, lat, R_AST)
    n = len(mass)
    T = np.zeros(n)

    try:
        from scipy.spatial import cKDTree
        tree = cKDTree(pos)
        for i, nbrs in enumerate(tree.query_ball_point(pos, D_EFF)):
            T[i] = G*mass[nbrs].sum()/(g_surface*D_EFF**2)
    except Exception:
        print("(scipy not found; using the slower numpy method, please wait)")
        for i in range(n):
            d = np.sqrt(((pos - pos[i])**2).sum(axis=1))
            T[i] = G*mass[d <= D_EFF].sum()/(g_surface*D_EFF**2)

    i = int(np.argmax(T))
    print(f"\nReal-catalogue trap number  (d_eff = {D_EFF:.0f} m, rho_b = {RHO_B:.0f} kg/m3)")
    print(f"  boulders read         : {n}")
    print(f"  surface gravity       : {g_surface:.3e} m/s^2")
    print(f"  T_max                 : {T.max():.3f}")
    print(f"  location of T_max     : lon {lon[i]:.1f}E, lat {lat[i]:.1f}, "
          f"largest local block {size[i]:.1f} m")
    print(f"  sites with T > 0.1    : {100*np.mean(T>0.1):.2f} %")
    print(f"  sites with T > 1      : {100*np.mean(T>1.0):.2f} %  (traps form)")

if __name__ == '__main__':
    main()
