"""A deliberately simple survey simulator.

Every object is followed with daily steps from 1998.0 to 2026.78.  For each
day we compute g(t) = V - H = 5 log10(r*Delta) + phase term + trailing loss, and
a solar-elongation weight E.  The survey is described, per era k, by a limiting
magnitude Vlim_k and a nightly success probability p_k:

    P(discovered by end of era K) = 1 - exp( sum_{k<=K} sum_days E * ln(1 - p_k * rho_b * S(H + g - Vlim_k)) )

with S a logistic roll-off of width 0.5 mag, set to zero for V fainter than a
hard per-era limit (VHARD_BASE), and rho_b = 1 within 15 deg of the
ecliptic and rho (fixed, not fitted) at higher ecliptic latitude.  E is zero south of
declination -30 deg and depends on solar elongation.

Elongation weight (changed in the revision).  The first version ASSUMED a ramp starting at
exactly 60 deg (elong_weight_ramp below).  Sun-Earth L4/L5 sit at 60 deg, so their detection
probability was set by that assumption.  The weight is now MEASURED: s02 compares the solar
elongation of real near-Earth asteroids on the night of first observation (s01b) with the
simulator's predicted opportunities under flat coverage, and writes the ratio per elongation
bin to data/elong_weight_measured.csv; load_elong_table() makes elong_weight() interpolate
that table (columns w_c, w_lo, w_hi).  The assumed ramp is kept for stage 1 of the
calibration and for sensitivity runs.

Because g does not depend on H, we store for each orbit an E-weighted histogram of g per
era and reuse it for any H (i.e. any size and albedo).  Only the five nightly probabilities
p_k are fitted, to the real discovery history of near-Earth asteroids (s02); the simulator
is then applied to hypothetical parked objects (s05) and interstellar trajectories (s09).
"""
import os

import numpy as np

from common import DATA, K_GAUSS, helio_xyz, observe, phase_HG, sky_rate_deg_day, trailing_loss, trailing_loss_theory

T0_YEAR = 1998.0
T1_YEAR = 2026.78
DAYS_PER_YEAR = 365.25
NDAYS = int((T1_YEAR - T0_YEAR) * DAYS_PER_YEAR)
T_DAYS = np.arange(NDAYS, dtype=float)
T_YEARS = T0_YEAR + T_DAYS / DAYS_PER_YEAR
ERA_EDGES = np.array([1998.0, 2005.0, 2011.0, 2016.0, 2021.0, 2027.0])
N_ERA = len(ERA_EDGES) - 1
ERA_IDX = np.clip(np.searchsorted(ERA_EDGES, T_YEARS, side="right") - 1, 0, N_ERA - 1)

G_MIN, G_MAX, G_STEP = -16.0, 12.0, 0.2
G_EDGES = np.arange(G_MIN, G_MAX + 1e-9, G_STEP)
G_CEN = 0.5 * (G_EDGES[1:] + G_EDGES[:-1])
N_G = len(G_CEN)
LOGISTIC_W = 0.5
# Hard faint limit per era: nothing fainter than this is ever detected, however many chances.
# 1998-2004 LINEAR (~20.5 on the best nights), 2005-2010 Catalina 1.5 m / Spacewatch (~22),
# 2011 onward Pan-STARRS (w ~ 22.5-23 on the best nights).
VHARD_BASE = np.array([20.5, 22.0, 23.0, 23.0, 23.0])
BETA_SPLIT = 15.0           # deg; geocentric ecliptic latitude separating the two sky bands

MU_EARTH = 3.04e-6          # (Earth+Moon)/Sun mass ratio
LAM_E0 = 0.0                # Earth's mean longitude at T0 (arbitrary: all orbit phases are random)


def earth_xyz(t_days=T_DAYS):
    lam = LAM_E0 + K_GAUSS * t_days
    return np.cos(lam), np.sin(lam), np.zeros_like(lam)


OBLIQ = np.radians(23.4393)
DEC_MIN = -30.0             # deg; southern declination limit of the (mostly northern) surveys
E0_DEFAULT, E1_DEFAULT = 60.0, 90.0   # deg; solar-elongation ramp of survey coverage
ELONG_TILT = 0.32   # survey coverage at elongation e0 relative to opposition (see elong_weight)


def elong_weight_ramp(elong_rad, e0=E0_DEFAULT, e1=E1_DEFAULT, tilt=ELONG_TILT):
    """ASSUMED survey coverage weight versus solar elongation (first version of the study):
    0 below e0, a linear ramp to e1, multiplied by a linear tilt that rises from `tilt` at e0
    to 1 at opposition.  Used for stage 1 of the calibration and for sensitivity runs."""
    el = np.degrees(elong_rad)
    ramp = np.clip((el - e0) / (e1 - e0), 0.0, 1.0)
    return ramp * (tilt + (1.0 - tilt) * np.clip((el - e0) / (180.0 - e0), 0.0, 1.0))


ELONG_EDGES = np.array([0.0, 45.0, 60.0, 70.0, 80.0, 90.0, 105.0, 120.0, 135.0, 150.0, 165.0, 180.0])
ELONG_CEN = 0.5 * (ELONG_EDGES[1:] + ELONG_EDGES[:-1])
N_ELB = len(ELONG_CEN)
ELONG_TABLE_FILE = os.path.join(DATA, "elong_weight_measured.csv")
_ELONG_TABLES = {}           # {survey variant: dict(c=..., lo=..., hi=...)} of weights at ELONG_CEN once loaded


def set_elong_table(variant, tab):
    """Activate measured weights for one survey variant from a table with columns w_c, w_lo, w_hi
    (one row per elongation bin, in order)."""
    _ELONG_TABLES[variant] = {k: np.asarray(tab["w_" + k], float).copy() for k in ("c", "lo", "hi")}


def load_elong_table(path=ELONG_TABLE_FILE):
    """Activate the measured elongation weights written by s02 (one block of rows per survey
    variant, because the weight is measured with each variant's own depth and nightly
    probabilities).  Returns True if loaded."""
    _ELONG_TABLES.clear()
    if not os.path.exists(path):
        return False
    import pandas as pd
    t = pd.read_csv(path)
    if "variant" not in t.columns:
        t = t.assign(variant="base")
    for v, g in t.groupby("variant", sort=False):
        set_elong_table(v, g.sort_values("elong_centre"))
    return True


def unload_elong_table():
    _ELONG_TABLES.clear()


def has_elong_table(variant):
    return variant in _ELONG_TABLES


def elong_weight(elong_rad, level="c", variant="base"):
    """The active survey coverage weight versus solar elongation: the measured table of the
    given survey variant if tables have been loaded (the base table for a variant without its
    own), otherwise the assumed ramp.  Linear interpolation between the bin centres; constant
    beyond the first and last centre."""
    if _ELONG_TABLES:
        tab = _ELONG_TABLES[variant] if variant in _ELONG_TABLES else _ELONG_TABLES["base"]
        return np.interp(np.degrees(elong_rad), ELONG_CEN, tab[level])
    return elong_weight_ramp(elong_rad)


def make_wfun(e0=None, e1=None, level="c", variant="base"):
    """Weight function for hist_from_positions: the assumed ramp e0-e1 if e0 is given, else the
    active function at the given level (c / lo / hi of the measured table) and survey variant."""
    if e0 is not None:
        return lambda el: elong_weight_ramp(el, e0, e1 if e1 is not None else e0 + 30.0)
    return lambda el: elong_weight(el, level, variant)


def _alloc(wfun, n):
    if isinstance(wfun, dict):
        return {k: np.zeros((n, N_ERA, 2, N_G)) for k in wfun}
    return np.zeros((n, N_ERA, 2, N_G))


def _store(out, sl, h):
    if isinstance(out, dict):
        for k in out:
            out[k][sl] = h[k]
    else:
        out[sl] = h


def hist_from_positions(x, y, z, extra_mag=0.0, wfun=None, rate_override=None,
                        dec_min=DEC_MIN, tmask=None, trail="neomod2"):
    """x,y,z: heliocentric positions [n_obj, NDAYS].  Returns hist [n_obj, N_ERA, 2, N_G]
    (third axis: 0 = within BETA_SPLIT deg of the ecliptic, 1 = higher latitude).
    wfun: elongation weight function of the elongation in rad (default: the active one), or a
    dict {name: function}, in which case a dict {name: hist} is returned (one geometry pass).
    trail: 'neomod2' (empirical Catalina law, fitted below 10 deg/day and extrapolated beyond)
    or 'theory' (Veres & Chesley 2017 detection loss for 30 s exposures and 2.5 arcsec seeing)."""
    ex, ey, ez = earth_xyz()
    r, d, alpha, elong, unit = observe((x, y, z), (ex[None, :], ey[None, :], ez[None, :]))
    rate = sky_rate_deg_day(unit) if rate_override is None else rate_override
    tl = trailing_loss(rate) if trail == "neomod2" else trailing_loss_theory(rate)
    g = 5.0 * np.log10(r * d) + phase_HG(alpha) + tl + extra_mag
    sin_dec = unit[2] * np.cos(OBLIQ) + unit[1] * np.sin(OBLIQ)
    w0 = (sin_dec > np.sin(np.radians(dec_min))).astype(float)
    if tmask is not None:
        w0 = w0 * tmask
    n_obj = x.shape[0]
    ib = np.clip(((g - G_MIN) / G_STEP).astype(np.int64), 0, N_G - 1)
    band = (np.abs(unit[2]) > np.sin(np.radians(BETA_SPLIT))).astype(np.int64)   # 0: near ecliptic, 1: high latitude
    flat = (((np.arange(n_obj)[:, None] * N_ERA + ERA_IDX[None, :]) * 2 + band) * N_G + ib).ravel()

    def one(f):
        h = np.bincount(flat, weights=(f(elong) * w0).ravel(), minlength=n_obj * N_ERA * 2 * N_G)
        return h.reshape(n_obj, N_ERA, 2, N_G)
    wf = elong_weight if wfun is None else wfun
    if isinstance(wf, dict):
        return {k: one(f) for k, f in wf.items()}
    return one(wf)


def hist_kepler(a, e, inc, Om, w, M0, chunk=200, t_start_days=None, **kw):
    """Histograms for fixed heliocentric Kepler orbits (angles in rad, M0 at T0).
    t_start_days: optional per-object day (from T0) before which the object does not count."""
    n = len(a)
    out = _alloc(kw.get("wfun"), n)
    for s in range(0, n, chunk):
        sl = slice(s, min(s + chunk, n))
        nmot = K_GAUSS / a[sl] ** 1.5
        M = M0[sl, None] + nmot[:, None] * T_DAYS[None, :]
        x, y, z = helio_xyz(a[sl, None], e[sl, None], inc[sl, None], Om[sl, None], w[sl, None], M)
        tm = None if t_start_days is None else (T_DAYS[None, :] >= np.asarray(t_start_days)[sl, None])
        _store(out, sl, hist_from_positions(x, y, z, tmask=tm, **kw))
    return out


def survival_log(hist, H, vlim, p, rho, vhard=None, width=None):
    """ln(1-P_detect) per era.  hist [n_obj, N_ERA, 2, N_G]; H scalar or [n_obj];
    vlim, p, rho: [N_ERA]; rho is the nightly-probability ratio high-latitude / near-ecliptic;
    width: width (mag) of the logistic roll-off of the detection efficiency (default LOGISTIC_W).
    Returns [n_obj, N_ERA]."""
    H = np.asarray(H, float)
    Hc = H[:, None, None, None] if H.ndim == 1 else H
    v = Hc + G_CEN[None, None, None, :] - np.asarray(vlim)[None, :, None, None]
    vh = VHARD_BASE if vhard is None else np.asarray(vhard)
    wd = LOGISTIC_W if width is None else float(width)
    vapp = Hc + G_CEN[None, None, None, :]
    Sg = np.where(vapp <= vh[None, :, None, None], 1.0 / (1.0 + np.exp(np.clip(v / wd, -60, 60))), 0.0)
    pb = np.stack([np.asarray(p), np.asarray(p) * np.asarray(rho)], axis=1)       # [era, band]
    return np.sum(hist * np.log1p(-pb[None, :, :, None] * Sg), axis=(2, 3))


def p_detect(hist, H, vlim, p, rho, upto_era=None, vhard=None, width=None):
    L = survival_log(hist, H, vlim, p, rho, vhard, width)
    if upto_era is not None:
        L = L[:, :upto_era + 1]
    return 1.0 - np.exp(L.sum(axis=1))


def fit_kw(f):
    """Keyword arguments (vhard, width) of a fitted survey variant (a dict from survey_fit.json)."""
    return dict(vhard=f["vhard"], width=f.get("width", LOGISTIC_W))


# ----------------------------------------------------------------------------
# co-orbital dynamics (averaged, low-eccentricity): phi = lambda - lambda_E,
# x = a - 1.  d2phi/dt2 = -(3/2) mu n^2 f'(phi),  f = 1/sin(phi/2) - 2 cos(phi),
# dphi/dt = -(3/2) n x.  f=1 at L4/L5, f=3 at L3 (tadpole/horseshoe separatrix).
# ----------------------------------------------------------------------------

def f_coorb(phi):
    return 1.0 / np.abs(np.sin(phi / 2.0)) - 2.0 * np.cos(phi)


def fprime_coorb(phi):
    s = np.sin(phi / 2.0)
    return -np.cos(phi / 2.0) * np.sign(s) / (2.0 * s * s) + 2.0 * np.sin(phi)


def coorbital_tracks(phi_turn, rng, burn_years=3000.0, dt=5.0):
    """Integrate the averaged co-orbital equation starting at the turning point
    phi_turn (rad, x=0) for a random time in [0, burn_years], then return daily
    phi(t), x(t) over the survey window.  Leapfrog; shape [n, NDAYS]."""
    n = len(phi_turn)
    phi = phi_turn.astype(float).copy()
    v = np.zeros(n)                     # dphi/dt  [rad/day]
    c = -1.5 * MU_EARTH * K_GAUSS ** 2
    nburn = (rng.uniform(0, burn_years, n) * DAYS_PER_YEAR / dt).astype(int)
    for k in range(int(nburn.max()) + 1):
        act = k < nburn
        if not act.any():
            break
        a_ = c * fprime_coorb(phi)
        vh = v + 0.5 * dt * a_
        ph = phi + dt * vh
        a2 = c * fprime_coorb(ph)
        vn = vh + 0.5 * dt * a2
        phi = np.where(act, ph, phi)
        v = np.where(act, vn, v)
    PHI = np.empty((n, NDAYS))
    X = np.empty((n, NDAYS))
    d1 = 1.0
    for k in range(NDAYS):
        PHI[:, k] = phi
        X[:, k] = -v / (1.5 * K_GAUSS)
        a_ = c * fprime_coorb(phi)
        vh = v + 0.5 * d1 * a_
        phi = phi + d1 * vh
        v = vh + 0.5 * d1 * c * fprime_coorb(phi)
    return PHI, X


def hist_coorbital(PHI, X, e, inc, Om, varpi, chunk=200, **kw):
    """Histograms for co-orbital objects with slowly varying a=1+x and mean longitude
    lambda_E + phi."""
    n = PHI.shape[0]
    out = _alloc(kw.get("wfun"), n)
    lamE = LAM_E0 + K_GAUSS * T_DAYS
    for s in range(0, n, chunk):
        sl = slice(s, min(s + chunk, n))
        a = 1.0 + X[sl]
        lam = lamE[None, :] + PHI[sl]
        M = lam - varpi[sl, None]
        w = varpi[sl, None] - Om[sl, None]
        x, y, z = helio_xyz(a, e[sl, None], inc[sl, None], Om[sl, None], w, M)
        _store(out, sl, hist_from_positions(x, y, z, **kw))
    return out


def hist_geocentric(R_km, inc, phase0, rng=None, chunk=200, **kw):
    """Objects on circular geocentric orbits of radius R_km (prograde, inclination inc to
    the ecliptic).  The sky rate is taken as the geocentric orbital rate."""
    from common import AU_KM
    n = len(R_km)
    out = _alloc(kw.get("wfun"), n)
    ex, ey, ez = earth_xyz()
    for s in range(0, n, chunk):
        sl = slice(s, min(s + chunk, n))
        P = 27.32166 * (R_km[sl] / 384400.0) ** 1.5      # days
        th = phase0[sl, None] + 2 * np.pi * T_DAYS[None, :] / P[:, None]
        R = (R_km[sl] / AU_KM)[:, None]
        x = ex[None, :] + R * np.cos(th)
        y = ey[None, :] + R * np.sin(th) * np.cos(inc[sl, None])
        z = R * np.sin(th) * np.sin(inc[sl, None])
        rate = np.broadcast_to((360.0 / P)[:, None], x.shape)
        _store(out, sl, hist_from_positions(x, y, z, rate_override=rate, **kw))
    return out