#!/usr/bin/env python3
# =====================================================================
# 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: Core model and Eq 5/8 on the 60x120 grid. Synthetic-model trap numbers (Section 3.2, Table 1).
# =====================================================================
"""
================================================================================
trap_number.py
Core computation of the trap number 𝒯 for rubble-pile asteroid surfaces.

Companion code for:
  "The Trap Number 𝒯: A Dimensionless Criterion for Boulder-Boulder
   Gravitational Dynamics on Near-Earth Rubble-Pile Asteroid Surfaces"
  Prepared for submission

The trap number is defined as:
    𝒯 = G * M_local / (g_surface * d_eff^2)
where:
    G        = gravitational constant (6.674e-11 N m^2 / kg^2)
    M_local  = total boulder mass within distance d_eff of a surface point (kg)
    g_surface = GM/R^2, the asteroid surface gravity (m/s^2)
    d_eff    = interaction radius (m), default 50 m

When 𝒯 > 1, boulder pull dominates the downslope drive and material tends to
accumulate near boulder clusters (a convergence zone); 𝒯 = 1 is a nominal
order-unity scale rather than a proven saddle-to-minimum bifurcation
(see manuscript Section 3.5). When 𝒯 << 1, the standard geopotential model suffices.

Usage:
    python trap_number.py
================================================================================
"""

import numpy as np

# Physical constant
G = 6.674e-11  # N m^2 / kg^2

# ------------------------------------------------------------------
# Asteroid physical parameters (from published literature)
# ------------------------------------------------------------------
ASTEROIDS = {
    'Ryugu': {
        'R': 450,          # mean radius (m)          [Watanabe et al. 2019]
        'M': 4.5e11,       # mass (kg)                [Watanabe et al. 2019]
        'g': 1.48e-4,      # surface gravity (m/s^2)
        'rho_b': 1600,     # boulder density (kg/m^3)
        'N': 4400,         # number of boulders modeled
        'd_min': 5,        # min boulder diameter (m)
        'd_max': 160,      # max boulder diameter (m)
        'alpha': -2.665,   # cumulative SFD power-law slope [Ray et al. 2026]
        'named': [         # (lat, lon, diameter, name) for named boulders
            (-80, 180, 140, 'Otohime'),
            (-30, 50, 60, 'Urashima'),
            (4, 296, 50, 'Cendrillon'),
            (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,
        'named': [(-40, 156, 51, 'Roc'), (-55, 310, 43, 'Simurgh'),
                  (23, 42, 21, 'BenBen')],
    },
    'Itokawa': {
        'R': 165, 'M': 3.51e10, 'g': 8.6e-5, 'rho_b': 2500,
        'N': 800, 'd_min': 1, 'd_max': 50, 'alpha': -3.05,
        'named': [(10, 250, 50, 'Yoshinodai')],
    },
}


def sample_powerlaw(N, d_min, d_max, alpha, rng):
    """Sample N boulder diameters from a cumulative power law of slope alpha."""
    beta = alpha - 1
    u = rng.uniform(0, 1, N)
    d = ((d_max**(beta + 1) - d_min**(beta + 1)) * u + d_min**(beta + 1))**(1 / (beta + 1))
    return np.clip(d, d_min, d_max)


def generate_boulders(params, seed=0):
    """
    Generate a synthetic boulder population for an asteroid.
    Named boulders are placed deterministically at their documented positions;
    the remaining population is sampled from the published size distribution.
    Returns: diameters (m), masses (kg), latitudes (deg), longitudes (deg).
    """
    rng = np.random.default_rng(seed)
    N = params['N']
    d = sample_powerlaw(N, params['d_min'], params['d_max'], params['alpha'], rng)
    lat = np.degrees(np.arcsin(rng.uniform(-1, 1, N)))
    lon = rng.uniform(0, 360, N)
    # Fix named boulders (these are real, documented features)
    for i, (la, lo, di, name) in enumerate(params['named']):
        if i < N:
            lat[i], lon[i], d[i] = la, lo, di
    mass = (4.0 / 3.0) * np.pi * params['rho_b'] * (d / 2)**3
    return d, mass, lat, lon


def latlon_to_xyz(lat, lon, R):
    """Convert latitude/longitude (deg) on a sphere of radius R to Cartesian."""
    la, lo = np.radians(lat), np.radians(lon)
    return np.array([R * np.cos(la) * np.cos(lo),
                     R * np.cos(la) * np.sin(lo),
                     R * np.sin(la)])


def compute_trap_number_map(params, lat_grid, lon_grid, d_eff=50.0, seed=0):
    """
    Compute the trap number 𝒯 on a latitude/longitude grid.
    For each grid point, sum the masses of all boulders whose centres lie
    within d_eff, then apply 𝒯 = G * M_local / (g * d_eff^2).
    Returns: 2D array of 𝒯 values, and the (lat, lon) of the maximum.
    """
    d, mass, blat, blon = generate_boulders(params, seed)
    R, g = params['R'], params['g']
    bpos = np.array([latlon_to_xyz(blat[i], blon[i], R) for i in range(len(blat))])

    Tmap = np.zeros((len(lat_grid), len(lon_grid)))
    for i, la in enumerate(lat_grid):
        for j, lo in enumerate(lon_grid):
            P = latlon_to_xyz(la, lo, R)
            dist = np.sqrt(np.sum((bpos - P)**2, axis=1))
            M_local = mass[dist < d_eff].sum()
            Tmap[i, j] = G * M_local / (g * d_eff**2)

    imax = np.unravel_index(np.argmax(Tmap), Tmap.shape)
    return Tmap, (lat_grid[imax[0]], lon_grid[imax[1]])


if __name__ == '__main__':
    lat_grid = np.linspace(-85, 85, 60)
    lon_grid = np.linspace(2.5, 357.5, 120)

    print(f"{'Asteroid':<10}{'T_max':>10}{'T_mean':>10}{'%>0.1':>10}")
    print("-" * 40)
    for name, params in ASTEROIDS.items():
        Tmap, loc = compute_trap_number_map(params, lat_grid, lon_grid)
        pct = 100 * np.mean(Tmap > 0.1)
        print(f"{name:<10}{Tmap.max():>10.3f}{Tmap.mean():>10.3f}{pct:>10.1f}")
