# =============================================================================
# code_fig2_trap_map.py
# Generates manuscript FIGURE 2: the real-catalogue trap-number field around
# Otohime on Ryugu, in three panels:
#   (a) the T field over a +/-180 m window with catalogued boulders overlaid
#       (sized by diameter), the d_eff = 50 m interaction circle, a size legend,
#       and the largest nearby blocks labelled by diameter;
#   (b) a zoom on Otohime showing the larger blocks within d_eff that build the
#       local mass M_local (the single 164 m block holds essentially all of it);
#   (c) the diagnostic vs. the true field: T equals the boulder-pull-to-surface-
#       gravity ratio evaluated at the FIXED reference distance r = d_eff, whereas
#       the true pull of Otohime rises as 1/r^2 closer in (exceeding surface
#       gravity within r ~ 41 m) and falls off beyond. This makes explicit that
#       the bright disk in (a)/(b) is the d_eff footprint of the dominant block,
#       a property of the T diagnostic, not a physical trap boundary.
#
# This module IS the generator for manuscript Figure 2; code_generate_figures.py
# imports build_fig2() so the all-in-one run and the paper figure stay identical.
# Reads the real catalogue Ryugu_Boulder_Attributes.csv (same folder).
# =============================================================================
import os
import numpy as np
import pandas as pd
from scipy.spatial import cKDTree
import matplotlib
matplotlib.use('Agg')
import matplotlib.pyplot as plt
import matplotlib.patheffects as pe
from matplotlib.patches import Circle
from matplotlib.lines import Line2D

_HALO = [pe.withStroke(linewidth=1.8, foreground='white')]

# physical constants (consistent with verify_from_scratch.py and the manuscript)
G = 6.674e-11        # gravitational constant
M_AST = 4.5e11       # Ryugu mass (kg)
R = 450.0            # Ryugu mean radius (m)
G_SURF = G * M_AST / R**2     # surface gravity (m/s^2)
RHO_B = 1600.0       # boulder density (kg/m^3)
D_EFF = 50.0         # interaction radius (m)
A_ELL, C_ELL = 502.0, 437.744  # ellipsoid semi-axes (m)

_HERE = os.path.dirname(os.path.abspath(__file__))
_CSV = os.path.join(_HERE, 'Ryugu_Boulder_Attributes.csv')


def _load_catalogue():
    df = pd.read_csv(_CSV)
    df = df[df['Length_m'] > 0].copy()
    lon = np.radians(df['Longitude_X'].values % 360.0)
    lat = np.radians(df['Latitude_Y'].values)
    L = df['Length_m'].values

    def pos_of(la, lo):
        r = 1.0 / np.sqrt((np.cos(la) / A_ELL)**2 + (np.sin(la) / C_ELL)**2)
        return np.array([r * np.cos(la) * np.cos(lo),
                         r * np.cos(la) * np.sin(lo),
                         r * np.sin(la)])

    P = np.array([pos_of(a, o) for a, o in zip(lat, lon)])
    m = (4.0 / 3.0) * np.pi * RHO_B * (L / 2.0)**3
    return L, P, m


def build_fig2(outpath=None):
    """Build the 3-panel manuscript Figure 2 and save it. Returns the save path."""
    L, P, m = _load_catalogue()
    io = int(np.argmax(L))           # Otohime = largest block
    r0 = P[io]; m_O = m[io]
    # local tangent frame at Otohime
    n = np.array([2*r0[0]/A_ELL**2, 2*r0[1]/A_ELL**2, 2*r0[2]/C_ELL**2]); n /= np.linalg.norm(n)
    east = np.cross([0, 0, 1.0], n); east /= np.linalg.norm(east)
    north = np.cross(n, east)
    dP = P - r0; E = dP @ east; N = dP @ north
    tree = cKDTree(P)

    def Tfield(half, res):
        gx = np.linspace(-half, half, res); gy = np.linspace(-half, half, res)
        T = np.zeros((res, res))
        for i, yy in enumerate(gy):
            rowpts = np.array([r0 + xx*east + yy*north for xx in gx])
            idxs = tree.query_ball_point(rowpts, D_EFF)
            T[i, :] = [G * m[ix].sum() / (G_SURF * D_EFF**2) for ix in idxs]
        return gx, gy, T

    fig = plt.figure(figsize=(13, 9.2))
    gs = fig.add_gridspec(2, 2, height_ratios=[1.55, 1.0], hspace=0.30, wspace=0.24)
    axL = fig.add_subplot(gs[0, 0]); axR = fig.add_subplot(gs[0, 1]); axB = fig.add_subplot(gs[1, :])

    # ---- (a) wide field ----
    half = 180.0; gx, gy, T = Tfield(half, 170)
    pcm = axL.pcolormesh(gx, gy, T, shading='auto', cmap='inferno', vmin=0, vmax=0.7)
    inwin = (np.abs(E) < half) & (np.abs(N) < half)
    axL.scatter(E[inwin], N[inwin], s=(L[inwin]*1.1)**1.35, facecolors='none',
                edgecolors='#39d4e6', linewidths=0.8, alpha=0.9)
    axL.add_patch(Circle((0, 0), D_EFF, fill=False, ec='white', ls='--', lw=1.4))
    axL.annotate(r'$d_{\rm eff}=50$ m', xy=(D_EFF*0.7, D_EFF*0.72), color='black',
                 fontsize=9, fontweight='bold', path_effects=_HALO)
    Ein, Nin, Lin = E[inwin], N[inwin], L[inwin]; order = np.argsort(Lin)[::-1]; lab = 0
    for k in order:
        if np.hypot(Ein[k], Nin[k]) < 3:
            continue
        if lab >= 6:
            break
        axL.annotate(f'{Lin[k]:.0f} m', xy=(Ein[k], Nin[k]), xytext=(Ein[k]+6, Nin[k]+6),
                     color='black', fontsize=8, fontweight='bold', path_effects=_HALO); lab += 1
    axL.scatter([0], [0], marker='*', s=240, facecolor='white', edgecolor='black', linewidths=0.8, zorder=6)
    axL.annotate('Otohime\n164 m,  ' + r'$\mathcal{T}\approx0.5$', xy=(0, 0), xytext=(-35, -150),
                 color='black', fontsize=9.5, fontweight='bold', path_effects=_HALO,
                 arrowprops=dict(arrowstyle='->', color='white', lw=1.2))
    axL.set_xlim(-half, half); axL.set_ylim(-half, half); axL.set_aspect('equal')
    axL.set_xlabel('East offset from Otohime (m)'); axL.set_ylabel('North offset (m)')
    axL.set_title('(a)  Real-catalogue $\\mathcal{T}$ field (%d boulders in view)' % inwin.sum(), fontsize=10.5)
    leg = [Line2D([0], [0], marker='o', ls='none', mfc='none', mec='#39d4e6',
                  ms=np.sqrt((d*1.1)**1.35)/2.2, label=f'{d} m') for d in (20, 60, 120)]
    axL.legend(handles=leg, title='boulder size', loc='lower left', fontsize=8,
               title_fontsize=8, framealpha=0.85, labelspacing=1.3, borderpad=1.0)
    fig.colorbar(pcm, ax=axL, fraction=0.046, pad=0.02).set_label(r'$\mathcal{T}$ (real boulders)')

    # ---- (b) zoom: only the contributors that matter ----
    halfz = 72.0; gx2, gy2, T2 = Tfield(halfz, 150)
    pcm2 = axR.pcolormesh(gx2, gy2, T2, shading='auto', cmap='inferno', vmin=0, vmax=0.7)
    zin = (np.abs(E) < halfz) & (np.abs(N) < halfz); Ez, Nz, Lz = E[zin], N[zin], L[zin]
    dist = np.hypot(Ez, Nz)
    axR.add_patch(Circle((0, 0), D_EFF, fill=False, ec='white', ls='--', lw=1.6))
    axR.scatter(Ez, Nz, s=(Lz*1.1)**1.4, facecolors='none', edgecolors='#39d4e6', linewidths=1.0)
    for k in range(len(Lz)):
        if np.hypot(Ez[k], Nz[k]) < 3:
            continue
        if dist[k] <= D_EFF and Lz[k] >= 7:
            axR.annotate(f'{Lz[k]:.0f}', xy=(Ez[k], Nz[k]), xytext=(Ez[k]+2.5, Nz[k]+2.5),
                         color='black', fontsize=8, fontweight='bold', path_effects=_HALO)
    axR.scatter([0], [0], marker='*', s=300, facecolor='white', edgecolor='black', linewidths=0.8, zorder=6)
    axR.annotate('164', xy=(0, 0), xytext=(4, 4), color='black', fontsize=9, fontweight='bold', path_effects=_HALO)
    Mloc = m[tree.query_ball_point(r0, D_EFF)].sum()
    n_in = int((dist <= D_EFF).sum())
    axR.annotate('within $d_{\\rm eff}$: %d boulders\nOtohime (164 m) holds ~%.0f%% of $M_{\\rm local}$'
                 % (n_in, 100*m_O/Mloc), xy=(0, 0), xytext=(-68, -66), color='white', fontsize=8.5,
                 bbox=dict(boxstyle='round', fc='#222222', ec='0.6', alpha=0.85))
    axR.set_xlim(-halfz, halfz); axR.set_ylim(-halfz, halfz); axR.set_aspect('equal')
    axR.set_xlabel('East offset (m)'); axR.set_ylabel('North offset (m)')
    axR.set_title('(b)  Zoom: larger blocks inside $d_{\\rm eff}$ (labels = diameter, m)', fontsize=10.5)
    fig.colorbar(pcm2, ax=axR, fraction=0.046, pad=0.02).set_label(r'$\mathcal{T}$')

    # ---- (c) true 1/r^2 fall-off vs the single-distance diagnostic ----
    r = np.linspace(35, 135, 400)
    ratio_true = G * m_O / (r**2 * G_SURF)
    T_diag = G * m_O / (G_SURF * D_EFF**2)
    r1 = np.sqrt(G * m_O / G_SURF)        # where Otohime pull == surface gravity
    axB.plot(r, ratio_true, color='#c0392b', lw=2.6,
             label=r"true ratio  $G m_{\rm Oto}/(r^2 g_{\rm surf})$  (1/$r^2$)")
    axB.hlines(T_diag, 35, D_EFF, color='#16a085', lw=2.4,
               label=r"diagnostic $\mathcal{T}=G m/(g\,d_{\rm eff}^2)=%.2f$" % T_diag)
    axB.axvline(D_EFF, color='0.5', ls='--', lw=1.2); axB.axhline(1.0, color='black', ls=':', lw=1.2)
    axB.plot([D_EFF], [T_diag], 'o', color='#16a085', ms=8, zorder=5)
    axB.annotate(r'$\mathcal{T}$ is the ratio evaluated at $r=d_{\rm eff}$',
                 xy=(D_EFF, T_diag), xytext=(D_EFF+6, T_diag+0.15), fontsize=9)
    if 35 <= r1 <= 135:
        axB.plot([r1], [1.0], 'o', color='black', ms=7, zorder=5)
        axB.annotate('Otohime pull = surface gravity\nat $r\\approx%.0f$ m' % r1,
                     xy=(r1, 1.0), xytext=(r1+5, 1.25), fontsize=9,
                     arrowprops=dict(arrowstyle='->', color='0.4', lw=0.8))
    axB.axvspan(35, 42, color='0.85', alpha=0.5)
    axB.annotate('near-field:\npoint-mass approx.', xy=(38, 0.2), fontsize=7.5, color='0.35')
    axB.set_xlim(35, 135); axB.set_ylim(0, 1.6)
    axB.set_xlabel('distance from Otohime centre  $r$ (m)'); axB.set_ylabel(r'boulder pull / surface gravity')
    axB.set_title(r'(c)  $\mathcal{T}$ is a single-distance proxy: the true Otohime pull rises as '
                  r'$1/r^2$ inside $d_{\rm eff}$ and falls off beyond', fontsize=10)
    axB.legend(loc='upper right', fontsize=8.5, framealpha=0.92); axB.grid(alpha=0.25)

    if outpath is None:
        figdir = os.path.normpath(os.path.join(_HERE, '..', '..', '01_Figures'))
        outpath = os.path.join(figdir, 'Fig2.png') if os.path.isdir(figdir) \
            else os.path.join(_HERE, 'Fig2.png')
    fig.savefig(outpath, dpi=190, bbox_inches='tight', pad_inches=0.1)
    plt.close(fig)
    return outpath


if __name__ == '__main__':
    p = build_fig2()
    print('saved manuscript Figure 2 ->', p)
