"""What happens at each candidate timestamp, and are the planets aligned on those dates?

For every candidate minute t: the embedding before and after (X*(t-1), X*(t)); the signed
column permutation P that best maps one onto the other; the residual after undoing P
(the genuine one-minute change); the Gram matrix's three leading eigenvalues.
Alignment on each date (00:00 TDB of the date, the four inner planets): heliocentric
ecliptic longitudes L_k; R1 = |mean exp(i L_k)| (all on one side of the Sun = 1);
R2 = |mean exp(2i L_k)| (all on one line through the Sun, either side = 1). Each date's
value is ranked against every day of the same year.

Usage: python -I diagnose_dates.py <horizons dir> <selection json> <out json>
"""
import datetime, itertools, json, pathlib, sys
sys.path.insert(0, str(pathlib.Path(__file__).resolve().parent))
import numpy as np
import config_dates as cd

INV = {v: k for k, v in cd.NAMES.items()}


def best_signed_perm(A, B):
    best = None
    for perm in itertools.permutations(range(3)):
        for signs in itertools.product((1, -1), repeat=3):
            M = B[:, perm] * np.array(signs)
            r = np.abs(M - A).max()
            if best is None or r < best[0]:
                best = (r, perm, signs)
    return best


def main():
    hdir, sel_path, out_path = sys.argv[1], sys.argv[2], sys.argv[3]
    sel = json.load(open(sel_path, encoding="utf-8"))
    out = []
    for yr in sel:
        year = yr["year"]; bodies = [INV[n] for n in yr["selected"]]
        pos = {b: cd.minute_positions(hdir, b, year) for b in set(bodies) | {"199", "299", "399", "499"}}
        X = np.stack([np.zeros_like(pos[bodies[0]])] + [pos[b] for b in bodies], axis=1)
        # alignment of the four inner planets at 00:00 of every day (row 1 + 1440*d)
        rows = 1 + 1440 * np.arange(X.shape[0] // 1440)
        L = np.stack([np.arctan2(pos[b][rows, 1], pos[b][rows, 0]) for b in ("199", "299", "399", "499")], axis=1)
        R1 = np.abs(np.exp(1j * L).mean(axis=1)); R2 = np.abs(np.exp(2j * L).mean(axis=1))
        days = [datetime.date(year, 1, 1) + datetime.timedelta(days=int(d)) for d in range(len(rows))]
        rank = lambda arr, i: float((arr < arr[i]).mean())
        for s in yr["stamps"]:
            t = datetime.datetime.strptime(s, "%Y-%m-%d %H:%M")
            k = int((t - datetime.datetime(year, 1, 1)).total_seconds() // 60)
            Xs, S = cd.embed(X[k:k + 2])                 # grid rows k, k+1 = minutes k-1, k
            r, perm, signs = best_signed_perm(Xs[1], Xs[0])
            d = (t.date() - datetime.date(year, 1, 1)).days
            out.append({"stamp": s, "jump_AU": float(np.abs(Xs[1] - Xs[0]).max()),
                        "after_undoing": {"perm": perm, "signs": signs, "residual_AU": float(r)},
                        "eigenvalues_AU2": [float(x) for x in S[1, :3]],
                        "R1_same_side": round(float(R1[d]), 3), "R1_pct_of_year": round(rank(R1, d), 3),
                        "R2_on_a_line": round(float(R2[d]), 3), "R2_pct_of_year": round(rank(R2, d), 3)})
    pathlib.Path(out_path).write_text(json.dumps(out, indent=1), encoding="utf-8")
    flips = sum(o["after_undoing"]["residual_AU"] < 3.5e-5 for o in out)
    print(f"{len(out)} candidates; {flips} are a pure sign flip or swap of singular vectors "
          f"(residual below the 3.5e-5 AU threshold once undone)")
    for o in out:
        print(o["stamp"], "jump %.3g" % o["jump_AU"], "signs", o["after_undoing"]["signs"], "perm", o["after_undoing"]["perm"],
              "resid %.2g" % o["after_undoing"]["residual_AU"], "| R1 %.2f (pct %.2f)  R2 %.2f (pct %.2f)" %
              (o["R1_same_side"], o["R1_pct_of_year"], o["R2_on_a_line"], o["R2_pct_of_year"]))


if __name__ == "__main__":
    main()