"""Experiment 3: does perfect state transfer survive on the real (4, 27) knot with distance-set couplings?

Places the 108 beads on the (4, 27) torus knot (ring radius 1, tube radius r), couples every pair
with J = (1/d)^alpha for alpha = 3 (dipolar) and 6 (Rydberg-type), no cutoff, normalised so the
bead 1 to bead 28 coupling (the square partner) is 1. Starts on bead 1 and records the best
probability on bead 55 for 0 <= t <= t_max, for r from 0.02 to 0.45.

  python experiment-3-leak-geometry.py                    # t_max = 40, writes leak-test-geometry.json
  python experiment-3-leak-geometry.py --t-max 8 --out leak-test-geometry-first-pass.json

The time window matters on thin tubes: the square's diagonal coupling delays the first complete
arrival from t = pi/2 to t = 12 pi (1/d^6) or about 27 (1/d^3), so t_max = 8 sees only the first
pass. The findings page uses t_max = 40, the same window as wound-chain-check.py.

Supplied with kickoff patch 03 (t_max = 8); the window became a parameter, and the default 40,
on the author's ruling of 2026-09-23. The physics is unchanged. Requires NumPy.
d28_over_d27 and d1_over_d27 are for bead 1, which sits on the outside of the tube.
"""
import argparse, json
import numpy as np

P, Q, N = 4, 27, 108


def distances(r, R=1.0):
    t = 2 * np.pi * np.arange(N) / N; u, v = P * t, Q * t
    pos = np.stack([(R + r * np.cos(v)) * np.cos(u), (R + r * np.cos(v)) * np.sin(u), r * np.sin(v)], -1)
    return np.linalg.norm(pos[:, None] - pos[None], axis=-1)


def best_transfer(D, alpha, t_max, samples):
    J = np.where(np.eye(N) == 1, 0.0, (1.0 / np.maximum(D, 1e-9)) ** alpha); J /= J[0, 27]
    w, V = np.linalg.eigh(-J)
    ts = np.linspace(0, t_max, samples)
    amp = (np.exp(-1j * np.outer(ts, w)) * V[0]) @ V[54]
    return float((np.abs(amp) ** 2).max())


if __name__ == "__main__":
    ap = argparse.ArgumentParser()
    ap.add_argument("--t-max", type=float, default=40.0)
    ap.add_argument("--out", default="leak-test-geometry.json")
    a = ap.parse_args()
    samples = int(250 * a.t_max) + 1                     # 2001 for t_max = 8, as supplied
    rs = np.linspace(0.02, 0.45, 44)
    out = {"t_max": a.t_max, "r_over_R": rs.tolist(), "fidelity_1_over_d3": [], "fidelity_1_over_d6": [],
           "d28_over_d27": [], "d1_over_d27": []}
    for r in rs:
        D = distances(r)
        out["fidelity_1_over_d3"].append(best_transfer(D, 3, a.t_max, samples))
        out["fidelity_1_over_d6"].append(best_transfer(D, 6, a.t_max, samples))
        out["d28_over_d27"].append(D[0, 28] / D[0, 27]); out["d1_over_d27"].append(D[0, 1] / D[0, 27])
    with open(a.out, "w") as f:
        json.dump(out, f)
    for r, x, y in zip(rs[::6], out["fidelity_1_over_d3"][::6], out["fidelity_1_over_d6"][::6]):
        print(f"r/R={r:.2f}  P55(1/d3)={x:.3f}  P55(1/d6)={y:.3f}")