"""The threshold rule of the author's own code, as a further, exploratory test (not part of the pre-registered one).

The preprint's section 2.2 compares every coordinate of the embedding from one minute to the next. The author's
planetary_dates.m (Zenodo record 21737429, version 1.1, MIT licence; we read it as text and did not run it) compares
only Earth's row: line 222, "maxDiff = max(abs(X(4,:) - Xold(4,:))); % Earth = row 4", with the Sun in row 1. It also
never tests the first minute of a year, because there is no previous minute in its loop. Everything else is as in
config_dates.py and select_dates.py, which this script imports unchanged: the same positions, Gram matrix, SVD,
threshold 0.35e-4 AU and nested body sets (the author's code stops with an error after Jupiter; so do we).

Writes selection-<tag>.json and dates-<tag>.json in the format of select_dates.py, so that overlap_test.py,
outsample_test.py and diagnose_dates.py can score and diagnose them. Needs NumPy.

Usage: python3 -I author_rule.py <horizons dir> <out dir> <first year> <last year> [method] [order]
  method: svd (default) or eigh; order: sunfirst (default), sunlast or reversed, as in config_dates.py
"""
import datetime, json, pathlib, sys
sys.path.insert(0, str(pathlib.Path(__file__).resolve().parent))
import numpy as np
import config_dates as cd

BASE = ["199", "299", "399", "499"]
SEQUENCE = [("301", 3), ("599", None)]          # the Moon after Earth, then Jupiter: the author's three sets


def earth_row(bodies, order):
    i = bodies.index("399")
    return {"sunfirst": 1 + i, "sunlast": i, "reversed": len(bodies) - i}[order]


def candidates(hdir, year, bodies, method="svd", order="sunfirst", chunk=44640):
    pos = [cd.minute_positions(hdir, b, year) for b in bodies]
    sun = np.zeros_like(pos[0])
    X = np.stack({"sunfirst": [sun] + pos, "sunlast": pos + [sun], "reversed": [sun] + pos[::-1]}[order], axis=1)
    r = earth_row(bodies, order)
    T = X.shape[0]; jumps = np.empty(T - 1)
    for a in range(0, T, chunk):
        b = min(T, a + chunk + 1)
        Xs, _ = cd.embed(X[a:b], method)
        jumps[a:b - 1] = np.abs(np.diff(Xs[:, r, :], axis=0)).max(axis=1)   # Earth's row only
    k = np.nonzero(jumps >= 0.35 * cd.EPS)[0]
    k = k[k > 0]                                  # the year's first minute is never tested
    jan1 = datetime.datetime(year, 1, 1)
    return [jan1 + datetime.timedelta(minutes=int(i)) for i in k], jumps


def select(hdir, year, method="svd", order="sunfirst", limit=100):
    bodies, steps = list(BASE), []
    for nxt in [None] + SEQUENCE:
        if nxt is not None:
            body, pos = nxt
            bodies = bodies[:pos] + [body] + bodies[pos:] if pos is not None else bodies + [body]
        stamps, jumps = candidates(hdir, year, bodies, method, order)
        steps.append({"bodies": [cd.NAMES[b] for b in bodies], "candidates": len(stamps)})
        if len(stamps) <= limit:
            return {"year": year, "rule": "author's code: Earth's row only, first minute untested",
                    "steps": steps, "selected": [cd.NAMES[b] for b in bodies],
                    "stamps": [s.strftime("%Y-%m-%d %H:%M") for s in stamps],
                    "dates": sorted({s.date().isoformat() for s in stamps}),
                    "jump_max_at_stamps_AU": [float(jumps[int((s - datetime.datetime(year, 1, 1)).total_seconds() // 60)])
                                              for s in stamps]}
    return {"year": year, "steps": steps, "selected": None, "stamps": [], "dates": [],
            "excluded": "no body set through Jupiter gave <= 100 candidates (the author's code stops with an error)"}


def main():
    hdir, out = sys.argv[1], pathlib.Path(sys.argv[2]); y0, y1 = int(sys.argv[3]), int(sys.argv[4])
    method = sys.argv[5] if len(sys.argv) > 5 else "svd"
    order = sys.argv[6] if len(sys.argv) > 6 else "sunfirst"
    out.mkdir(parents=True, exist_ok=True)
    tag = f"{y0}-{y1}" + ("" if (method, order) == ("svd", "sunfirst") else f"-{method}-{order}")
    res = [select(hdir, y, method, order) for y in range(y0, y1 + 1)]
    for name, obj in ((f"selection-{tag}.json", res), (f"dates-{tag}.json", {str(r["year"]): r["dates"] for r in res})):
        with open(out / name, "w", encoding="utf-8", newline="\n") as f:
            json.dump(obj, f, indent=1); f.write("\n")
    for r in res:
        print(r["year"], " -> ".join(f"{'+'.join(s['bodies'])}: {s['candidates']}" for s in r["steps"]),
              "| dates:", len(r["dates"]), r.get("excluded", ""))


if __name__ == "__main__":
    main()