"""Calibrate the survey simulator on the real discovery history of near-Earth asteroids.

Truth population: orbits from the Granvik et al. (2018) debiased model
realisation (1,500 orbits drawn per unit-magnitude bin, 17<H<25; for H>25 the
24<H<25 orbits are reused), with the number per H bin taken from NEOMOD2
(primary) or Harris & Chodas 2021 (alternative).
Data: number of catalogued near-Earth asteroids per 0.5-mag bin discovered
before 2005, 2011, 2016, 2021 and 2026.8 (from s01).

Stages (2-4 are new in the revision):
  1. Fit with the ASSUMED elongation ramp (coverage from 60 deg, simlib.elong_weight_ramp).
  2. MEASURE the coverage weight versus solar elongation: the number of real first
     observations per elongation bin (s01b; 17.25 <= H < 27.75; objects closer than 0.005 au
     left out) is divided by the number the simulator predicts if coverage did not depend on
     elongation (sum over model objects and days of survival x nightly probability x
     detection efficiency), and normalised to 1 at elongations above 150 deg.  The survival
     term depends on the weight, so the measurement is iterated to a fixed point, the
     simulator is refitted with the measured weight and the measurement repeated.
  3. The same measurement for each depth variant with its own limiting magnitudes and
     nightly probabilities: the discoveries at low elongation constrain the product
     weight x nightly probability x efficiency, not the weight alone.
  4. Final fits of all variants with their measured weights (hc21 and ecliptic use the base
     table).  Weights with 95% Poisson bounds: data/elong_weight_measured.csv.

Fit: the nightly success probability in each of 5 eras (5 parameters).  The
limiting magnitudes are NOT fitted: a free fit is degenerate (a deep limit with a
tiny nightly probability fits as well as a realistic limit), so they are fixed at
VLIM_BASE + dv.  Revision: the central variant is now dv = 0, i.e. the published
effective depths (21.0-21.8 after 2005); the first version used dv = +1.0
(22.5-22.8), which fits the completeness history better but predicts discoveries
far fainter than the real ones (discovery-magnitude test in s05).  Variants:
  base      dv =  0.0   (19.5, 21.0, 21.5, 21.7, 21.8)
  shallow   dv = -0.5   (limits 21.0-21.3 after 2011)
  deep      dv = +1.0   (22.5-22.8; the baseline of the first version)
  vdeep     dv = +1.5, hard limit +0.5 (23.0-23.3; deeper than any published single-exposure depth)
  hard22p5  dv =  0.0, hard limit 22.5 instead of 23.0
  hc21      Harris & Chodas 2021 populations
  ecliptic  nightly probability at |ecliptic latitude| > 15 deg multiplied by 0.3
The completeness history and the discovery-magnitude test pull in opposite directions,
which signals a mis-specified model (one efficiency curve per era cannot represent many
surveys); the spread shallow..vdeep is carried as the depth systematic.
Nothing fainter than a hard per-era limit (simlib.VHARD_BASE, 23.0 after 2011) is ever
detected.
Validation (not used in the fit): completeness split by orbit class
(Amor / Apollo / Aten), for i<5 deg, i>25 deg and for Earth-like orbits.

Outputs: data/survey_fit.json, data/survey_fit_curves.csv, data/survey_validation.csv,
         data/elong_weight_measured.csv, data/survey_fit_stage1.json,
         data/neo_orbit_hists.npz (cached geometry histograms, measured base weights),
         RAW/neo_orbit_hists_stage1.npz (cached, assumed ramp),
         RAW/neo_orbit_hists_variants.npz (cached, measured weights of the depth variants)
"""
import gzip
import json
import os

import numpy as np
import pandas as pd
from scipy.optimize import least_squares

import simlib as S
from common import DATA, RAW, SEED, K_GAUSS, helio_xyz, observe, phase_HG, sky_rate_deg_day, trailing_loss
from s01_neo_completeness import EPOCH_ENDS, load_neo

N_PER_BIN = 1500
# Effective limiting V magnitude of the deepest wide-area NEO survey in each era:
# 1998-2004 LINEAR (~19.5, Stokes et al. 2000); 2005-2010 Catalina G96/703 + Spacewatch (~21);
# 2011-2015 Pan-STARRS1 + G96 (~21.5); 2016-2020 upgraded G96 + PS1 (~21.7); 2021-2026 + PS2 (~21.8).
VLIM_BASE = np.array([19.5, 21.0, 21.5, 21.7, 21.8])
HBINS_GRANVIK = np.arange(17, 25)         # lower edges of unit-magnitude orbit bins
I_EDGES = [0.0, 10.0, 25.0, 180.0]
N_I = len(I_EDGES) - 1


def orbit_class(a, e):
    q = a * (1 - e)
    Q = a * (1 + e)
    return np.where(q > 1.017, "AMO", np.where(a > 1.0, "APO", np.where(Q > 0.983, "ATE", "IEO")))


def subset_masks(df):
    return {
        "AMO": (df.cls == "AMO").values,
        "APO": (df.cls == "APO").values,
        "ATE": (df.cls == "ATE").values,
        "i<5deg": (df.i < 5).values,
        "i>25deg": (df.i > 25).values,
        "earthlike(0.9<a<1.1,e<0.2,i<10)": ((df.a > 0.9) & (df.a < 1.1) & (df.e < 0.2) & (df.i < 10)).values,
    }


def load_granvik_sample():
    cache = os.path.join(DATA, "granvik_sample.csv")
    fcache = os.path.join(DATA, "granvik_class_fractions.csv")
    if os.path.exists(cache) and os.path.exists(fcache):
        return pd.read_csv(cache)
    rng = np.random.default_rng(SEED)
    arr = np.loadtxt(gzip.open(os.path.join(RAW, "granvik2018.dat.gz"), "rt"))
    df = pd.DataFrame(arr, columns=["a", "e", "i", "Om", "w", "M", "H"])
    print(f"[sim] Granvik realisation: {len(df)} synthetic NEOs, H range {df.H.min():.2f}-{df.H.max():.2f}")
    df["hb"] = np.floor(df.H).astype(int)
    df["cls"] = orbit_class(df.a, df.e)
    df["icls"] = np.clip(np.searchsorted(I_EDGES, df.i.values, side="right") - 1, 0, N_I - 1)
    parts, fr = [], []
    for hb in HBINS_GRANVIK:
        sub = df[df.hb == hb]
        parts.append(sub.sample(n=min(N_PER_BIN, len(sub)), random_state=int(rng.integers(1 << 31))))
        row = dict(hb=hb, n=len(sub))
        for k in range(N_I):
            row[f"f_icls{k}"] = float((sub.icls == k).mean())
        for name, m in subset_masks(sub).items():
            row["f_" + name] = float(m.mean())
        fr.append(row)
    out = pd.concat(parts).reset_index(drop=True)
    out.to_csv(cache, index=False)
    pd.DataFrame(fr).to_csv(fcache, index=False)
    return out


def get_hists(sample, cache, wfun=None, tag=""):
    """Geometry histograms of the sample orbits for the given elongation weight function.
    The cache is reused only if it was made for the same sample and the same weights."""
    d2r = np.pi / 180
    sig = np.round((S.elong_weight if wfun is None else wfun)(np.radians(S.ELONG_CEN)), 6)
    if os.path.exists(cache):
        z = np.load(cache)
        if z["hist"].shape == (len(sample), S.N_ERA, 2, S.N_G) and "wsig" in z and np.allclose(z["wsig"], sig):
            return z["hist"].astype(float)
    h = S.hist_kepler(sample.a.values, sample.e.values, sample.i.values * d2r, sample.Om.values * d2r,
                      sample.w.values * d2r, sample.M.values * d2r, wfun=wfun)
    np.savez_compressed(cache, hist=h.astype(np.float32), wsig=sig)
    return h


def unpack(theta, rho_fixed, dv):
    """Returns vlim, p, rho for a parameter vector (the hard limits are S.VHARD_BASE + dvh,
    handled by the callers)."""
    vlim = VLIM_BASE + dv
    p = 1.0 / (1.0 + np.exp(-theta[:5]))
    rho = np.full(5, rho_fixed)
    return vlim, p, rho


def main(model="neomod2", rho_fixed=1.0, dv=0.0, dvh=0.0, hist=None, sample=None, quiet=False):
    vhard = S.VHARD_BASE + dvh
    comp = pd.read_csv(os.path.join(DATA, "neo_completeness.csv"))
    comp = comp[comp.H1 >= 17.25].reset_index(drop=True)
    if sample is None:
        sample = load_granvik_sample()
    if hist is None:
        S.load_elong_table()
        hist = get_hists(sample, os.path.join(DATA, "neo_orbit_hists.npz"), wfun=S.make_wfun(variant="base"))
    fr = pd.read_csv(os.path.join(DATA, "granvik_class_fractions.csv")).set_index("hb")
    neo = load_neo()
    neo["cls"] = orbit_class(neo.a, neo.e)
    Hc = ((comp.H1 + comp.H2) / 2).values
    hb_of = np.clip(np.floor(Hc).astype(int), 17, 24)
    dN = comp["dN_" + model].values
    nb = len(comp)
    cnt = np.stack([comp[f"known_in_bin_lt{int(y)}"].values for y in EPOCH_ENDS], axis=1).astype(float)
    obs = cnt / dN[:, None]
    obs_c = np.clip(obs, 1e-5, 0.985)
    wgt = np.sqrt(np.minimum(cnt[:, -1:], 50.0) / 50.0) * np.ones_like(obs)
    idx = {hb: np.where(sample.hb.values == hb)[0] for hb in HBINS_GRANVIK}
    hsubs = [hist[idx[hb_of[b]]] for b in range(nb)]

    def predict(theta):
        vlim, p, rho = unpack(theta, rho_fixed, dv)
        out = np.empty_like(obs)
        for b in range(nb):
            L = S.survival_log(hsubs[b], Hc[b], vlim, p, rho, vhard)
            out[b] = (1.0 - np.exp(np.cumsum(L, axis=1))).mean(axis=0)
        return out

    def resid(theta):
        m = np.clip(predict(theta), 1e-5, 0.985)
        return ((np.log(m / (1 - m)) - np.log(obs_c / (1 - obs_c))) * wgt).ravel()

    p0 = np.array([0.05, 0.1, 0.15, 0.2, 0.25])
    theta0 = np.log(p0 / (1 - p0))
    fit = least_squares(resid, theta0, bounds=(np.full(5, -9.0), np.full(5, 3.0)), diff_step=1e-3, max_nfev=120)
    vlim, p, rho = unpack(fit.x, rho_fixed, dv)
    pred = predict(fit.x)
    rms = float(np.sqrt(np.mean(fit.fun ** 2)))
    if not quiet:
      print(f"[sim:{model}, rho={rho_fixed}, dV={dv:+.1f}, dVhard={dvh:+.1f}] Vlim per era {np.round(vlim, 2)}  nightly p {np.round(p, 4)}  "
          f"weighted rms logit resid {rms:.3f}")
    rows = []
    for b in range(nb):
        row = dict(model=model, rho=rho_fixed, dv=dv, H=Hc[b], obs_final=obs[b, -1], pred_final=pred[b, -1])
        for j, y in enumerate(EPOCH_ENDS):
            row[f"obs_lt{int(y)}"] = obs[b, j]
            row[f"pred_lt{int(y)}"] = pred[b, j]
        rows.append(row)
    curves = pd.DataFrame(rows)

    # ---------------- validation on subsets that were not fitted ----------------
    sample["cls"] = orbit_class(sample.a, sample.e)
    mn, ms = subset_masks(neo), subset_masks(sample)
    val = []
    for name in mn:
        for (h1, h2) in [(17.25, 19.25), (19.25, 21.25), (21.25, 23.25), (23.25, 25.25), (25.25, 27.75)]:
            n_known = int((mn[name] & (neo.H >= h1).values & (neo.H < h2).values).sum())
            n_true = n_pred = 0.0
            for b in range(nb):
                if comp.H1[b] < h1 or comp.H2[b] > h2:
                    continue
                ii = np.where((sample.hb.values == hb_of[b]) & ms[name])[0]
                if len(ii) == 0:
                    continue
                ntrue_b = dN[b] * fr.loc[hb_of[b], "f_" + name]
                n_true += ntrue_b
                n_pred += ntrue_b * S.p_detect(hist[ii], Hc[b], vlim, p, rho, vhard=vhard).mean()
            val.append(dict(model=model, rho=rho_fixed, dv=dv, subset=name, H1=h1, H2=h2, n_known=n_known, n_model_population=n_true,
                            n_predicted=n_pred, compl_observed=n_known / n_true if n_true else np.nan,
                            compl_predicted=n_pred / n_true if n_true else np.nan,
                            ratio_pred_over_known=n_pred / n_known if n_known else np.nan))
    return (dict(model=model, rho_fixed=rho_fixed, dv=dv, dvh=dvh, vhard=vhard.tolist(), vlim=vlim.tolist(), p=p.tolist(), rho=rho.tolist(), rms_logit=rms,
                 width=S.LOGISTIC_W, era_edges=S.ERA_EDGES.tolist()), curves, pd.DataFrame(val))


H_W_MIN, H_W_MAX = 17.25, 27.75      # absolute-magnitude range used to measure the elongation weight
N_ORB_W = 1000                       # orbits per unit-magnitude bin used for the prediction


def poisson95(n):
    """95% interval for the mean of a Poisson variable observed as n (n may be non-integer:
    an effective count)."""
    from scipy.stats import chi2
    lo = 0.0 if n <= 0 else 0.5 * chi2.ppf(0.025, 2 * n)
    hi = 0.5 * chi2.ppf(0.975, 2 * (n + 1))
    return lo, hi


_GEOM = {}        # per absolute-magnitude bin: (g, elongation, declination ok) of the prediction orbits, float32


def _geometry(sample, hb):
    """Daily geometry of the first N_ORB_W sample orbits of magnitude bin hb (kept in memory:
    the weight measurement calls the prediction many times)."""
    if hb not in _GEOM:
        d2r = np.pi / 180
        ex, ey, ez = S.earth_xyz()
        sub = sample[sample.hb == hb].iloc[:N_ORB_W]
        gs, es, oks = [], [], []
        for s0 in range(0, len(sub), 200):
            c = sub.iloc[s0:s0 + 200]
            nmot = K_GAUSS / c.a.values ** 1.5
            M = c.M.values[:, None] * d2r + nmot[:, None] * S.T_DAYS[None, :]
            xyz = helio_xyz(c.a.values[:, None], c.e.values[:, None], c.i.values[:, None] * d2r, c.Om.values[:, None] * d2r,
                            c.w.values[:, None] * d2r, M)
            r, d, alpha, elong, unit = observe(xyz, (ex[None, :], ey[None, :], ez[None, :]))
            gs.append((5.0 * np.log10(r * d) + phase_HG(alpha) + trailing_loss(sky_rate_deg_day(unit))).astype(np.float32))
            es.append(elong.astype(np.float32))
            oks.append((unit[2] * np.cos(S.OBLIQ) + unit[1] * np.sin(S.OBLIQ)) > np.sin(np.radians(S.DEC_MIN)))
        _GEOM[hb] = (np.concatenate(gs), np.concatenate(es), np.concatenate(oks))
    return _GEOM[hb]


def predicted_flat(sample, fit, variant="base"):
    """Number of discoveries per solar-elongation bin that the simulator predicts if the survey
    coverage did not depend on elongation: sum over model objects and days of
        survival(t) x nightly probability x S(V)      (declination above the survey limit),
    where survival(t) = exp(-cumulative hazard) is computed with the currently active elongation
    weight of the variant (so that objects are removed from the pool once discovered)."""
    comp = pd.read_csv(os.path.join(DATA, "neo_completeness.csv"))
    comp = comp[(comp.H1 >= H_W_MIN - 1e-6) & (comp.H2 <= H_W_MAX + 1e-6)]
    vlim, p, vhard = np.array(fit["vlim"]), np.array(fit["p"]), np.array(fit["vhard"])
    Hc_all = 0.5 * (comp.H1.values + comp.H2.values)
    hb_all = np.clip(np.floor(Hc_all).astype(int), 17, 24)
    pred = np.zeros((2, S.N_ELB))          # two interleaved halves of the orbit sample (Monte Carlo error)
    pe, vl, vh = p[S.ERA_IDX][None, :], vlim[S.ERA_IDX][None, :], vhard[S.ERA_IDX][None, :]
    for hb in sorted(set(hb_all)):
        g_all, el_all, ok_all = _geometry(sample, hb)
        n_sub = g_all.shape[0]
        for s0 in range(0, n_sub, 200):
            g = g_all[s0:s0 + 200].astype(float)
            elong = el_all[s0:s0 + 200].astype(float)
            ok = ok_all[s0:s0 + 200]
            w_cur = S.elong_weight(elong, variant=variant) * ok
            ie = np.clip(np.searchsorted(S.ELONG_EDGES, np.degrees(elong), side="right") - 1, 0, S.N_ELB - 1).ravel()
            for Hc, dN in zip(Hc_all[hb_all == hb], comp.dN_neomod2.values[hb_all == hb]):
                V = Hc + g
                Sg = np.where(V <= vh, 1.0 / (1.0 + np.exp(np.clip((V - vl) / S.LOGISTIC_W, -60, 60))), 0.0)
                pn = pe * Sg
                haz = -w_cur * np.log1p(-pn)
                surv = np.exp(-(np.cumsum(haz, axis=1) - haz))             # survival up to the previous day
                contrib = surv * pn * ok
                for half in (0, 1):
                    pred[half] += dN / n_sub * np.bincount(ie.reshape(contrib.shape)[half::2].ravel(), weights=contrib[half::2].ravel(),
                                                           minlength=S.N_ELB)
    return pred


def measure_elong_weight(sample, fit, variant="base"):
    """Coverage weight versus solar elongation: real first observations per elongation bin (s01b)
    divided by the prediction under flat coverage, normalised to 1 above 150 deg."""
    disc = pd.read_csv(os.path.join(DATA, "discovery_geometry.csv"))
    pred2 = predicted_flat(sample, fit, variant)
    pred = pred2.sum(axis=0)
    # objects closer than 0.005 au on the first night are left out: their sky position changes by tens of
    # degrees within hours, so the elongation computed for 08:00 UT is not that of the observation
    d = disc[(disc.H >= H_W_MIN) & (disc.H < H_W_MAX) & (disc.delta >= 0.005)].copy()
    d["ib"] = np.clip(np.searchsorted(S.ELONG_EDGES, d.elong.values, side="right") - 1, 0, S.N_ELB - 1)
    n_raw = np.bincount(d.ib.values, minlength=S.N_ELB).astype(float)
    n_w = np.bincount(d.ib.values, weights=d.weight.values, minlength=S.N_ELB)
    n_w2 = np.bincount(d.ib.values, weights=d.weight.values ** 2, minlength=S.N_ELB)
    n_eff = np.where(n_w2 > 0, n_w ** 2 / np.maximum(n_w2, 1e-12), 0.0)       # effective number of objects
    ref = S.ELONG_EDGES[:-1] >= 150.0
    norm = n_w[ref].sum() / pred[ref].sum()
    # Monte Carlo error of the prediction (finite orbit sample): the two halves of the sample give two
    # independent estimates of (prediction in the bin) / (prediction above 150 deg)
    ra, rb = pred2[0] / pred2[0][ref].sum(), pred2[1] / pred2[1][ref].sum()
    mc = np.abs(ra - rb) / np.maximum(ra + rb, 1e-300)                         # ~1 sigma relative error of the full sample
    rows = []
    for b in range(S.N_ELB):
        lo, hi = poisson95(n_eff[b])
        scale = (n_w[b] / n_eff[b]) if n_eff[b] > 0 else 1.0
        if n_eff[b] > 0 and lo > 0:
            # 95% bounds: Poisson error of the real objects and twice the Monte Carlo error, in quadrature (log space)
            lo = n_eff[b] * np.exp(-np.hypot(np.log(n_eff[b] / lo), 2 * mc[b]))
            hi = n_eff[b] * np.exp(np.hypot(np.log(hi / n_eff[b]), 2 * mc[b]))
        rows.append(dict(variant=variant, elong_lo=S.ELONG_EDGES[b], elong_hi=S.ELONG_EDGES[b + 1], elong_centre=S.ELONG_CEN[b],
                         n_objects=int(n_raw[b]), n_weighted=n_w[b], n_effective=n_eff[b],
                         V_median=float(d.V[d.ib == b].median()) if n_raw[b] else np.nan,
                         H_median=float(d.H[d.ib == b].median()) if n_raw[b] else np.nan,
                         predicted_flat_coverage=pred[b] * norm, mc_rel_err=mc[b],
                         w_c=n_w[b] / (pred[b] * norm), w_lo=lo * scale / (pred[b] * norm), w_hi=hi * scale / (pred[b] * norm)))
    return pd.DataFrame(rows)


def rel_change(new, old):
    big = (old.n_objects >= 10).values
    return float(np.max(np.abs(new.w_c.values[big] - old.w_c.values[big]) / np.maximum(old.w_c.values[big], 1e-9)))


def iterate_weight(sample, fit, variant="base", nmax=6, tol=0.02):
    """Fixed-point iteration of the measured weight of one survey variant for a given simulator
    fit.  Starts from whatever weight is active for that variant (the assumed ramp in the very
    first round; the base table for the other depth variants)."""
    tab = measure_elong_weight(sample, fit, variant)
    for it in range(nmax):
        S.set_elong_table(variant, tab)
        new = measure_elong_weight(sample, fit, variant)
        ch = rel_change(new, tab)
        tab = new
        print(f"[sim] elongation weight ({variant}): fixed-point iteration {it + 1}, largest relative change {ch:.3f}", flush=True)
        if ch < tol:
            break
    S.set_elong_table(variant, tab)
    return tab


def get_hists_multi(sample, cache, names):
    """Geometry histograms of the sample orbits for the measured weights of several survey
    variants, in one geometry pass.  Cached (in the raw directory: about 30 MB)."""
    d2r = np.pi / 180
    sig = {v: np.round(S.elong_weight(np.radians(S.ELONG_CEN), variant=v), 6) for v in names}
    if os.path.exists(cache):
        z = np.load(cache)
        if all(("hist_" + v) in z and z["hist_" + v].shape[0] == len(sample) and np.allclose(z["wsig_" + v], sig[v]) for v in names):
            return {v: z["hist_" + v].astype(float) for v in names}
    h = S.hist_kepler(sample.a.values, sample.e.values, sample.i.values * d2r, sample.Om.values * d2r,
                      sample.w.values * d2r, sample.M.values * d2r, wfun={v: S.make_wfun(variant=v) for v in names})
    np.savez_compressed(cache, **{"hist_" + v: h[v].astype(np.float32) for v in names}, **{"wsig_" + v: sig[v] for v in names})
    return h


VARIANTS = {"base": ("neomod2", 1.0, 0.0, 0.0), "shallow": ("neomod2", 1.0, -0.5, 0.0),
            "deep": ("neomod2", 1.0, 1.0, 0.0), "vdeep": ("neomod2", 1.0, 1.5, 0.5),
            "hard22p5": ("neomod2", 1.0, 0.0, -0.5),
            "hc21": ("hc21", 1.0, 0.0, 0.0), "ecliptic": ("neomod2", 0.3, 0.0, 0.0)}
DEPTH_VARIANTS = ["shallow", "deep", "vdeep", "hard22p5"]      # these get their own measured elongation weight


def run_all():
    sample = load_granvik_sample()
    pd.set_option("display.width", 250)
    # ---- stage 1: assumed ramp ----
    S.unload_elong_table()
    h1 = get_hists(sample, os.path.join(RAW, "neo_orbit_hists_stage1.npz"), wfun=S.elong_weight_ramp)
    f1, _, _ = main("neomod2", 1.0, 0.0, 0.0, hist=h1, sample=sample)
    json.dump(f1, open(os.path.join(DATA, "survey_fit_stage1.json"), "w"), indent=1)
    del h1
    # ---- stage 2: measure the elongation weight of the base variant.  The survival term of the prediction
    #      depends on the weight itself, so the measurement is iterated to a fixed point; then the simulator is
    #      re-fitted with the measured weight and the measurement repeated with that fit (at most three rounds). ----
    fit_w, tabs = f1, {}
    for rnd in range(3):
        tab = iterate_weight(sample, fit_w, "base")
        hist = get_hists(sample, os.path.join(DATA, "neo_orbit_hists.npz"), wfun=S.make_wfun(variant="base"))
        fb, _, _ = main("neomod2", 1.0, 0.0, 0.0, hist=hist, sample=sample, quiet=True)
        ch = rel_change(measure_elong_weight(sample, fb, "base"), tab)
        print(f"[sim] elongation weight (base), round {rnd + 1}: largest relative change of w_c when re-measured with the simulator "
              f"re-fitted on these weights (bins with >= 10 objects): {ch:.3f}", flush=True)
        if ch < 0.05:
            break
        fit_w = fb
    tab["closure_rel_change"] = ch
    tabs["base"] = tab
    # ---- stage 3: the other depth variants.  What the discoveries at low elongation constrain is the product
    #      (weight x nightly probability x efficiency), so the weight is measured again with each variant's own
    #      limiting magnitudes and nightly probabilities (first fitted with the base weights). ----
    for name in DEPTH_VARIANTS:
        model, rho, dv, dvh = VARIANTS[name]
        f0, _, _ = main(model, rho, dv, dvh, hist=hist, sample=sample, quiet=True)
        S.set_elong_table(name, tabs["base"])
        tabs[name] = iterate_weight(sample, f0, name)
    hv = get_hists_multi(sample, os.path.join(RAW, "neo_orbit_hists_variants.npz"), DEPTH_VARIANTS)
    # ---- stage 4: final fits ----
    fits, curves, vals = {}, [], []
    for name, (model, rho, dv, dvh) in VARIANTS.items():
        f, c, v = main(model, rho, dv, dvh, hist=hv.get(name, hist), sample=sample)
        f["variant"] = name
        f["elongation_weight_table"] = name if name in tabs else "base"
        c["variant"] = name
        v["variant"] = name
        fits[name] = f
        curves.append(c)
        vals.append(v)
    del hv
    json.dump(fits, open(os.path.join(DATA, "survey_fit.json"), "w"), indent=1)
    # closure of the depth variants: one re-measurement with the final fit
    for name in DEPTH_VARIANTS:
        ch = rel_change(measure_elong_weight(sample, fits[name], name), tabs[name])
        tabs[name]["closure_rel_change"] = ch
        print(f"[sim] elongation weight ({name}): largest relative change when re-measured with the final fit: {ch:.3f}", flush=True)
    T = pd.concat([tabs[v] for v in ["base"] + DEPTH_VARIANTS], ignore_index=True)
    T["w_assumed_ramp"] = S.elong_weight_ramp(np.radians(T.elong_centre.values))
    T.to_csv(S.ELONG_TABLE_FILE, index=False)
    print("[sim] measured survey coverage weight versus solar elongation (17.25 <= H < 27.75; normalised to >= 150 deg), base variant:")
    print(T[T.variant == "base"].drop(columns=["variant"]).round(4).to_string(index=False))
    print("[sim] the same for every depth variant (w_c; 95% bounds in the file):")
    print(T.pivot(index="elong_centre", columns="variant", values="w_c")[["base"] + DEPTH_VARIANTS].round(4).to_string())
    c = pd.concat(curves)
    c.to_csv(os.path.join(DATA, "survey_fit_curves.csv"), index=False)
    v = pd.concat(vals)
    v.to_csv(os.path.join(DATA, "survey_validation.csv"), index=False)
    print(c[c.variant == "base"][["H", "obs_lt2005", "pred_lt2005", "obs_lt2016", "pred_lt2016",
                                                     "obs_final", "pred_final"]].round(4).to_string(index=False))
    print(v[v.variant.isin(["base"])].drop(columns=["model", "rho", "dv"]).round(4).to_string(index=False))
    rng_ = v[v.n_known >= 10].groupby("subset").ratio_pred_over_known.agg(["min", "max"])
    print("[sim] predicted / catalogued, range over H bins and variants (bins with >= 10 known objects):")
    print(rng_.round(2).to_string())


if __name__ == "__main__":
    run_all()