"""Blind replication of four quantum tests on the 108-bead (4, 27) network.

Requires only Python 3 and NumPy. Prints a JSON report to stdout.
See replication-protocol.md for the definitions. This file contains no expected values.
"""
import json, math
import numpy as np


def ring_hamiltonian(n, hops):
    """hops: dict {step: strength}. H[k, k±step] -= strength."""
    h = np.zeros((n, n))
    for step, t in hops.items():
        if t:
            for k in range(n):
                h[k, (k + step) % n] -= t
                h[k, (k - step) % n] -= t
    return h


def evolver(h):
    """Exact evolution via diagonalisation of a real symmetric H."""
    w, v = np.linalg.eigh(h)

    def u(psi0, t):
        return v @ (np.exp(-1j * w * t) * (v.conj().T @ psi0))

    return u


report = {}

# ---------------- Test 1: relative phase ----------------
N = 108
u = evolver(ring_hamiltonian(N, {27: 1}))
ts = np.linspace(0, 4, 4001)
t1 = {}
for deg in (0, 45, 90, 135, 180):
    psi0 = np.zeros(N, complex)
    psi0[0] = 1
    psi0[54] = np.exp(1j * math.radians(deg))
    psi0 /= np.linalg.norm(psi0)
    p_empty = np.array([abs(u(psi0, t)[27]) ** 2 + abs(u(psi0, t)[81]) ** 2 for t in ts])
    entry = {"max_P_empty": float(p_empty.max())}
    if deg == 0:
        i = int(np.argmax(p_empty > p_empty.max() - 1e-5))  # first time the maximum is reached (grid tolerance)
        entry["first_max_time"] = float(ts[i])
        entry["first_max_value"] = float(p_empty[i])
    t1[f"phase_{deg}_deg"] = entry
report["test1_relative_phase"] = t1

# ---------------- Test 2: factorisation ----------------
u2 = evolver(ring_hamiltonian(N, {27: 1, 4: 1}))
psi0 = np.zeros(N, complex); psi0[0] = 1
ts2 = np.linspace(0, 4 * math.pi, 2000)
F = np.array([abs(np.vdot(psi0, u2(psi0, t))) ** 2 for t in ts2])
u4 = evolver(ring_hamiltonian(4, {1: 1})); d4 = np.zeros(4, complex); d4[0] = 1
u27 = evolver(ring_hamiltonian(27, {1: 1})); d27 = np.zeros(27, complex); d27[0] = 1
F4 = np.array([abs(np.vdot(d4, u4(d4, t))) ** 2 for t in ts2])
F27 = np.array([abs(np.vdot(d27, u27(d27, t))) ** 2 for t in ts2])
report["test2_factorisation"] = {
    "max_abs_diff_F_vs_F4_times_F27": float(np.max(np.abs(F - F4 * F27))),
    "F_at_pi_multiples": {f"{m}pi": float(np.interp(m * math.pi, ts2, F)) for m in (1, 2, 3, 4)},
}

# ---------------- Test 3: quantum vs classical walk ----------------
u3 = evolver(ring_hamiltonian(N, {1: 1}))
psi0 = np.zeros(N, complex); psi0[0] = 1
k = np.arange(N); d = np.minimum(k, N - k).astype(float)
L = np.zeros((N, N))
for i in range(N):
    L[i, (i + 1) % N] += 1; L[i, (i - 1) % N] += 1; L[i, i] -= 2
wl, vl = np.linalg.eigh(L); p0 = np.zeros(N); p0[0] = 1
t3 = {}
for t in (5, 10, 20):
    sq = math.sqrt(float(np.sum(np.abs(u3(psi0, t)) ** 2 * d ** 2)))
    pc = vl @ (np.exp(wl * t) * (vl.T @ p0))
    sc = math.sqrt(max(0.0, float(np.sum(pc * d ** 2))))
    t3[f"t={t}"] = {"quantum_spread": sq, "classical_spread": sc}
# simple fits over t in [1, 20]
tt = np.linspace(1, 20, 40)
sq_all = np.array([math.sqrt(float(np.sum(np.abs(u3(psi0, t)) ** 2 * d ** 2))) for t in tt])
sc_all = np.array([math.sqrt(max(0.0, float(np.sum((vl @ (np.exp(wl * t) * (vl.T @ p0))) * d ** 2)))) for t in tt])
t3["fit_quantum_linear_a"] = float(np.sum(sq_all * tt) / np.sum(tt * tt))
t3["fit_classical_sqrt_a"] = float(np.sum(sc_all * np.sqrt(tt)) / np.sum(tt))
report["test3_walks"] = t3

# ---------------- Test 4: other bead counts ----------------
def coprime_splits(n):
    return [(p, n // p) for p in range(2, math.isqrt(n) + 1) if n % p == 0 and math.gcd(p, n // p) == 1]

t4 = {}
for n in (27, 54, 108, 1008):
    splits = coprime_splits(n)
    if not splits:
        t4[str(n)] = "no coprime split"
        continue
    rows = []
    for p, q in splits:
        h = ring_hamiltonian(n, {q: 1})
        w, v = np.linalg.eigh(h)
        a, b = 0, ((p // 2) * q) % n
        psi0 = np.zeros(n, complex); psi0[a] = 1
        if b != a:
            psi0[b] = 1
        psi0 /= np.linalg.norm(psi0)
        c = v.conj().T @ psi0
        sites = [a] + ([b] if b != a else [])
        tt4 = np.linspace(0, 12, 1201)
        P = []
        for t in tt4:
            psi = v @ (np.exp(-1j * w * t) * c)
            P.append(sum(abs(psi[s]) ** 2 for s in sites))
        P = np.array(P)
        # frequencies present in P_start(t): eigenvalue differences weighted by overlaps
        amp = {}
        idx = np.where(np.abs(c) > 1e-9)[0]
        for i in idx:
            for j in idx:
                f = round(abs(w[i] - w[j]), 4)
                if f == 0:
                    continue
                wgt = abs(sum(c[i] * np.conj(c[j]) * v[s, i] * np.conj(v[s, j]) for s in sites))
                amp[f] = amp.get(f, 0.0) + wgt
        tot = sum(amp.values()) or 1.0
        nfreq = sum(1 for f, x in amp.items() if x > 0.01 * tot)
        rows.append({"split": f"{p} x {q}", "cycle_length": p, "min_P_start": float(P.min()),
                     "n_frequencies_above_1pct": int(nfreq)})
    t4[str(n)] = rows
report["test4_other_counts"] = t4

report["method"] = "exact diagonalisation (numpy.linalg.eigh); classical walk via matrix exponential of the Laplacian"
print(json.dumps(report, indent=1))
