"""Build N_cov(L, nu): the number of stars for which a continuous narrowband transmitter of
EIRP >= L at frequency nu would have been detected by the published null searches.

Inputs : data/radio_model_components.csv (hand-compiled from the literature; see surveys_compilation.csv),
         raw VizieR tables (fetch_data.py), Galaxy star-count model (galaxy.py).
Outputs: data/ncov_by_component.csv, data/neff_table.csv, data/coverage_intervals.csv,
         data/galaxy_model_validation.csv, data/ata_targets_with_distances.csv, data/bl_targets.csv,
         data/drift_factors.csv, data/bycatch_calibration.csv, data/bycatch_pointings.csv

Explicit per-survey factors added in the revision
  * drift: a transmitter with fractional drift rate a/c (ACC) smears a carrier over fdot*t/dnu channels of a
    survey whose spectra are dnu wide and accumulated over t without drift correction; the flux threshold is
    multiplied by max(1, fdot*t/dnu) (columns chan_hz, tres_s of the component table).  SETI@home uses its
    published barycentric / non-barycentric sensitivity columns instead.
  * Doppler window: for a beacon at the HI rest frequency of its own star ('HI_rest'), sky surveys with a
    narrow window (META +-42 km/s, Big Ear +-53 km/s, SETI@home -178..+350 km/s) only count stars whose
    line-of-sight velocity falls inside the window (kinematic model in galaxy.py).
  * in-beam stars: the analytic star counts are rescaled, distance shell by distance shell, to the published
    Besancon counts of Mason, Garrett & Siemion (2025, Table 1).
"""
import os
import numpy as np
import pandas as pd

from common import DATA, RAW, LGRID, PC, FOURPI, eirp_min, dmax_pc
import galaxy as G

VIZ = os.path.join(RAW, "vizier")
ARCMIN = np.pi / 180.0 / 60.0
DEG2 = (np.pi / 180.0) ** 2
C_KMS = 299792.458
F_HI = 1420.40575          # MHz

# ---------------------------------------------------------------- frequency priors (MHz)
PRIORS = {
    # name: (kind, f_lo, f_hi)
    "HI_line": ("uniform", 1420.2, 1420.6),        # 21-cm line as received, +-0.2 MHz: beacon tuned to OUR frame
    "HI_rest": ("line", F_HI, F_HI),               # 21-cm line in the rest frame of the transmitter's own star
    "water_hole": ("uniform", 1420.0, 1666.0),     # H to OH
    "log_1_10GHz": ("log", 1000.0, 10000.0),       # terrestrial microwave window
    "log_0p1_100GHz": ("log", 100.0, 100000.0),    # wide plausible band
}
AUX_PRIORS = {"uniform_1_2GHz": ("uniform", 1000.0, 2000.0)}   # only for the comparison with Grimaldi & Marcy (2018)

# ---------------------------------------------------------------- transmitter drift classes: a/c in 1/s
ACC = {
    "inertial": 0.0,       # acceleration removed by the transmitter, or free-flying beyond ~3 AU (a < 5e-4 m/s2)
    "orbit": 1.98e-11,     # Earth-like orbital acceleration only (5.9e-3 m/s2): 0.028 Hz/s at 1.42 GHz
    "surface": 1.13e-10,   # Earth-like equatorial rotation (3.4e-2 m/s2): 0.16 Hz/s at 1.42 GHz
}


def prior_weight(kind, f1, f2, a, b):
    if kind == "line":
        return 1.0 if a <= f1 < b else 0.0
    lo, hi = max(a, f1), min(b, f2)
    if hi <= lo:
        return 0.0
    if kind == "uniform":
        return (hi - lo) / (f2 - f1)
    return np.log(hi / lo) / np.log(f2 / f1)


def parse_bands(s):
    out = []
    for part in str(s).split(";"):
        a, b = part.split("-")
        out.append((float(a), float(b)))
    return out


def drift_loss(c, drift):
    """Factor by which the flux threshold of component row c rises for a transmitter of drift class `drift`,
    evaluated at the geometric centre of the component's band."""
    if ACC[drift] == 0.0 or c["chan_hz"] in ("", None) or pd.isna(c["chan_hz"]):
        return 1.0
    b = parse_bands(c["bands_mhz"])
    fmid = np.sqrt(min(x[0] for x in b) * max(x[1] for x in b)) * 1e6
    return float(max(1.0, ACC[drift] * fmid * float(c["tres_s"]) / float(c["chan_hz"])))


# ---------------------------------------------------------------- helpers for N(L)
def n_from_eirp_list(eirp, weights=None):
    """Cumulative (weighted) number of stars with EIRP_min <= L on LGRID."""
    eirp = np.asarray(eirp, float)
    w = np.ones_like(eirp) if weights is None else np.asarray(weights, float)
    o = np.argsort(eirp)
    cs = np.concatenate([[0.0], np.cumsum(w[o])])
    return cs[np.searchsorted(eirp[o], LGRID, side="right")]


def n_from_shells(C, smin, scint=False):
    """Stars per distance shell C[j] (already multiplied by solid angle) -> N(L): step at d_max(L)
    (or exp(-(d/dmax)^2) weighting for saturated scintillation)."""
    dm = dmax_pc(LGRID, smin)
    if scint:
        return np.array([np.sum(C * np.exp(-(G.D_MID / x) ** 2)) for x in dm])
    cum = np.concatenate([[0.0], np.cumsum(C)])
    return np.interp(np.log(dm), np.log(G.D_EDGES), cum)


def n_bycatch_dm(Cs, dm, fwhm_arcmin, scint=False):
    """Cs[..., j]: stars per steradian per distance shell -> expected number of in-beam stars detectable when a
    star on the beam axis is detectable out to dm (pc).  Gaussian beam, stars counted out to the half-power
    radius; an off-axis star at gain g is detectable out to dm*sqrt(g).  dm may be an array (result gets a
    trailing axis)."""
    th2 = (0.5 * fwhm_arcmin * ARCMIN) ** 2
    dm = np.atleast_1d(np.asarray(dm, float))
    u = (np.arange(40) + 0.5) / 40.0         # theta^2/theta_h^2, uniform in area
    out = np.zeros(Cs.shape[:-1] + (len(dm),))
    for k, x in enumerate(dm):
        a = (G.D_MID / x) ** 2               # required gain fraction
        if scint:
            w = np.mean(np.exp(-a[:, None] * 2.0 ** u[None, :]), axis=1)
        else:
            w = np.clip(np.log(1.0 / np.maximum(a, 1e-300)) / np.log(2.0), 0.0, 1.0) * (a < 1.0)
        out[..., k] = np.pi * th2 * (Cs @ w)
    return out


def n_bycatch(Cs, smin, fwhm_arcmin, scint=False):
    return n_bycatch_dm(Cs, dmax_pc(LGRID, smin), fwhm_arcmin, scint=scint)


def n_anchors(spec, sfac=1.0):
    """Published (EIRP, number of stars) anchor points, log-log interpolated; zero below the first
    anchor, constant above the last.  sfac shifts the EIRP anchors (threshold scaling)."""
    pts = [p.split("=") for p in spec.split(";")]
    L = np.array([float(p[0]) for p in pts]) * sfac; N = np.array([float(p[1]) for p in pts])
    out = np.zeros(len(LGRID))
    m = LGRID >= L[0] * 0.999
    out[m] = 10 ** np.interp(np.log10(LGRID[m]), np.log10(L), np.log10(N))
    return out


# ---------------------------------------------------------------- data loaders
def load_ws():
    d = pd.read_csv(os.path.join(VIZ, "ws2020_catalog.csv"))
    d.columns = [c.strip() for c in d.columns]
    for c in ("InstS", "InstT", "Price-target?", "Target"):
        d[c] = d[c].astype(str).str.strip()
    return d


def bl_target_rows(ws, inst):
    """Rows of the Wlodarczyk-Sroka+ 2020 catalogue flagged as primary targets for one receiver, one row per
    target name (the flag is set on more than one Gaia source for some targets; the lowest EIRP_min is kept)."""
    t = ws[(ws["InstS"] == inst) & (ws["Price-target?"] == "true")]
    return t.sort_values("min-detec-EIRP").drop_duplicates("Target"), len(t)


def load_ata():
    """Harp+ 2016 targets with Gaia DR3 distances from the positional cross-match (20 arcsec).
    Match rule: brightest Gaia source within 20" for HabCat/Tycho/HZ targets (bright, possibly
    high proper motion); nearest source within 5" for Kepler targets (faint, crowded field)."""
    t7 = pd.read_csv(os.path.join(VIZ, "harp2016_table7.csv"))
    t7["id"] = "t7_" + t7["recno"].astype(str)
    t6 = pd.read_csv(os.path.join(VIZ, "harp2016_table6.csv"))
    t6["id"] = "t6_" + t6["recno"].astype(str); t6["S"] = "Z"; t6["BW"] = 8000.0
    t = pd.concat([t7[["id", "S", "Name", "BW"]], t6[["id", "S", "Name", "BW"]]], ignore_index=True)
    x = pd.read_csv(os.path.join(VIZ, "harp2016_x_gaiadr3.csv"))
    x = x.merge(t[["id", "S"]], on="id")
    kep = x["S"] == "K"
    xa = x[kep & (x["angDist"] <= 5.0)].sort_values("angDist").groupby("id").head(1)
    xb = x[~kep].sort_values("Gmag").groupby("id").head(1)
    m = pd.concat([xa, xb])[["id", "angDist", "Plx", "e_Plx", "Gmag"]]
    t = t.merge(m, on="id", how="left")
    good = (t["Plx"] > 0) & (t["Plx"] / t["e_Plx"] > 5)
    t["d_pc"] = np.where(good, 1000.0 / t["Plx"], np.nan)
    t["fweight"] = np.clip(t["BW"] / 8000.0, 0, 1)
    return t


def cap_los(l0, b0, area_deg2):
    """Sample lines of sight over a small circular cap (equal-area rings). Returns l, b, weights(sr)."""
    theta = np.degrees(np.arccos(1.0 - area_deg2 * DEG2 / (2 * np.pi)))
    ls, bs = [l0], [b0]
    nr = 6
    rings = [(theta * np.sqrt((i + 0.5) / nr), 6 * (i + 1)) for i in range(nr)]
    for r, n in rings:
        ang = np.arange(n) * 2 * np.pi / n
        bs += list(b0 + r * np.sin(ang)); ls += list(l0 + r * np.cos(ang) / np.cos(np.radians(b0)))
    ls = np.array(ls) % 360.0; bs = np.clip(np.array(bs), -89.9, 89.9)
    w = np.full(len(ls), area_deg2 * DEG2 / len(ls))
    return ls, bs, w


# ---------------------------------------------------------------- Mason+ 2025 calibration of in-beam counts
# Mason, Garrett & Siemion (2025, arXiv:2511.20231) Table 1: cumulative number of Besancon-model stars (N_B)
# inside the FWHM of the Breakthrough Listen pointings with EIRP_min below 1e14 ... 1e18 W, and the distance d
# at which a star on the beam axis has that EIRP_min.  Their simulation stops at 25 kpc and omits two
# pointings 'along b ~ 0 towards the Galactic Centre' that were too dense to simulate.
MASON = {"GBT": dict(d=[290.0, 930.0, 2940.0, 9310.0, 25000.0], NB=[1828.0, 36071.0, 367616.0, 3091788.0, 5115832.0],
                     NG=[2599, 25192, 147976, 246245, 246492]),
         "Parkes": dict(d=[160.0, 500.0, 1600.0, 5050.0, 16000.0], NB=[98.0, 2019.0, 26613.0, 267747.0, 1178528.0],
                        NG=[82, 695, 5225, 30434, 40592])}
D_SIM_MAX = 25000.0


def _shell_factor(fk, dnodes):
    """Piecewise-constant factor per distance shell: fk[i] applies to (dnodes[i-1], dnodes[i]]; the last value
    is kept beyond the last node."""
    idx = np.searchsorted(np.asarray(dnodes), G.D_MID, side="left")
    return np.asarray(fk)[np.minimum(idx, len(fk) - 1)]


def calibrate_bycatch(sets, tel, drop=None):
    """sets: list of (Cs_pt (n_pointings, n_shell), fwhm_arcmin).  Finds shell factors such that the model,
    counted exactly as Mason+ count (Gaussian taper inside the FWHM, d < 25 kpc), reproduces their cumulative
    N_B.  drop: optional list of boolean masks (one per set) of pointings left out of the comparison."""
    dn, NB = np.array(MASON[tel]["d"]), np.array(MASON[tel]["NB"])
    insim = (G.D_MID <= D_SIM_MAX)
    tot = []
    for i, (Cs, fw) in enumerate(sets):
        keep = np.ones(len(Cs), bool) if drop is None else ~drop[i]
        tot.append((Cs[keep].sum(axis=0) * insim, fw))
    model = lambda f: sum(n_bycatch_dm(c * f, dn, fw) for c, fw in tot)
    raw = model(np.ones(len(G.D_MID)))
    fk = np.ones(len(dn))
    for _ in range(12):
        cur = model(_shell_factor(fk, dn))
        fk = fk * np.diff(np.concatenate([[0.0], NB])) / np.diff(np.concatenate([[0.0], cur]))
    fin = model(_shell_factor(fk, dn))
    return dict(f=_shell_factor(fk, dn), fk=fk, raw=raw, fin=fin, d=dn, NB=NB)


# ---------------------------------------------------------------- main builder
_PREP = None


def prep():
    """Heavy, option-independent precomputation (cached): catalogues and line-of-sight star counts."""
    global _PREP
    if _PREP is not None:
        return _PREP
    comps = pd.read_csv(os.path.join(DATA, "radio_model_components.csv"), keep_default_na=False)
    for col in ("smin_wm2", "fwhm_arcmin", "chan_hz", "tres_s"):
        comps[col] = pd.to_numeric(comps[col], errors="coerce")
    G.calibrate_local()
    ws = load_ws()
    isa = pd.read_csv(os.path.join(VIZ, "isaacson2017_table1.csv"))
    isa_d = np.sort(isa["Dist"].dropna().values); isa_d = isa_d[isa_d <= 50.0]
    ata = load_ata()
    gcs = pd.read_csv(os.path.join(VIZ, "baumgardt2018_gc.csv"))
    gcs["Cluster"] = gcs["Cluster"].astype(str).str.strip()
    # Doppler windows (km/s, LSR) of the sky surveys whose band contains the HI line
    vwin = {}
    for _, c in comps.iterrows():
        if c["geom"] in ("sky", "skytiers"):
            for lo, hi in parse_bands(c["bands_mhz"]):
                if lo < F_HI < hi and (hi - lo) < 10.0:
                    vwin[c["comp_id"]] = (round(C_KMS * (F_HI - hi) / F_HI, 1), round(C_KMS * (F_HI - lo) / F_HI, 1))
    windows = sorted(set(vwin.values()))
    # sky-grid accumulation for all 'sky', 'skytiers' and large 'cap' components
    gl, gb, gom = G.sky_grid()
    gra, gdec = G.gal_to_equ(gl, gb)
    sky_w = {}
    for _, c in comps.iterrows():
        a = str(c["geom_args"]).split(";")
        if c["geom"] == "cap" and float(a[2]) > 20.0:
            l0, b0, area = float(a[0]), float(a[1]), float(a[2])
            cs = (np.sin(np.radians(gb)) * np.sin(np.radians(b0)) +
                  np.cos(np.radians(gb)) * np.cos(np.radians(b0)) * np.cos(np.radians(gl - l0)))
            w = gom * (cs >= 1.0 - area * DEG2 / (2 * np.pi))
            sky_w[c["comp_id"]] = w * (area * DEG2 / w.sum())
    sky_w["__allsky__"] = gom
    # declination zones bounded by the edges of all sky surveys: inside a zone every sky survey either
    # covers it or not, so the union of overlapping sky surveys can be taken zone by zone
    zed = sorted({-90.0, 90.0} | {float(x) for _, c in comps.iterrows() if c["geom"] in ("sky", "skytiers")
                                  for x in str(c["geom_args"]).split(";")[:2]})
    zones = [(zed[i], zed[i + 1]) for i in range(len(zed) - 1)]
    vkeys = ["__allsky__"]
    for i, (lo, hi) in enumerate(zones):
        sky_w[f"__zone{i}__"] = gom * ((gdec >= lo) & (gdec < hi if hi < 90.0 else gdec <= hi))
        vkeys.append(f"__zone{i}__")
    shells = {k: np.zeros(len(G.D_MID)) for k in sky_w}
    shells_v = {(w, k): np.zeros(len(G.D_MID)) for w in windows for k in vkeys}
    step = 8000
    for i0 in range(0, len(gl), step):
        sl = slice(i0, i0 + step)
        cnt, cntv = G.los_counts_vwin(gl[sl], gb[sl], windows)
        for k, w in sky_w.items():
            shells[k] += w[sl] @ cnt
        for w, cv in zip(windows, cntv):
            for k in vkeys:
                shells_v[(w, k)] += sky_w[k][sl] @ cv
    n_mw_total = float(shells["__allsky__"].sum())
    bl = {}
    for _, c in comps.iterrows():
        a = str(c["geom_args"]).split(";")
        if c["geom"] == "bl_bycatch":
            w = ws[ws["InstS"] == a[0]]
            pt = w.groupby("Target")[["RATdeg", "DETdeg"]].first()
            l, b = G.equ_to_gal(pt["RATdeg"].values, pt["DETdeg"].values)
            nt = w[w["Price-target?"] != "true"]
            bl[c["comp_id"]] = dict(Cs_pt=G.los_counts_per_sr(l, b), l=l, b=b, names=list(pt.index), n_pointings=len(pt),
                                    fwhm=float(c["fwhm_arcmin"]), gaia_inbeam=n_from_eirp_list(nt["min-detec-EIRP"].values))
        elif c["geom"] == "cap" and float(a[2]) <= 20.0:
            l, b, w = cap_los(float(a[0]), float(a[1]), float(a[2]))
            shells[c["comp_id"]] = w @ G.los_counts_per_sr(l, b)
    # --- calibration of the in-beam counts on Mason+ 2025 (GBT = L-band pointings at 8.4' + S-band-only at 5.5')
    L, S, P_ = bl["bl_gbtL_bycatch"], bl["bl_gbtS_bycatch"], bl["bl_pks10_bycatch"]
    s_only = ~np.isin(S["names"], L["names"])
    gbt_sets = [(L["Cs_pt"], 8.4), (S["Cs_pt"][s_only], 5.5)]
    cal = {"none": {"GBT": np.ones(len(G.D_MID)), "Parkes": np.ones(len(G.D_MID))}, "mason": {}, "mason_excl2": {}}
    cal_diag = {}
    cg = calibrate_bycatch(gbt_sets, "GBT"); cp = calibrate_bycatch([(P_["Cs_pt"], 6.4)], "Parkes")
    cal["mason"] = {"GBT": cg["f"], "Parkes": cp["f"]}
    cal_diag["mason"] = {"GBT": cg, "Parkes": cp}
    # variant: leave the two densest low-latitude inner-Galaxy pointings out of the comparison (Mason+ omitted two)
    tot_pt = np.concatenate([L["Cs_pt"].sum(axis=1) * 8.4 ** 2, S["Cs_pt"][s_only].sum(axis=1) * 5.5 ** 2])
    lb = np.concatenate([np.c_[L["l"], L["b"]], np.c_[S["l"][s_only], S["b"][s_only]]])
    inner = (np.abs(lb[:, 1]) < 2.0) & ((lb[:, 0] < 45.0) | (lb[:, 0] > 315.0))
    top2 = np.argsort(np.where(inner, tot_pt, -1.0))[-2:]
    dropm = np.zeros(len(tot_pt), bool); dropm[top2] = True
    cg2 = calibrate_bycatch(gbt_sets, "GBT", drop=[dropm[:len(L["names"])], dropm[len(L["names"]):]])
    cal["mason_excl2"] = {"GBT": cg2["f"], "Parkes": cp["f"]}
    cal_diag["mason_excl2"] = {"GBT": cg2, "Parkes": cp}
    names_all = list(L["names"]) + list(np.array(S["names"])[s_only])
    cal_diag["dropped_pointings"] = [(names_all[i], float(lb[i, 0]), float(lb[i, 1])) for i in top2]
    _PREP = dict(comps=comps, ws=ws, isa_d=isa_d, ata=ata, gcs=gcs, shells=shells, shells_v=shells_v, vwin=vwin, bl=bl,
                 zones=zones, n_mw_total=n_mw_total, stars_per_msun=G.n_local() / G.RHO_STAR_LOCAL, cal=cal,
                 cal_diag=cal_diag, s_only=s_only)
    return _PREP


def build(scint=False, smin_scale=None, n_scale=None, drift="inertial", bycatch_cal="mason"):
    """Evaluate N(L) for every component.

    scint      : apply saturated-scintillation detection weighting exp(-(d/dmax)^2) to model-counted
                 (in-beam and sky-survey) stars instead of a sharp distance cut.
    smin_scale : dict {comp_id | 'group_<g>' | 'tier_<t>' | 'all': factor} multiplying the flux threshold.
    n_scale    : dict {comp_id | 'group_<g>': factor} multiplying the star counts.
    drift      : transmitter drift class ('inertial', 'orbit', 'surface'); see ACC and drift_loss().
    bycatch_cal: 'mason' (in-beam counts rescaled to Mason+ 2025), 'mason_excl2', or 'none' (raw analytic model).
    Returns (dict comp_id -> dict(N_mw, N_mw_v, N_x, bands, group, tier, label, geom[, zones, zones_v]), meta).
    The *_v arrays count only stars inside the survey's Doppler window (used for the 'HI_rest' prior)."""
    P = prep()
    smin_scale = smin_scale or {}
    n_scale = n_scale or {}
    comps, ws, ata, gcs, shells = P["comps"], P["ws"], P["ata"], P["gcs"], P["shells"]
    out = {}
    diag = {}
    for _, c in comps.iterrows():
        cid, geom = c["comp_id"], c["geom"]
        gargs = c["geom_args_drift"] if (drift != "inertial" and str(c["geom_args_drift"]) != "") else c["geom_args"]
        a = str(gargs).split(";")
        dfac = drift_loss(c, drift)
        sfac = dfac * smin_scale.get(cid, smin_scale.get("tier_" + str(c["tier"]), smin_scale.get("group_" + c["group"],
                                     smin_scale.get("all", 1.0))))
        smin = (float(c["smin_wm2"]) if pd.notna(c["smin_wm2"]) else np.nan) * sfac
        N_mw = np.zeros(len(LGRID)); N_x = np.zeros(len(LGRID)); N_v = None
        if geom == "bl_targets":
            t, nrows = bl_target_rows(ws, a[0])
            N_mw = n_from_eirp_list(t["min-detec-EIRP"].values * sfac)
            diag[cid] = dict(n_rows_flagged=nrows, n_targets=len(t))
        elif geom == "bl_bycatch":
            B = P["bl"][cid]
            f = P["cal"][bycatch_cal]["Parkes" if a[0].startswith("Parkes") else "GBT"]
            N_mw = n_bycatch(B["Cs_pt"].sum(axis=0) * f, smin, B["fwhm"], scint=scint)
            diag[cid] = dict(n_pointings=B["n_pointings"], gaia_inbeam=B["gaia_inbeam"])
        elif geom == "isaacson_scaled":
            # a = [stars < 50 pc, stars < 1 kpc, fraction of the band searched]
            dm = dmax_pc(LGRID, smin)
            N_mw = float(a[0]) * np.searchsorted(P["isa_d"], dm, side="right") / len(P["isa_d"])
            if len(a) > 1:
                N_mw = N_mw + (float(a[1]) - float(a[0])) * np.clip(np.log(dm / 50.0) / np.log(1000.0 / 50.0), 0.0, 1.0)
            if len(a) > 2:
                N_mw = N_mw * float(a[2])
        elif geom == "ata":
            ok = ata["d_pc"].notna()
            N_mw = n_from_eirp_list(eirp_min(ata.loc[ok, "d_pc"].values, smin), ata.loc[ok, "fweight"].values)
            diag[cid] = dict(n_total=len(ata), n_with_distance=int(ok.sum()),
                             mean_fweight=float(ata["fweight"].mean()))
        elif geom == "volume":
            N_mw = float(a[0]) * np.clip(dmax_pc(LGRID, smin) / float(a[1]), 0, 1) ** 3
        elif geom == "list":
            t = pd.read_csv(os.path.join(DATA, a[0]))
            N_mw = n_from_eirp_list(eirp_min(t["d_pc"].values, smin))
        elif geom == "gcs":
            g = gcs[gcs["Cluster"].isin([x.strip() for x in a])]
            nst = 0.5 * g["Mass"].values * P["stars_per_msun"]
            N_mw = n_from_eirp_list(eirp_min(g["Dist"].values * 1000.0, smin), nst)
            diag[cid] = dict(n_clusters=len(g), stars=float(nst.sum()))
        elif geom == "anchors":
            N_mw = n_anchors(c["geom_args"], sfac)
        elif geom == "xgal":
            N_x = n_anchors(c["geom_args"], sfac)
        elif geom == "cap":
            C = shells[cid]
            N_mw = n_from_shells(C, smin, scint=scint)
            diag[cid] = dict(stars_to_8p5kpc=float(C[G.D_MID < 8500].sum()), stars_all=float(C.sum()))
        elif geom in ("sky", "skytiers"):
            # a = [dec_min, dec_max, filling fraction, (S_min=cumulative sky fraction; ...)]
            lo, hi, frac = float(a[0]), float(a[1]), float(a[2])
            tiers = [(smin, 1.0)] if geom == "sky" else [(float(t.split("=")[0]) * sfac, float(t.split("=")[1])) for t in a[3:]]
            vw = P["vwin"].get(cid)
            zN = np.zeros((len(P["zones"]), len(LGRID))); zNv = np.zeros_like(zN)
            for iz, (zlo, zhi) in enumerate(P["zones"]):
                if zlo >= lo and zhi <= hi:
                    prev = 0.0
                    for s_, f_ in tiers:
                        zN[iz] += frac * (f_ - prev) * n_from_shells(shells[f"__zone{iz}__"], s_, scint=scint)
                        cv = P["shells_v"][(vw, f"__zone{iz}__")] if vw is not None else shells[f"__zone{iz}__"]
                        zNv[iz] += frac * (f_ - prev) * n_from_shells(cv, s_, scint=scint)
                        prev = f_
            N_mw = zN.sum(axis=0); N_v = zNv.sum(axis=0)
        else:
            raise ValueError(geom)
        nfac = n_scale.get(cid, n_scale.get("group_" + c["group"], 1.0))
        out[cid] = dict(N_mw=N_mw * nfac, N_mw_v=(N_mw if N_v is None else N_v) * nfac, N_x=N_x * nfac,
                        bands=parse_bands(c["bands_mhz"]), group=c["group"], tier=c["tier"], label=c["label"],
                        geom=geom, drift_loss=dfac)
        if geom in ("sky", "skytiers"):
            out[cid]["zones"] = zN * nfac
            out[cid]["zones_v"] = zNv * nfac
    meta = dict(n_mw_total=P["n_mw_total"], n_local=G.n_local(), diag=diag, ata=ata, ws=ws)
    return out, meta


def combine(comps, include=None, n_cap=None, vel=False):
    """Combine components over frequency.

    Returns dict with elementary frequency intervals and, for each interval and L,
      fid  : sum over overlap-groups of the max within the group (fiducial; avoids double counting
             surveys that share targets / sky area; for the sky surveys the max is taken separately in
             each declination zone, so complementary northern and southern surveys add),
      lo   : max over all components (strict lower bound on the union),
      hi   : plain sum (upper bound on the union; capped at n_cap = number of stars in the Galaxy),
    separately for Milky Way stars (mw) and for all stars including external galaxies (all).
    vel=True uses the Doppler-window-weighted counts of the sky surveys (for the 'HI_rest' prior)."""
    kN, kZ = ("N_mw_v", "zones_v") if vel else ("N_mw", "zones")
    ids = [k for k in comps if include is None or include(k, comps[k])]
    edges = sorted({e for k in ids for b in comps[k]["bands"] for e in b})
    iv = [(edges[i], edges[i + 1]) for i in range(len(edges) - 1)]
    nL = len(LGRID)
    res = {k: np.zeros((len(iv), nL)) for k in ("fid_mw", "lo_mw", "hi_mw", "fid_all", "lo_all", "hi_all")}
    for i, (a, b) in enumerate(iv):
        mid = 0.5 * (a + b)
        cov = [k for k in ids if any(lo <= mid <= hi for lo, hi in comps[k]["bands"])]
        if not cov:
            continue
        for tag in ("mw", "all"):
            arrs = {k: (comps[k][kN] if tag == "mw" else comps[k][kN] + comps[k]["N_x"]) for k in cov}
            stack = np.array(list(arrs.values()))
            res["lo_" + tag][i] = stack.max(axis=0)
            res["hi_" + tag][i] = stack.sum(axis=0)
            groups = {}
            for k in cov:
                groups.setdefault(comps[k]["group"], []).append(k)
            tot = np.zeros(nL)
            for ks in groups.values():
                if all("zones" in comps[k] for k in ks):      # sky surveys: best survey per declination zone
                    tot += np.max(np.array([comps[k][kZ] for k in ks]), axis=0).sum(axis=0)
                else:
                    tot += np.max(np.array([arrs[k] for k in ks]), axis=0)
            res["fid_" + tag][i] = tot
    if n_cap is not None:        # no union of Milky Way stars can exceed the number of stars in the Galaxy
        for k in ("fid_mw", "lo_mw", "hi_mw"):
            res[k] = np.minimum(res[k], n_cap)
    res["intervals"] = iv
    return res


def neff(res, prior, which="fid_mw"):
    kind, f1, f2 = (PRIORS.get(prior) or AUX_PRIORS[prior])
    w = np.array([prior_weight(kind, f1, f2, a, b) for a, b in res["intervals"]])
    return w @ res[which]


def neff_all(comps, include=None, n_cap=None, which="fid_mw", priors=None):
    """N_eff(L) for every frequency prior; 'HI_rest' uses the Doppler-window-weighted sky-survey counts."""
    res = combine(comps, include=include, n_cap=n_cap)
    resv = combine(comps, include=include, n_cap=n_cap, vel=True)
    out = {p: neff(resv if p == "HI_rest" else res, p, which) for p in (priors or PRIORS)}
    return out, res


NO_SAH = lambda k, c: k != "setiathome"     # baseline record: SETI@home is not yet a null


def bycatch_pointing_table(L_ref=1e17):
    """Per-pointing expected number of in-beam stars detectable at EIRP L_ref (calibrated counts)."""
    P = prep()
    rows = []
    for cid in ("bl_gbtL_bycatch", "bl_gbtS_bycatch", "bl_pks10_bycatch"):
        B = P["bl"][cid]
        c = P["comps"].set_index("comp_id").loc[cid]
        f = P["cal"]["mason"]["Parkes" if cid.startswith("bl_pks") else "GBT"]
        dm = float(dmax_pc(L_ref, float(c["smin_wm2"])))
        n = n_bycatch_dm(B["Cs_pt"] * f, dm, B["fwhm"])[:, 0]
        # weighted median distance of the detectable in-beam stars of each pointing set
        a = (G.D_MID / dm) ** 2
        w = np.clip(np.log(1.0 / np.maximum(a, 1e-300)) / np.log(2.0), 0.0, 1.0) * (a < 1.0)
        wd = (B["Cs_pt"] * f * w).sum(axis=0)
        cum = np.cumsum(wd) / wd.sum()
        dmed = float(np.interp(0.5, cum, G.D_MID))
        for name, l, b, x in zip(B["names"], B["l"], B["b"], n):
            rows.append(dict(comp_id=cid, target=name, l_deg=l, b_deg=b, n_inbeam_detectable=x, median_d_pc_of_set=dmed))
    t = pd.DataFrame(rows)
    t["share_of_set"] = t["n_inbeam_detectable"] / t.groupby("comp_id")["n_inbeam_detectable"].transform("sum")
    return t.sort_values(["comp_id", "n_inbeam_detectable"], ascending=[True, False])


def validate_model(meta, comps):
    """Compare model in-beam counts with Mason+ 2025 (Besancon model; their Table 1), counted the way they count
    (Gaussian taper inside the FWHM, d < 25 kpc), before and after calibration; and with the star numbers quoted
    by Gajjar+ 2021 and Tremblay & Tingay 2020."""
    P = prep()
    rows = []
    d = meta["diag"]
    cd = P["cal_diag"]
    for tel in ("GBT", "Parkes"):
        c0, c2 = cd["mason"][tel], cd["mason_excl2"][tel]
        for i in range(len(c0["d"])):
            rows.append(dict(check=f"BL {tel} in-beam stars with EIRP_min <= 1e{14 + i} W (axis distance {c0['d'][i]:.0f} pc)",
                             this_model=c0["raw"][i], reference=c0["NB"][i], reference_name="Mason+ 2025 Besancon (N_B)",
                             ratio=c0["raw"][i] / c0["NB"][i], calibrated_model=c0["fin"][i],
                             shell_factor=c0["fk"][i], shell_factor_excl2=c2["fk"][i], gaia_count_ref=MASON[tel]["NG"][i]))
    for cid, name, ref in (("bl_gc_gbtC", "Gajjar+ 2021 estimate, 4'x4' field to 8.5 kpc", 6e5),
                           ("bl_gc_parkes", "Gajjar+ 2021 estimate, two 30'x30' fields", 6e7)):
        rows.append(dict(check=f"{cid}: stars in field out to 8.5 kpc", this_model=d[cid]["stars_to_8p5kpc"],
                         reference=ref, reference_name=name, ratio=d[cid]["stars_to_8p5kpc"] / ref))
    rows.append(dict(check="mwa_vela: stars in 400 deg2 field (all distances)", this_model=d["mwa_vela"]["stars_all"],
                     reference=10355066, reference_name="Tremblay & Tingay 2020 Gaia DR2 stars with distances (magnitude-limited)",
                     ratio=d["mwa_vela"]["stars_all"] / 10355066, gaia_count_ref=10355066))
    tot = G.total_stars()
    rows.append(dict(check="Total stars in model Galaxy (R,z integral)", this_model=tot["total"], reference=1e11,
                     reference_name="canonical 1-4e11", ratio=tot["total"] / 1e11))
    rows.append(dict(check="Total stars in model Galaxy (sky-grid line-of-sight integral)",
                     this_model=meta["n_mw_total"], reference=tot["total"], reference_name="R,z integral",
                     ratio=meta["n_mw_total"] / tot["total"]))
    rows.append(dict(check="Stellar mass of the model Galaxy at the local mass-per-object ratio (Msun)",
                     this_model=meta["n_mw_total"] / P["stars_per_msun"], reference=5.5e10,
                     reference_name="5-6e10 Msun (e.g. McMillan 2017: 5.4e10)", ratio=meta["n_mw_total"] / P["stars_per_msun"] / 5.5e10))
    for w, v in sorted(P["vwin"].items()):
        fr = P["shells_v"][(v, "__allsky__")].sum() / meta["n_mw_total"]
        rows.append(dict(check=f"Fraction of Milky Way stars with LSR velocity inside the {w} window ({v[0]:+.0f}..{v[1]:+.0f} km/s)",
                         this_model=fr, reference=0.33 if abs(v[1]) < 60 else np.nan,
                         reference_name="referee Monte Carlo (+-42 km/s): 0.33" if abs(v[1]) < 60 else "",
                         ratio=fr / 0.33 if abs(v[1]) < 60 else np.nan))
    return pd.DataFrame(rows), tot


def main():
    P = prep()
    comps, meta = build()
    val, tot = validate_model(meta, comps)
    val.to_csv(os.path.join(DATA, "galaxy_model_validation.csv"), index=False, float_format="%.5g")
    pd.set_option("display.width", 250); pd.set_option("display.max_colwidth", 80)
    print(val.to_string(index=False, float_format=lambda x: f"{x:.4g}"))
    print("model components:", {k: f"{v:.3e}" for k, v in tot.items()}, " local n =", round(meta["n_local"], 4))
    print("pointings left out in the 'mason_excl2' calibration:", P["cal_diag"]["dropped_pointings"])
    # calibration table
    crow = []
    for mode in ("mason", "mason_excl2"):
        for tel in ("GBT", "Parkes"):
            c = P["cal_diag"][mode][tel]
            for i in range(len(c["d"])):
                crow.append(dict(calibration=mode, telescope=tel, EIRPmin_W=10.0 ** (14 + i), d_axis_pc=c["d"][i],
                                 N_Besancon_Mason2025=c["NB"][i], N_model_raw=c["raw"][i], N_model_calibrated=c["fin"][i],
                                 shell_factor=c["fk"][i]))
    pd.DataFrame(crow).to_csv(os.path.join(DATA, "bycatch_calibration.csv"), index=False, float_format="%.5g")
    bp = bycatch_pointing_table()
    bp.groupby("comp_id").head(15).to_csv(os.path.join(DATA, "bycatch_pointings.csv"), index=False, float_format="%.5g")
    # drift factors per component
    cc = P["comps"]
    dr = pd.DataFrame(dict(comp_id=cc["comp_id"], label=cc["label"], chan_hz=cc["chan_hz"], tres_s=cc["tres_s"],
                           threshold_factor_orbit=[drift_loss(c, "orbit") for _, c in cc.iterrows()],
                           threshold_factor_surface=[drift_loss(c, "surface") for _, c in cc.iterrows()],
                           source=cc["drift_source"]))
    dr.to_csv(os.path.join(DATA, "drift_factors.csv"), index=False, float_format="%.4g")
    # per-component table
    tab = pd.DataFrame({"EIRP_W": LGRID})
    for k, c in comps.items():
        tab[k] = c["N_mw"] + c["N_x"]
    tab.to_csv(os.path.join(DATA, "ncov_by_component.csv"), index=False, float_format="%.5g")
    # derived target tables
    ata = meta["ata"]
    ata[["id", "S", "Name", "BW", "angDist", "Plx", "e_Plx", "Gmag", "d_pc", "fweight"]].to_csv(
        os.path.join(DATA, "ata_targets_with_distances.csv"), index=False, float_format="%.6g")
    ws = meta["ws"]
    pd.concat([bl_target_rows(ws, inst)[0] for inst in ("GBT L-band", "GBT S-band", "Parkes 10-cm")])[
        ["Target", "RATdeg", "DETdeg", "rest", "InstS", "min-detec-EIRP"]].to_csv(
        os.path.join(DATA, "bl_targets.csv"), index=False, float_format="%.6g")
    print("ATA:", meta["diag"]["ata_harp2016"], " FAST GCs:", meta["diag"]["fast_gcs"])
    print("BL target rows / unique targets:", {k: (v["n_rows_flagged"], v["n_targets"]) for k, v in meta["diag"].items() if "n_targets" in v})
    rows = pd.DataFrame({"EIRP_W": LGRID})
    res = None
    for tag, inc in (("", NO_SAH), ("_withSAH", None)):
        for w in ("fid_mw", "lo_mw", "hi_mw", "fid_all"):
            ne, r = neff_all(comps, include=inc, n_cap=meta["n_mw_total"], which=w)
            for p in PRIORS:
                rows[f"Neff_{p}_{w}{tag}"] = ne[p]
            if tag == "" and w == "fid_mw":
                res = r
    rows["Nmax_any_freq_mw"] = res["fid_mw"].max(axis=0)
    rows["Nmax_any_freq_all"] = res["fid_all"].max(axis=0)
    rows["N_MW_model"] = meta["n_mw_total"]
    rows.to_csv(os.path.join(DATA, "neff_table.csv"), index=False, float_format="%.5g")
    iv = pd.DataFrame(res["intervals"], columns=["f_lo_MHz", "f_hi_MHz"])
    for j, L in enumerate(LGRID):
        if j % 8 == 0:
            iv[f"N_mw_L{np.log10(L):.0f}"] = res["fid_mw"][:, j]
    iv.to_csv(os.path.join(DATA, "coverage_intervals.csv"), index=False, float_format="%.5g")
    sel = [np.argmin(abs(LGRID - x)) for x in (1e10, 1e12, 1e13, 1e15, 1e17, 1e20, 1e22, 1e26)]
    print(rows.iloc[sel][["EIRP_W", "Neff_HI_line_fid_mw", "Neff_HI_rest_fid_mw", "Neff_HI_line_fid_mw_withSAH", "Neff_water_hole_fid_mw",
                          "Neff_log_1_10GHz_fid_mw", "Neff_log_1_10GHz_lo_mw", "Neff_log_1_10GHz_hi_mw",
                          "Neff_log_0p1_100GHz_fid_mw"]].to_string(index=False, float_format=lambda x: f"{x:.3g}"))
    return comps, meta, res


if __name__ == "__main__":
    main()