"""zepto: ZeptoPhysics formulas as plain Python (numpy only).

The same equations that power the calculators on zeptophysics.com, so that an AI
assistant (or you) can compute instead of guess. Units are stated in every docstring.
Licence: MIT. Source: https://zeptophysics.com/lab/
"""
import math

try:
    import numpy as np  # optional, only used where noted
except Exception:  # pragma: no cover
    np = None

# ---------------------------------------------------------------- constants (SI unless noted)
c = 2.99792458e8            # m/s
G = 6.67430e-11             # m^3 kg^-1 s^-2
h = 6.62607015e-34          # J s
hbar = 1.054571817e-34      # J s
k_B = 1.380649e-23          # J/K
e = 1.602176634e-19         # C
m_e = 9.1093837015e-31      # kg
m_p = 1.67262192369e-27     # kg
eps0 = 8.8541878128e-12     # F/m
mu0 = 1.25663706212e-6      # N/A^2
M_sun = 1.98847e30          # kg
pc = 3.0856775814913673e16  # m
Mpc = 1e6 * pc
year = 3.15576e7            # s (Julian)
GM_earth = 3.986004418e14   # m^3/s^2
R_earth = 6.371e6           # m
HBARC_GeV_m = 1.973269804e-16  # hbar*c in GeV*m
MeV_fm = 197.3269804           # hbar*c in MeV*fm

# ---------------------------------------------------------------- units (multiply to convert INTO SI)
m = meter = metre = 1.0
km = 1e3; cm = 1e-2; mm = 1e-3; um = micron = 1e-6; nm = 1e-9; pm = 1e-12; fm = 1e-15; am = 1e-18; zm = 1e-21
AU = 1.495978707e11; ly = light_year = 9.4607304725808e15; kpc = 1e3 * pc; Gpc = 1e9 * pc
s = second = 1.0; ms = 1e-3; us = 1e-6; ns = 1e-9; ps = 1e-12; fs = 1e-15; attosecond = 1e-18
minute = 60.0; hour = 3600.0; day = 86400.0; Gyr = 1e9 * year
Hz = 1.0; kHz = 1e3; MHz = 1e6; GHz = 1e9; THz = 1e12
kg = 1.0; g_mass = 1e-3; Msun = M_solar = solar_mass = M_sun
J = 1.0; eV = 1.602176634e-19; keV = 1e3 * eV; MeV = 1e6 * eV; GeV = 1e9 * eV; TeV = 1e12 * eV; erg = 1e-7
W = 1.0; kW = 1e3; MW = 1e6; GW_power = 1e9; TW = 1e12; PW = 1e15
T = tesla = 1.0; gauss = 1e-4
LIGO_arm_m = 4000.0
g = g0 = g_earth = 9.80665   # standard gravity, m/s^2


def help():
    """List the available functions."""
    names = [n for n, f in globals().items() if callable(f) and not n.startswith('_') and f.__module__ == __name__]
    for n in sorted(names):
        doc = (globals()[n].__doc__ or '').strip().splitlines()[0]
        print(f"{n:28s} {doc}")


# ---------------------------------------------------------------- scale translator
def energy_to_length(E_GeV):
    """Length (m) probed by an energy/mass scale E (GeV): lambda = hbar c / E."""
    return HBARC_GeV_m / E_GeV


def length_to_energy(L_m):
    """Energy scale (GeV) corresponding to a length L (m)."""
    return HBARC_GeV_m / L_m


# ---------------------------------------------------------------- pulsars / cosmology / black holes
def pulsar(P_s, Pdot, I=1e38):
    """Dipole spin-down from period P (s) and Pdot (s/s). I in kg m^2 (1e38 = 1e45 g cm^2).
    Returns dict: tau_c_yr, B_surface_G, Edot_W, R_lc_m."""
    tau = P_s / (2 * Pdot)
    B = 3.2e19 * math.sqrt(P_s * Pdot)  # gauss
    Edot = 4 * math.pi ** 2 * I * Pdot / P_s ** 3
    return dict(tau_c_yr=tau / year, B_surface_G=B, Edot_W=Edot, R_lc_m=c * P_s / (2 * math.pi))


def lcdm_distances(z, H0=67.4, Om=0.315, n=4000):
    """Flat LambdaCDM distances. Returns dict: D_C_Mpc, D_L_Mpc, D_A_Mpc, lookback_Gyr."""
    OL = 1 - Om
    E = lambda zz: math.sqrt(Om * (1 + zz) ** 3 + OL)
    if z <= 0:
        return dict(D_C_Mpc=0.0, D_L_Mpc=0.0, D_A_Mpc=0.0, lookback_Gyr=0.0)
    n += n % 2
    hstep = z / n
    ic = it = 0.0
    for i in range(n + 1):
        zz = i * hstep
        w = 1 if i in (0, n) else (4 if i % 2 else 2)
        ic += w / E(zz)
        it += w / ((1 + zz) * E(zz))
    DH = 299792.458 / H0
    tH = 977.7922 / H0
    DC = DH * ic * hstep / 3
    return dict(D_C_Mpc=DC, D_L_Mpc=DC * (1 + z), D_A_Mpc=DC / (1 + z), lookback_Gyr=tH * it * hstep / 3)


def schwarzschild(M_solar):
    """Non-spinning black hole of mass M (solar masses). Returns dict: R_s_m, R_photon_m, R_isco_m, T_hawking_K."""
    Rs = 2 * G * M_solar * M_sun / c ** 2
    return dict(R_s_m=Rs, R_photon_m=1.5 * Rs, R_isco_m=3 * Rs, T_hawking_K=6.17e-8 / M_solar)


def frb_delay(DM, f_lo_MHz, f_hi_MHz):
    """Cold-plasma dispersion delay (s) between two frequencies (MHz) for DM (pc cm^-3)."""
    return 4.148808e3 * DM * (f_lo_MHz ** -2 - f_hi_MHz ** -2)


# ---------------------------------------------------------------- gravitational waves
def chirp_mass(m1, m2):
    """Chirp mass (solar masses) of a binary with component masses m1, m2 (solar masses)."""
    return (m1 * m2) ** 0.6 / (m1 + m2) ** 0.2


def gw_strain(m1, m2, D_Mpc, f_Hz):
    """Leading-order strain amplitude h0 (optimal orientation) at GW frequency f (Hz), distance D (Mpc)."""
    GMc2 = G * chirp_mass(m1, m2) * M_sun / c ** 2
    return 4 * GMc2 ** (5 / 3) * (math.pi * f_Hz / c) ** (2 / 3) / (D_Mpc * Mpc)


def time_to_merger(m1, m2, f_Hz):
    """Newtonian time to coalescence (s) from GW frequency f (Hz)."""
    return 2.18 * (1.21 / chirp_mass(m1, m2)) ** (5 / 3) * (100 / f_Hz) ** (8 / 3)


def f_isco(m1, m2):
    """GW frequency (Hz) at the innermost stable circular orbit."""
    return 4400.0 / (m1 + m2)


def gw_luminosity(m1, m2, f_Hz):
    """Leading-order GW luminosity (W) at GW frequency f (Hz)."""
    x = math.pi * f_Hz * G * chirp_mass(m1, m2) * M_sun / c ** 3
    return 32 / 5 * c ** 5 / G * x ** (10 / 3)


# ---------------------------------------------------------------- plasma
def synchrotron(gamma, B_G):
    """Electron synchrotron (90 deg pitch), CGS field B in gauss. Returns nu_c_Hz, t_cool_s, r_L_cm."""
    return dict(nu_c_Hz=4.2e6 * gamma ** 2 * B_G, t_cool_s=5.1e8 / (gamma * B_G ** 2), r_L_cm=1.7e3 * gamma / B_G)


def plasma_scales(n_e_cm3, T_K):
    """Debye length (cm) and electron plasma frequency (Hz)."""
    return dict(debye_cm=6.9 * math.sqrt(T_K / n_e_cm3), f_pe_Hz=8.98e3 * math.sqrt(n_e_cm3))


# ---------------------------------------------------------------- colliders
def collider(E_beam_TeV, pT_GeV=None, B_T=None, m_GeV=0.93827):
    """Symmetric collider: sqrt_s_TeV, gamma; optional track radius (m) for pT (GeV) in field B (T)."""
    out = dict(sqrt_s_TeV=2 * E_beam_TeV, gamma=E_beam_TeV * 1e3 / m_GeV)
    if pT_GeV and B_T:
        out['radius_m'] = pT_GeV / (0.3 * B_T)
    return out


def synchrotron_loss_per_turn(E_GeV, rho_m, m_GeV=0.000511):
    """Energy radiated per turn (eV) by a particle of energy E (GeV) on bending radius rho (m)."""
    return 88.5e3 * E_GeV ** 4 / rho_m * (0.000511 / m_GeV) ** 4


def lifetime(width_GeV):
    """Mean lifetime (s) from a decay width (GeV): tau = hbar / Gamma."""
    return 6.582119569e-25 / width_GeV


# ---------------------------------------------------------------- lasers
def laser(E_J, tau_fs, spot_um, lam_um=0.8):
    """Focused Gaussian pulse. Returns P_W, I_Wcm2, a0, E_field_Vm, ratio_to_Schwinger, n_crit_cm3."""
    P = 0.94 * E_J / (tau_fs * 1e-15)
    d = spot_um * 1e-4
    I = 0.88 * P / d ** 2
    a0 = 0.855 * lam_um * math.sqrt(I / 1e18)
    EL = 2745 * math.sqrt(I)
    return dict(P_W=P, I_Wcm2=I, a0=a0, E_field_Vm=EL, ratio_to_Schwinger=EL / 1.323e18, n_crit_cm3=1.115e21 / lam_um ** 2)


def wakefield(P_W, n_e_cm3, lam_um=0.8):
    """Bubble-regime LWFA scalings (Lu et al. 2007). Returns lambda_p_m, E0_GVm, P_crit_W, energy_gain_GeV."""
    nc = 1.115e21 / lam_um ** 2
    return dict(lambda_p_m=33.4e-6 * math.sqrt(1e18 / n_e_cm3), E0_GVm=96 * math.sqrt(n_e_cm3 / 1e18),
                P_crit_W=17e9 * nc / n_e_cm3,
                energy_gain_GeV=1.7 * (P_W / 1e14) ** (1 / 3) * (1e18 / n_e_cm3) ** (2 / 3) * (0.8 / lam_um) ** (4 / 3))


def hhg_cutoff(I_Wcm2, Ip_eV=15.76, lam_um=0.8):
    """High-harmonic cutoff photon energy (eV): Ip + 3.17 Up."""
    return Ip_eV + 3.17 * 9.33e-14 * I_Wcm2 * lam_um ** 2


# ---------------------------------------------------------------- fusion
def dt_reactivity(T_keV):
    """D-T <sigma v> (m^3/s), Bosch-Hale (1992) fit, valid ~0.2-100 keV."""
    BG, mrc2 = 34.3827, 1124656
    C1, C2, C3, C4, C5, C6, C7 = 1.17302e-9, 1.51361e-2, 7.51886e-2, 4.60643e-3, 1.35e-2, -1.0675e-4, 1.366e-5
    th = T_keV / (1 - (T_keV * (C2 + T_keV * (C4 + T_keV * C6))) / (1 + T_keV * (C3 + T_keV * (C5 + T_keV * C7))))
    xi = (BG ** 2 / (4 * th)) ** (1 / 3)
    return C1 * th * math.sqrt(xi / (mrc2 * T_keV ** 3)) * math.exp(-3 * xi) * 1e-6


def tokamak(n_m3, T_keV, tau_E_s, V_m3, f_DT=0.8, B_T=None):
    """0-D D-T power balance with flat profiles. Returns P_fus_W, P_fus_density_W_m3, Q, triple_product, beta."""
    keV = 1.602176634e-16
    pf = 0.25 * (f_DT * n_m3) ** 2 * dt_reactivity(T_keV) * 17.59e3 * keV
    pa = pf * 3.52 / 17.59
    pext = 3 * n_m3 * T_keV * keV / tau_E_s - pa
    out = dict(P_fus_W=pf * V_m3, P_fus_density_W_m3=pf, Q=(pf / pext) if pext > 0 else float('inf'), triple_product=n_m3 * T_keV * tau_E_s)
    if B_T:
        out['beta'] = 2 * mu0 * (2 * n_m3 * T_keV * keV) / B_T ** 2
    return out


# ---------------------------------------------------------------- clocks & relativity
def clock_redshift(dh_m):
    """Fractional frequency shift for a height difference dh (m) near Earth's surface (exact potential)."""
    return GM_earth * (1 / R_earth - 1 / (R_earth + dh_m)) / c ** 2


def orbit_clock(r_km):
    """Satellite clock rate vs. ground (us/day): returns gravity, velocity, net."""
    r = r_km * 1e3
    g = GM_earth * (1 / R_earth - 1 / r) / c ** 2 * 86400e6
    v = -(GM_earth / r) / (2 * c ** 2) * 86400e6
    return dict(gravity_us_day=g, velocity_us_day=v, net_us_day=g + v)


def ultralight_dm(m_eV, v=2.2e5):
    """Ultralight dark matter: Compton frequency (Hz), period (s), coherence time (s), de Broglie wavelength (m)."""
    f = m_eV * e / h
    return dict(f_Hz=f, period_s=1 / f, coherence_s=1e6 / f, lambda_dB_m=h / (m_eV * 1.78266192e-36 * v))


def atom_interferometer(T_s, lam_m=780e-9, N=1e6, two_photon=True):
    """Gravity phase (rad) and shot-noise dg/g per shot for interrogation time T (s)."""
    k = (4 if two_photon else 2) * math.pi / lam_m
    phi = k * 9.80665 * T_s ** 2
    return dict(phase_rad=phi, dg_over_g=1 / (math.sqrt(N) * phi), launch_height_m=9.80665 * T_s ** 2 / 2)


# ---------------------------------------------------------------- XFEL / magnets
def undulator_wavelength(E_GeV, lambda_u_m, K):
    """FEL resonant wavelength (m) for electron energy E (GeV), undulator period (m), parameter K."""
    g = E_GeV / 0.000510999
    return lambda_u_m / (2 * g ** 2) * (1 + K ** 2 / 2)


def magnet(B_T, r_m=None, L_m=None):
    """Magnetic pressure (Pa) = energy density (J/m^3); stored energy (J) for a solenoid bore r, length L."""
    p = B_T ** 2 / (2 * mu0)
    out = dict(pressure_Pa=p)
    if r_m and L_m:
        out['stored_J'] = p * math.pi * r_m ** 2 * L_m
    return out


def __getattr__(name):  # PEP 562: friendly errors for guessed names
    import difflib
    names = [n for n in globals() if not n.startswith('_')]
    close = difflib.get_close_matches(name, names, n=4, cutoff=0.5)
    hint = f" Did you mean: {', '.join('zepto.' + c for c in close)}?" if close else ''
    raise AttributeError(f"zepto has no '{name}'.{hint} Call zepto.help() to list functions; functions take the units named in their arguments (e.g. solar masses, Mpc, Hz), so pass numbers directly.")