#!/usr/bin/env python3
"""Recount of the TESS superflare study on sciencejournal.ai from its own tables: the numbers of ours on
Fact-Checked Science's page https://factchecked.science/tess-superflare-rate/ .

Our analysis. Python 3, standard library only. It reads the study's tables as data and never runs the study's
code. Run it with `python3 -I`:

    python3 -I recount_from_tables.py <folder> [output.json]

<folder> (searched with its subfolders) must hold these five files of the study's bundle,
sha256:e214a7787ce60a51854fac1ab0385f59f930e0d3d56b0350a4cb0f733c916f75, each served by its SHA-256 at
https://sciencejournal.ai/api/v1/files/sha256:<hash> (save each under the name given):

    lcs.csv.gz        data/run/lcs.csv.gz       15.3 MB  3d02037de085014bec40c708dd2ed481706d774cf0427959976f28245e6f0d13
    cands.csv.gz      data/run/cands.csv.gz     14.8 MB  354946a22be3e8ff66ee25f97f9970041b050dd69d0557304ff933b4d65bc4ea
    inj.csv.gz        data/run/inj.csv.gz       32.5 MB  7ef8bb85eb0bc4cbc367e05cb26ab9422febbcd675c584825ca62575a679c7da
    sample_stars.csv  data/sample_stars.csv     12.9 MB  296553f8c82fba49891b9300d676c1e4f5498b36c1d20ac5c82bf5576f413b4c
    localization.csv  results/localization.csv  18 KB    0bc6f7fd1da8af9126050db031fc432fb3cbcc6c132c39e1bd0f6ded3600e9da

The study's data are CC BY-SA 4.0 (sciencejournal.ai reference agent, 2026); sample_stars.csv carries Gaia DR3
values (credit ESA/Gaia/DPAC, CC BY-SA 3.0 IGO). This script's output is derived from them and is offered under
CC BY-SA 4.0. The script checks every file's SHA-256 and stops if one differs. It prints a summary and writes
recount-from-tables.json (into the folder, or to the path given). A run takes about a minute. The resampling of
stars in step 6 uses Python's random module with the fixed seed 20261031, so it gives the same numbers each run.

What it does, following the rules the study's Methods state
  1. Stars set aside: TESS fast rotators (any sector with a Lomb-Scargle peak at 0.2-12 days, power above 0.3
     and semi-amplitude above 0.1%) and eclipsers (any sector with two or more dips of ED 20 s or more).
  2. Star-years: the years of good data of the light curves left.
  3. Counted superflares: vetted flares of 1e34 erg or more on the stars left, with at most one other star
     having a candidate of either sign in the same sector and camera within three cadences of the peak; then
     without the one event the study's SAP sign check removed (TIC 469390625, sector 16). That check needs the
     light curves, so the event is taken from the study's record (its deviations), not tested here.
  4. The count Vasilyev et al. (2024)'s rate predicts: the study's Eq. (2), sum over energy bins of the rate's
     flares per star-year in the 0.2-dex bin centred on each injected energy (the top bin taking the tail),
     times the years of good data times the fraction of injections there recovered at a measured 1e34 erg or
     more. The same for other Kepler rates and slopes, and for the study's shorter and longer flares.
  5. Exact 95% Poisson intervals, tail probabilities, the spread of the flares over stars (the study's C6),
     the stars that hold most of them, their steady short periods, offset tests, crowding and Gaia values, the
     split by Gaia's Ca II activity index (C5), and the planned subsets of the study's Table 1: the Sun-like
     stars (C3; 5500 K < Teff < 6000 K, as Vasilyev et al.'s Table 1 writes it, and median TESS variability
     below 0.3%), the brightness thirds (by the target lists' TESS magnitude, as the study's Methods use it),
     the flares that pass the offset test, the stars with no sign of a companion, and the higher thresholds.
  6. An interval that allows for the flares' clustering on a few stars: the stars resampled with replacement,
     2,000 times, and the ratio of counted to predicted flares in each resample.
"""

import collections
import csv
import gzip
import hashlib
import json
import math
import os
import random
import statistics
import sys

YEAR = 365.25
CADENCE_DAYS = 120.0 / 86400.0
FILES = {
    "lcs.csv.gz": "3d02037de085014bec40c708dd2ed481706d774cf0427959976f28245e6f0d13",
    "cands.csv.gz": "354946a22be3e8ff66ee25f97f9970041b050dd69d0557304ff933b4d65bc4ea",
    "inj.csv.gz": "7ef8bb85eb0bc4cbc367e05cb26ab9422febbcd675c584825ca62575a679c7da",
    "sample_stars.csv": "296553f8c82fba49891b9300d676c1e4f5498b36c1d20ac5c82bf5576f413b4c",
    "localization.csv": "0bc6f7fd1da8af9126050db031fc432fb3cbcc6c132c39e1bd0f6ded3600e9da",
}
SEED, RESAMPLES = 20261031, 2000
URL = "https://sciencejournal.ai/api/v1/files/sha256:"
SIGN_CHECK_REMOVED = (469390625, 16)   # the study's deviations: the one counted flare with a falling SAP flux
F34_TABLE1 = 8.63e-3 + 3.90e-4          # Vasilyev et al. (2024), Table 1, full sample, above 1e34 erg
F35_TABLE1 = 3.90e-4


# ---------- statistics, standard library only ----------

def poisson_logpmf(k, mu):
    return -mu + k * math.log(mu) - math.lgamma(k + 1) if mu > 0 else (0.0 if k == 0 else -math.inf)


def poisson_cdf(k, mu):
    return sum(math.exp(poisson_logpmf(i, mu)) for i in range(k + 1))


def poisson_log10_sf(n, mu):
    """log10 P(X >= n) for X ~ Poisson(mu), summed in logs."""
    terms = [poisson_logpmf(k, mu) for k in range(n, n + 2000)]
    m = max(terms)
    return (m + math.log(sum(math.exp(t - m) for t in terms))) / math.log(10)


def bisect(f, lo, hi, tol=1e-10):
    for _ in range(200):
        mid = 0.5 * (lo + hi)
        if f(mid) > 0:
            hi = mid
        else:
            lo = mid
        if hi - lo < tol:
            break
    return 0.5 * (lo + hi)


def poisson_ci(n, cl=0.95):
    """Exact (Garwood) interval: the means for which n is just in the upper or lower 2.5% tail."""
    a = (1 - cl) / 2
    lo = 0.0 if n == 0 else bisect(lambda m: (1 - poisson_cdf(n - 1, m)) - a, 0.0, n + 50.0)
    hi = bisect(lambda m: a - poisson_cdf(n, m), 0.0, n * 3 + 50.0)
    return lo, hi


def binom_logpmf(k, n, p):
    return math.lgamma(n + 1) - math.lgamma(k + 1) - math.lgamma(n - k + 1) + k * math.log(p) + (n - k) * math.log1p(-p)


def binom_test_two_sided(k, n, p):
    """As scipy.stats.binomtest: the sum of the probabilities no larger than that of the observed count."""
    obs = binom_logpmf(k, n, p)
    return min(1.0, sum(math.exp(binom_logpmf(i, n, p)) for i in range(n + 1) if binom_logpmf(i, n, p) <= obs + 1e-7))


def fisher_two_sided(a, b, c, d):
    """Fisher's exact test of [[a, b], [c, d]], two-sided, as scipy.stats.fisher_exact."""
    r1, c1, n = a + b, a + c, a + b + c + d

    def logp(x):
        return (math.lgamma(r1 + 1) - math.lgamma(x + 1) - math.lgamma(r1 - x + 1)
                + math.lgamma(n - r1 + 1) - math.lgamma(c1 - x + 1) - math.lgamma(n - r1 - c1 + x + 1)
                - math.lgamma(n + 1) + math.lgamma(c1 + 1) + math.lgamma(n - c1 + 1))
    obs = logp(a)
    lo, hi = max(0, c1 - (n - r1)), min(r1, c1)
    return min(1.0, sum(math.exp(logp(x)) for x in range(lo, hi + 1) if logp(x) <= obs + 1e-7))


def ratio_row(n, N, F34=F34_TABLE1):
    lo, hi = poisson_ci(n)
    return {"counted": n, "predicted": round(N, 3), "ratio": round(n / N, 3),
            "ratio_ci95": [round(lo / N, 3), round(hi / N, 3)],
            "rate_per_star_year": round(n / N * F34, 5), "years_per_superflare": round(N / (n * F34), 1),
            "years_ci95": [round(N / (hi * F34), 1), round(N / (lo * F34), 1)] if n else None}


# ---------- reading ----------

def find(folder, name):
    for root, _dirs, files in os.walk(folder):
        if name in files:
            return os.path.join(root, name)
    sys.exit(f"{name} not found under {folder}: download it from {URL}{FILES[name]} and save it as {name}")


def checked(path, want):
    h = hashlib.sha256()
    with open(path, "rb") as fh:
        for chunk in iter(lambda: fh.read(1 << 20), b""):
            h.update(chunk)
    if h.hexdigest() != want:
        sys.exit(f"{path}: SHA-256 {h.hexdigest()} is not the study's {want}")
    return path


def rows(path):
    opener = gzip.open if path.endswith(".gz") else open
    with opener(path, "rt", encoding="utf-8", newline="") as fh:
        yield from csv.DictReader(fh)


def fnum(x):
    try:
        return float(x)
    except (TypeError, ValueError):
        return math.nan


# ---------- the recount ----------

def cum_rate(F34, beta, loge):
    return F34 * 10 ** (-beta * (loge - 34.0))


def predicted(expo, F34, beta):
    """The study's Eq. (2): bins of 0.2 dex centred on the injected energies, the top bin taking the tail."""
    es = sorted(expo)
    total = 0.0
    for k, e in enumerate(es):
        lo, hi = e - 0.1, e + 0.1
        dF = cum_rate(F34, beta, lo) if k == len(es) - 1 else cum_rate(F34, beta, lo) - cum_rate(F34, beta, hi)
        total += dF * expo[e]
    return total


def main():
    if len(sys.argv) < 2:
        sys.exit(__doc__)
    folder = sys.argv[1]
    out_path = sys.argv[2] if len(sys.argv) > 2 else os.path.join(folder, "recount-from-tables.json")
    paths = {name: checked(find(folder, name), sha) for name, sha in FILES.items()}
    out = {"inputs": {name: {"sha256": sha, "url": URL + sha} for name, sha in FILES.items()},
           "licence": "Derived from the study's data (CC BY-SA 4.0) and Gaia DR3 values (ESA/Gaia/DPAC, CC BY-SA 3.0 IGO); offered under CC BY-SA 4.0"}

    # light curves
    lcs = []
    for r in rows(paths["lcs.csv.gz"]):
        lcs.append({"file": r["file"], "tic": int(r["tic"]), "sector": int(r["sector"]), "days": fnum(r["days_good"]),
                    "rvar": fnum(r["rvar"]), "p": fnum(r["p_rot"]), "pow": fnum(r["ls_power"]), "amp": fnum(r["ls_amp"]),
                    "crowd": fnum(r["crowdsap"])})
    stars_all = {l["tic"] for l in lcs}
    fast = {l["tic"] for l in lcs if 0.2 <= l["p"] <= 12 and l["pow"] > 0.3 and l["amp"] > 0.001}

    # candidates
    cands = []
    for r in rows(paths["cands.csv.gz"]):
        cands.append({"file": r["file"], "tic": int(r["tic"]), "sector": int(r["sector"]), "camera": int(r["camera"]), "kind": r["kind"],
                      "tpk": fnum(r["tpk"]), "ed": fnum(r["ed"]), "energy": fnum(r["energy"]),
                      "vetted": r["vetted"] == "True", "fwhm": fnum(r["fwhm"])})
    dips_by = collections.Counter((c["tic"], c["sector"]) for c in cands if c["kind"] == "dip" and c["ed"] >= 20)
    eclipsers = {t for (t, _s), n in dips_by.items() if n >= 2}
    used_stars = stars_all - fast - eclipsers
    used = [l for l in lcs if l["tic"] in used_stars]
    out["sample"] = {"confirmatory_stars": len(stars_all), "light_curves": len(lcs), "fast_rotators": len(fast),
                     "eclipsers": len(eclipsers), "eclipsers_not_fast": len(eclipsers - fast),
                     "stars_used": len(used_stars), "light_curves_used": len(used),
                     "star_years": round(sum(l["days"] for l in used) / YEAR, 2)}

    # the coincidence rule: candidates of either sign by sector and camera
    by_cam = collections.defaultdict(list)
    for c in cands:
        by_cam[(c["sector"], c["camera"])].append((c["tpk"], c["tic"]))
    for v in by_cam.values():
        v.sort()

    def others(c):
        import bisect as bs
        v = by_cam[(c["sector"], c["camera"])]
        i = bs.bisect_left(v, (c["tpk"] - 3 * CADENCE_DAYS - 1e-9, -1))
        seen = set()
        while i < len(v) and v[i][0] <= c["tpk"] + 3 * CADENCE_DAYS + 1e-9:
            if v[i][1] != c["tic"]:
                seen.add(v[i][1])
            i += 1
        return len(seen)

    def counted(kind, stars):
        return [c for c in cands if c["kind"] == kind and c["vetted"] and c["energy"] >= 1e34
                and c["tic"] in stars and others(c) <= 1]

    flares_reg = counted("flare", used_stars)
    flares = [c for c in flares_reg if (c["tic"], c["sector"]) != SIGN_CHECK_REMOVED]
    not_fast = stars_all - fast
    out["counts"] = {"counted_registered": len(flares_reg), "counted_after_sign_check": len(flares),
                     "dips_on_stars_used": len(counted("dip", used_stars)),
                     "dips_before_eclipser_rule": len(counted("dip", not_fast)),
                     "flares_before_eclipser_rule": len([c for c in counted("flare", not_fast)
                                                        if (c["tic"], c["sector"]) != SIGN_CHECK_REMOVED]),
                     "median_fwhm_minutes": round(statistics.median(c["fwhm"] for c in flares) * 1440, 2),
                     "energy_quartiles_erg": [float(f"{q:.4g}") for q in statistics.quantiles([c["energy"] for c in flares], n=4)],
                     "energies": {"1-2e34": sum(1e34 <= c["energy"] < 2e34 for c in flares),
                                  "2-5e34": sum(2e34 <= c["energy"] < 5e34 for c in flares),
                                  "5e34-1e35": sum(5e34 <= c["energy"] < 1e35 for c in flares),
                                  ">=1e35": sum(c["energy"] >= 1e35 for c in flares)},
                     "largest_erg": float(f"{max(c['energy'] for c in flares):.4g}")}

    # injections: exposure (years of good data times the fraction recovered at >= 1e34 erg), per model
    days = {l["file"]: l["days"] for l in used}
    tic_of = {l["file"]: l["tic"] for l in used}
    expo = {m: collections.defaultdict(float) for m in ("base", "short", "long")}
    rec_any = {m: collections.defaultdict(float) for m in ("base", "short", "long")}
    tw = {m: collections.defaultdict(float) for m in ("base", "short", "long")}
    files_by_model = {m: set() for m in ("base", "short", "long")}
    per_star_rows = collections.defaultdict(lambda: collections.defaultdict(float))   # tic -> log_e -> exposure (base)
    expo_hi = {"345": collections.defaultdict(float), "350": collections.defaultdict(float)}   # base, >= 1e34.5 / 1e35 erg
    sub_rows = []   # (file, model, log_e, exposure) for the light curves injected with all three models
    for r in rows(paths["inj.csv.gz"]):
        f = r["file"]
        if f not in days:
            continue
        m, le = r["model"], round(float(r["log_e"]), 1)
        n_inj = float(r["n_inj"])
        if n_inj == 0:   # 106 of the 2,379,692 rows hold no injection; they add nothing, as in the study's sums
            continue
        t = days[f] / YEAR
        e = t * float(r["n_rec_340"]) / n_inj
        expo[m][le] += e
        rec_any[m][le] += t * float(r["n_rec_any"]) / n_inj
        tw[m][le] += t
        files_by_model[m].add(f)
        if m == "base":
            per_star_rows[tic_of[f]][le] += e
            for k in expo_hi:
                expo_hi[k][le] += t * float(r[f"n_rec_{k}"]) / n_inj
        sub_rows.append((f, m, le, e))
    out["completeness"] = {m: {f"{le:.1f}": {"recovered_at_ge_1e34": round(expo[m][le] / tw[m][le], 4),
                                             "recovered_at_any_energy": round(rec_any[m][le] / tw[m][le], 4)}
                               for le in sorted(tw[m])} for m in expo}

    beta_t1 = math.log10(F34_TABLE1 / F35_TABLE1)
    N = predicted(expo["base"], F34_TABLE1, beta_t1)
    n = len(flares)
    out["primary"] = ratio_row(n, N)
    out["primary"]["log10_p_at_least_if_kepler"] = round(poisson_log10_sf(n, N), 1)
    out["earlier_rates"] = {}
    for name, F in (("1 per 2,500 years (the study's earlier rate)", 1 / 2500), ("1 per 6,000 years (Okamoto et al. 2021)", 1 / 6000)):
        Ne = predicted(expo["base"], F, beta_t1)
        out["earlier_rates"][name] = {"predicted": round(Ne, 3), "ratio": round(n / Ne, 1),
                                      "log10_p_at_least_observed": round(poisson_log10_sf(n, Ne), 1)}

    # the study's duration rows: the prediction scaled by the completeness in the light curves with all three
    both = files_by_model["short"] & files_by_model["long"] & files_by_model["base"]
    sub = {m: collections.defaultdict(float) for m in ("base", "short", "long")}
    for f, m, le, e in sub_rows:
        if f in both:
            sub[m][le] += e
    Nsub = {m: predicted(sub[m], F34_TABLE1, beta_t1) for m in sub}
    out["durations"] = {"light_curves_with_all_three": len(both)}
    for m in ("short", "long"):
        Nm = N * Nsub[m] / Nsub["base"]
        out["durations"][m] = ratio_row(n, Nm)

    # other Kepler normalisations and slopes
    variants = {
        "Table 1, as the study (F34 9.02e-3, beta 1.364)": (F34_TABLE1, beta_t1),
        "the text's 8.63e-3 above 1e34 erg, same slope": (8.63e-3, beta_t1),
        "their corrected Fig. 3, the study's reading (0.0121 + 0.000288; 0.000288 above 1e35)": (0.0121 + 0.000288, math.log10((0.0121 + 0.000288) / 0.000288)),
        "their corrected Fig. 3, our reading (0.012666; 0.000298 above 1e35)": (0.012666, math.log10(0.012666 / 0.0002982)),
        "Table 1 normalisation, cumulative slope 0.97": (F34_TABLE1, 0.97),
        "Table 1 normalisation, cumulative slope 1.97": (F34_TABLE1, 1.97),
        "the Sun-like row (9.31e-3 + 3.59e-4) applied to all stars": (9.31e-3 + 3.59e-4, math.log10((9.31e-3 + 3.59e-4) / 3.59e-4)),
    }
    out["kepler_variants"] = {}
    for name, (F, b) in variants.items():
        row = ratio_row(n, predicted(expo["base"], F, b), F)
        row.update({"F34": round(F, 6), "beta": round(b, 4)})
        out["kepler_variants"][name] = row

    # the spread over stars (C6)
    per_star = collections.Counter(c["tic"] for c in flares)
    mu_k = {t: predicted(per_star_rows[t], F34_TABLE1, beta_t1) for t in used_stars}
    tot = sum(mu_k.values())
    dist = [1.0] + [0.0] * 13
    for t in used_stars:
        mu = mu_k[t] / tot * n
        p2 = 1 - math.exp(-mu) * (1 + mu)
        new = [x * (1 - p2) for x in dist]
        for i in range(1, len(dist)):
            new[i] += dist[i - 1] * p2
        new[-1] += dist[-1] * p2
        dist = new
    multi = {t: k for t, k in per_star.items() if k >= 2}
    out["spread"] = {"stars_with_flares": len(per_star), "stars_with_two_or_more": len(multi),
                     "flares_on_them": sum(multi.values()), "most_on_one_star": max(per_star.values()),
                     "expected_stars_with_two_or_more_at_one_shared_rate": round(sum(1 - math.exp(-mu_k[t] / tot * n) * (1 + mu_k[t] / tot * n) for t in used_stars), 3),
                     "log10_p_12_or_more": round(math.log10(sum(dist[12:])), 1),
                     "top_stars": [{"tic": t, "flares": k, "expected_at_kepler_rate": round(mu_k[t], 4),
                                    "expected_at_one_shared_rate": round(mu_k[t] / tot * n, 4)} for t, k in per_star.most_common(4)]}

    # the two stars that hold 18 flares, and steady short periods across sectors
    sectors = collections.defaultdict(list)
    for l in used:
        sectors[l["tic"]].append(l)
    steady = set()
    multi_sector = 0
    for t, ls in sectors.items():
        if len(ls) < 3:
            continue
        multi_sector += 1
        med = statistics.median(l["p"] for l in ls)
        frac = sum(abs(l["p"] / med - 1) <= 0.03 for l in ls) / len(ls)
        if frac >= 0.5 and med < 11:
            steady.add(t)
    stars_meta = {}
    for r in rows(paths["sample_stars.csv"]):
        t = int(r["tic"])
        if t in used_stars:
            stars_meta[t] = {"teff": fnum(r["Teff"]), "rad": fnum(r["Rad"]), "tmag": fnum(r["Tmag"]), "tmag_list": fnum(r["tmag_list"]),
                             "act": fnum(r["activityindex_espcs"]), "vbroad": fnum(r["vbroad"]), "ruwe": fnum(r["ruwe"]),
                             "plx": fnum(r["parallax"]), "nss": fnum(r["non_single_star"]), "rvp": fnum(r["rv_chisq_pvalue"]),
                             "rva": fnum(r["rv_amplitude_robust"])}
    loc = {(r["file"], round(float(r["tpk"]), 5)): {"offset_px": float(r["offset_px"]), "offset_sigma": float(r["offset_sigma"]),
                                                      "on_target": r["on_target"] == "True"} for r in rows(paths["localization.csv"])}
    loc_counted = [loc[(c["file"], round(c["tpk"], 5))] for c in flares]
    top2 = [t for t, _k in per_star.most_common(2)]
    out["top_two"] = []
    for t in top2:
        ls = sectors[t]
        med = statistics.median(l["p"] for l in ls)
        near = [l for l in ls if abs(l["p"] / med - 1) <= 0.03]
        half = [l for l in ls if abs(l["p"] / (med / 2) - 1) <= 0.03]
        s = stars_meta[t]
        v_eq = 2 * math.pi * s["rad"] * 695700.0 / (med * 86400.0)
        out["top_two"].append({
            "tic": t, "flares": per_star[t], "sectors": len(ls), "sectors_with_flares": len({c["sector"] for c in flares if c["tic"] == t}),
            "years_of_data": round(sum(l["days"] for l in ls) / YEAR, 2),
            "median_peak_period_days": round(med, 3), "sectors_within_3pc": len(near),
            "period_range_within_3pc": [round(min(l["p"] for l in near), 3), round(max(l["p"] for l in near), 3)],
            "sectors_at_half_the_period": len(half),
            "power_median": round(statistics.median(l["pow"] for l in ls), 3), "power_max": round(max(l["pow"] for l in ls), 3),
            "semi_amplitude_percent": [round(100 * min(l["amp"] for l in ls), 4), round(100 * max(l["amp"] for l in ls), 4)],
            "gaia_activity_index_nm": s["act"], "gaia_vbroad_km_s": round(s["vbroad"], 2), "gaia_ruwe": round(s["ruwe"], 3),
            "tic_teff_K": s["teff"], "tic_radius_rsun": s["rad"], "tmag": s["tmag"], "distance_pc_from_parallax": round(1000 / s["plx"], 1),
            "equatorial_speed_at_that_period_km_s": round(v_eq, 1),
            "sin_i_at_most_vbroad_over_speed": round(s["vbroad"] / v_eq, 3),
            "i_at_most_deg": round(math.degrees(math.asin(min(1.0, s["vbroad"] / v_eq))), 1),
            "chance_of_i_that_low_for_a_random_orientation": round(1 - math.cos(math.asin(min(1.0, s["vbroad"] / v_eq))), 3),
            "crowdsap_range": [round(min(l["crowd"] for l in ls), 3), round(max(l["crowd"] for l in ls), 3)],
            "gaia_non_single_star": s["nss"], "gaia_rv_chisq_pvalue": round(s["rvp"], 3), "gaia_rv_amplitude_robust_km_s": round(s["rva"], 2),
            "offset_test": [{"sector": c["sector"], "offset_px": round(loc[(c["file"], round(c["tpk"], 5))]["offset_px"], 2),
                             "offset_sigma": round(loc[(c["file"], round(c["tpk"], 5))]["offset_sigma"], 1),
                             "passes": loc[(c["file"], round(c["tpk"], 5))]["on_target"]} for c in flares if c["tic"] == t]})
    out["offset_test_all_counted"] = {"flares": len(loc_counted), "pass": sum(x["on_target"] for x in loc_counted),
                                      "median_offset_px": round(statistics.median(x["offset_px"] for x in loc_counted), 2)}
    flaring_multi = {t for t in per_star if len(sectors[t]) >= 3}
    a = len(flaring_multi & steady)
    b = len(flaring_multi - steady)
    quiet_multi = {t for t, ls in sectors.items() if len(ls) >= 3} - flaring_multi
    c = len(quiet_multi & steady)
    d = len(quiet_multi - steady)
    out["steady_periods"] = {"stars_with_3_or_more_sectors": multi_sector, "steady": len(steady),
                             "flaring_stars_with_3_or_more_sectors": len(flaring_multi), "flaring_and_steady": a,
                             "flares_on_steady_stars": sum(per_star[t] for t in steady if t in per_star),
                             "fisher_table": [[a, b], [c, d]], "odds_ratio": round((a * d) / (b * c), 2),
                             "fisher_p_two_sided": round(fisher_two_sided(a, b, c, d), 4)}

    def without(drop):
        keep = used_stars - set(drop)
        e = collections.defaultdict(float)
        for t in keep:
            for le, v in per_star_rows[t].items():
                e[le] += v
        k = sum(1 for c in flares if c["tic"] in keep)
        return ratio_row(k, predicted(e, F34_TABLE1, beta_t1))
    out["without_the_two_stars"] = without(top2)
    out["without_steady_period_stars"] = without(steady)

    # C5: split at the median Gaia DR3 Ca II activity index of the stars used
    acts = [s["act"] for s in stars_meta.values() if not math.isnan(s["act"])]
    med_act = statistics.median(acts)

    def split(drop):
        res = {}
        for name, cond in (("below", lambda x: x < med_act), ("at_or_above", lambda x: x >= med_act)):
            ids = {t for t, s in stars_meta.items() if not math.isnan(s["act"]) and cond(s["act"]) and t not in drop}
            e = collections.defaultdict(float)
            for t in ids:
                for le, v in per_star_rows[t].items():
                    e[le] += v
            k = sum(1 for c in flares if c["tic"] in ids)
            res[name] = ratio_row(k, predicted(e, F34_TABLE1, beta_t1))
        p0 = res["below"]["predicted"] / (res["below"]["predicted"] + res["at_or_above"]["predicted"])
        res["binomial_p_two_sided"] = round(binom_test_two_sided(res["below"]["counted"], res["below"]["counted"] + res["at_or_above"]["counted"], p0), 4)
        return res
    out["activity_split"] = {"median_index_nm": med_act, "flares_on_stars_without_an_index":
                             sum(1 for c in flares if math.isnan(stars_meta[c["tic"]]["act"])),
                             "all": split(set()), "without_the_two_stars": split(set(top2))}

    # brightness thirds (C4) and the Sun-like subset (C3)
    def subset(ids, F34=F34_TABLE1, beta=beta_t1):
        e = collections.defaultdict(float)
        for t in ids:
            for le, v in per_star_rows[t].items():
                e[le] += v
        k = sum(1 for c in flares if c["tic"] in ids)
        row = ratio_row(k, predicted(e, F34, beta), F34)
        row["stars"] = len(ids)
        return row
    thirds = {}
    for name, lo, hi in (("T below 8.5", -99, 8.5), ("T 8.5-9.5", 8.5, 9.5), ("T 9.5-10.5", 9.5, 10.5)):
        thirds[name] = subset({t for t, s in stars_meta.items() if lo <= s["tmag_list"] < hi})
    obs = [r["counted"] for r in thirds.values()]
    pre = [r["predicted"] for r in thirds.values()]
    exp_ = [p / sum(pre) * sum(obs) for p in pre]
    chi = sum((o - x) ** 2 / x for o, x in zip(obs, exp_))
    thirds["chi_square_2_dof"] = {"chi2": round(chi, 2), "p": round(math.exp(-chi / 2), 3)}
    out["brightness_thirds"] = thirds
    rvar_med = {t: statistics.median(l["rvar"] for l in ls) for t, ls in sectors.items()}
    sunlike = {t for t, s in stars_meta.items() if 5500 < s["teff"] < 6000 and rvar_med.get(t, 1) < 0.003}
    F_sun = 9.31e-3 + 3.59e-4
    out["sunlike_subset"] = subset(sunlike, F_sun, math.log10(F_sun / 3.59e-4))

    # more planned subsets of the study's Table 1
    k_off = sum(1 for x in loc_counted if x["on_target"])
    out["offset_test_subset"] = ratio_row(k_off, N)
    single = {t for t, s in stars_meta.items() if (s["nss"] == 0 or math.isnan(s["nss"])) and not (s["rvp"] < 0.01)}
    out["no_sign_of_a_companion_subset"] = subset(single)
    for k, thr in (("345", 10 ** 34.5), ("350", 1e35)):
        kk = sum(1 for c in flares if c["energy"] >= thr)
        out[f"above_1e{k[:2]}.{k[2]}_subset"] = ratio_row(kk, predicted(expo_hi[k], F34_TABLE1, beta_t1))

    # an interval that allows for the clustering: resample the stars with replacement
    ids = sorted(used_stars)
    cnt = [per_star.get(t, 0) for t in ids]
    mus = [mu_k[t] for t in ids]
    rng = random.Random(SEED)
    rs = []
    for _ in range(RESAMPLES):
        pick = rng.choices(range(len(ids)), k=len(ids))
        rs.append(sum(map(cnt.__getitem__, pick)) / sum(map(mus.__getitem__, pick)))
    rs.sort()
    out["star_resampled_interval"] = {"seed": SEED, "resamples": RESAMPLES, "ratio": round(n / N, 3),
                                      "ratio_2.5_97.5_percent": [round(rs[int(0.025 * RESAMPLES)], 2), round(rs[int(0.975 * RESAMPLES) - 1], 2)],
                                      "sum_of_squared_counts_per_star": sum(k * k for k in per_star.values())}

    with open(out_path, "w", encoding="utf-8", newline="\n") as fh:
        json.dump(out, fh, indent=1)
        fh.write("\n")
    p = out["primary"]
    print(f"stars {out['sample']['stars_used']:,}, light curves {out['sample']['light_curves_used']:,}, "
          f"star-years {out['sample']['star_years']:,}")
    print(f"counted {p['counted']} (registered {out['counts']['counted_registered']}), Kepler rate predicts {p['predicted']}: "
          f"ratio {p['ratio']} ({p['ratio_ci95'][0]}-{p['ratio_ci95'][1]}), one per {p['years_per_superflare']} years")
    for k, v in out["kepler_variants"].items():
        print(f"  {k}: {v['predicted']} -> {v['ratio']} ({v['ratio_ci95'][0]}-{v['ratio_ci95'][1]})")
    for m in ("short", "long"):
        v = out["durations"][m]
        print(f"  flares {'half' if m == 'short' else '2.5 times'} as long: {v['predicted']} -> {v['ratio']} ({v['ratio_ci95'][0]}-{v['ratio_ci95'][1]})")
    print(f"  without the two stars: {out['without_the_two_stars']['ratio']}; without steady-period stars: "
          f"{out['without_steady_period_stars']['ratio']} ({out['steady_periods']['steady']} stars)")
    print(f"  stars resampled ({RESAMPLES}, seed {SEED}): {out['star_resampled_interval']['ratio_2.5_97.5_percent']}")
    print(f"  Sun-like {out['sunlike_subset']['ratio']} on {out['sunlike_subset']['stars']:,} stars; thirds "
          f"{[thirds[k]['ratio'] for k in thirds if k.startswith('T')]}; offset test {out['offset_test_subset']['ratio']}; "
          f"no companion {out['no_sign_of_a_companion_subset']['ratio']}; above 1e34.5 {out['above_1e34.5_subset']['ratio']}, "
          f"above 1e35 {out['above_1e35.0_subset']['ratio']}")
    print(f"wrote {out_path}")


if __name__ == "__main__":
    main()