#!/usr/bin/env python3
"""
coverage_density_probe.py — CoverageSystem braid falsification gate.

Assembles the 3-column coverage-density matrix (Goormaghtigh, LonelyRunner,
SpherionTwinPrime) and runs two diagnostics:
  (A) Spectral effective-rank (participation ratio) — resolves the structural-
      identity claim: rank → 1 means same operator; rank = 3 means independent.
  (B) Householder-QR of the coupling matrix — qr_error / tau per crossing gives
      the lossless oracle for CoverageSystem.DimensionalTransition in Lean.

All three densities are evaluated at integer w = 1..W, treating w as the
shared "target value" axis:

  d_G(w)  Goormaghtigh   : #{(x,m): x,m≥2, repunit(x,m)=w}
                           repunit(x,m) = (x^m - 1)/(x - 1)
  d_S(w)  SpherionTwin   : #{(a,b,sheet): obstruction(a,b,s)=w, a,b≥1}
                           obstruction = 6ab ± a ± b  (4 sign sheets)
  d_L(w)  LonelyRunner   : mean Φ(t_w, ·) over 128 circle points, k=3 runners
                           t_w = (w/W) * T_max  (sampling one period)

Claim under test: d_G, d_S, d_L are structurally identical (same operator at
different dimensionalities).  Effective rank < 2 → confirmed.

Braid crossing difficulty (Sidon slack proxy):
  Strand 0 ↔ Goormaghtigh  (Sidon address 1, slack 127)
  Strand 3 ↔ LonelyRunner   (Sidon address 8, slack 120)
  Strand 6 ↔ SpherionTwin   (Sidon address 64, slack 64)
  → Δ(0,3)=7, Δ(3,6)=56, Δ(0,6)=63  → prove C₀₃ first.

Usage:
  python3 coverage_density_probe.py [--W 200] [--T 5.0] [--k 3] [--json]
"""
from __future__ import annotations

import argparse
import hashlib
import json
import os
import sys
import time
from pathlib import Path

import numpy as np

ROOT = Path(__file__).resolve().parent
sys.path.insert(0, str(ROOT))

# ---------------------------------------------------------------------------
# 1. Goormaghtigh density
# ---------------------------------------------------------------------------

def repunit(x: int, m: int) -> int:
    """repunit(x,m) = (x^m - 1) // (x - 1)  for x >= 2, m >= 2."""
    return (x**m - 1) // (x - 1)


def goormaghtigh_density(W: int) -> np.ndarray:
    """d_G(w) = #{(x,m): x,m ≥ 2, repunit(x,m) = w} for w = 1..W."""
    d = np.zeros(W + 1, dtype=np.float64)
    # x^m grows fast; outer loop on x from 2 upward until x^2-1 > W*(x-1)
    x = 2
    while True:
        # minimum repunit value for this x is repunit(x,2) = x+1
        if x + 1 > W:
            break
        m = 2
        while True:
            rv = repunit(x, m)
            if rv > W:
                break
            d[rv] += 1.0
            m += 1
        x += 1
    return d[1:]  # return w=1..W  (0-indexed)


# ---------------------------------------------------------------------------
# 2. SpherionTwinPrime density
# ---------------------------------------------------------------------------

def spherion_density(W: int) -> np.ndarray:
    """d_S(w) = #{(a,b,sheet): obstruction(a,b,s) = w} for w = 1..W."""
    d = np.zeros(W + 1, dtype=np.float64)
    SIGNS = [(1, 1), (1, -1), (-1, 1), (-1, -1)]
    # bound: 6ab - a - b ≤ W  →  a ≤ (W + b) / (6b - 1)  (rough: ab ≤ W/4)
    max_ab = W // 4 + 2
    for s1, s2 in SIGNS:
        for a in range(1, max_ab + 1):
            for b in range(1, max_ab + 1):
                v = 6 * a * b + s1 * a + s2 * b
                if v < 1:
                    continue
                if v > W:
                    break
                d[v] += 1.0
    return d[1:]


# ---------------------------------------------------------------------------
# 3. LonelyRunner density
# ---------------------------------------------------------------------------

def lonely_density(W: int, k: int = 3, T_max: float = 5.0, N_theta: int = 128) -> np.ndarray:
    """d_L(w) = mean Φ(t_w, ·) over N_theta circle points, w = 1..W.

    t_w = (w / W) * T_max  (uniform time samples across one period).
    k runners with speeds 0, 1, ..., k-1.  δ = 1/(k+1).
    """
    speeds = list(range(k))
    delta = 1.0 / (k + 1.0)
    theta = np.linspace(0, 1.0, N_theta, endpoint=False)

    d = np.zeros(W, dtype=np.float64)
    for i, w in enumerate(range(1, W + 1)):
        t = (w / W) * T_max
        # runner positions mod 1
        pos = np.array([(s * t) % 1.0 for s in speeds])
        # coverage density Φ(t, θ) for each θ
        diff = np.abs(theta[:, np.newaxis] - pos[np.newaxis, :])   # (N_theta, k)
        dist = np.minimum(diff, 1.0 - diff)
        phi = (dist < delta).sum(axis=1)   # (N_theta,)
        d[i] = phi.mean()

    return d


# ---------------------------------------------------------------------------
# 4. Assemble 3-column matrix, run spectral probe + QR oracle
# ---------------------------------------------------------------------------

def participation_ratio(eig: np.ndarray) -> float:
    pos = eig[eig > 1e-9]
    if len(pos) == 0:
        return 0.0
    return float(pos.sum() ** 2 / np.square(pos).sum())


def householder_qr(M: np.ndarray) -> tuple[np.ndarray, np.ndarray, float, list[float]]:
    """Householder QR of M (n×3 or 3×3 coupling matrix).

    Returns Q, R, qr_error, list of tau values.
    τ = 2 and sparse nonzero pattern → lossless crossing.
    """
    Q, R = np.linalg.qr(M, mode="complete")
    reconstructed = Q @ R
    qr_error = float(np.max(np.abs(reconstructed[:R.shape[0], :] - M)))

    # Compute Householder reflector parameters manually for the 3-column case
    taus = []
    A = M.copy().astype(np.float64)
    n, p = A.shape
    for j in range(min(n, p)):
        x = A[j:, j].copy()
        norm_x = np.linalg.norm(x)
        if norm_x < 1e-14:
            taus.append(0.0)
            continue
        s = -np.sign(x[0]) if x[0] != 0 else -1.0
        u1 = x[0] - s * norm_x
        v = x.copy()
        v[0] = u1
        tau_j = 2.0 * u1**2 / np.dot(v, v) if np.dot(v, v) > 1e-14 else 0.0
        taus.append(float(tau_j))
        # apply reflector to update A
        if tau_j != 0.0:
            A[j:, j:] = A[j:, j:] - tau_j * np.outer(v, v @ A[j:, j:] / u1**2 * u1**2)

    return Q, R, qr_error, taus


def run_probe(W: int, k: int, T_max: float) -> dict:
    print(f"Building coverage-density matrix  W={W}, k={k}, T_max={T_max} ...")

    t0 = time.perf_counter()
    d_G = goormaghtigh_density(W)
    t_G = time.perf_counter() - t0
    print(f"  Goormaghtigh  : {d_G.sum():.0f} hits,  nonzero={int((d_G > 0).sum())}  ({t_G:.2f}s)")

    t0 = time.perf_counter()
    d_S = spherion_density(W)
    t_S = time.perf_counter() - t0
    print(f"  SpherionTwin  : {d_S.sum():.0f} hits,  nonzero={int((d_S > 0).sum())}  ({t_S:.2f}s)")

    t0 = time.perf_counter()
    d_L = lonely_density(W, k=k, T_max=T_max)
    t_L = time.perf_counter() - t0
    print(f"  LonelyRunner  : mean={d_L.mean():.4f}, std={d_L.std():.4f}  ({t_L:.2f}s)")

    # 3-column matrix: (W, 3)
    M = np.column_stack([d_G, d_S, d_L])

    # ── coupling matrix C = corrcoef of columns ──────────────────────────
    # drop rows that are all-zero to avoid degenerate correlations
    nonzero_rows = (M != 0).any(axis=1)
    M_nz = M[nonzero_rows]
    print(f"\n  Non-zero sample points: {M_nz.shape[0]} / {W}")

    C = np.corrcoef(M_nz, rowvar=False)
    print(f"\n  Coupling matrix C (corrcoef):")
    labels = ["Goor", "Spher", "LR"]
    print(f"         {'  '.join(f'{l:>8}' for l in labels)}")
    for i, row in enumerate(C):
        print(f"  {labels[i]:>6} " + "  ".join(f"{v:+8.4f}" for v in row))

    eig = np.sort(np.linalg.eigvalsh(C))[::-1]
    pr = participation_ratio(eig)
    print(f"\n  Eigenvalues: {', '.join(f'{e:.4f}' for e in eig)}")
    print(f"  Effective rank (participation ratio): {pr:.4f}")

    # ── Householder-QR of the 3×3 coupling matrix ────────────────────────
    Q, R, qr_error, taus = householder_qr(C)
    print(f"\n  Householder-QR of C  (lossless oracle):")
    print(f"    qr_error = {qr_error:.2e}   (0 = exact = lossless crossing)")
    print(f"    tau per reflector: {[f'{t:.4f}' for t in taus]}")
    print(f"    tau=2.0 → orthogonal reflector → lossless;  tau<2 → lossy")

    # ── braid crossing assessment ─────────────────────────────────────────
    # C₀₃ = Goor ↔ LR (strands 0,3), C₃₆ = LR ↔ Spher (strands 3,6)
    c03 = float(abs(C[0, 2]))   # Goor-LR correlation
    c36 = float(abs(C[2, 1]))   # LR-Spher correlation
    c06 = float(abs(C[0, 1]))   # Goor-Spher correlation

    SIDON_SLACKS = {0: 127, 3: 120, 6: 64}
    delta_03 = SIDON_SLACKS[0] - SIDON_SLACKS[3]   # 7
    delta_36 = SIDON_SLACKS[3] - SIDON_SLACKS[6]   # 56
    delta_06 = SIDON_SLACKS[0] - SIDON_SLACKS[6]   # 63

    print(f"\n  Braid crossing assessment (C_ij = |corr|, higher = tighter coupling):")
    print(f"    C₀₃  Goor↔LR   : |corr|={c03:.4f}  Sidon_Δ={delta_03}")
    print(f"    C₃₆  LR↔Spher  : |corr|={c36:.4f}  Sidon_Δ={delta_36}")
    print(f"    C₀₆  Goor↔Spher: |corr|={c06:.4f}  Sidon_Δ={delta_06}")

    # ── verdict ───────────────────────────────────────────────────────────
    print(f"\n{'='*66}")
    print(f"VERDICT:")
    if pr < 1.5:
        verdict = "SAME_OPERATOR — effective rank near 1; structural identity CONFIRMED"
    elif pr < 2.5:
        verdict = "PARTIAL_IDENTITY — rank ~2; one pair structurally identical, third independent"
    else:
        verdict = "INDEPENDENT — effective rank 3; structural identity NOT confirmed on this parameterization"
    print(f"  {verdict}")
    print(f"  qr_error={qr_error:.2e} → C₀₃ lossless: {'YES (exact)' if qr_error < 1e-10 else 'NO (lossy)'}")
    proof_order = sorted([("C₀₃", c03, delta_03), ("C₃₆", c36, delta_36), ("C₀₆", c06, delta_06)],
                         key=lambda x: -x[1])
    print(f"  Proof order (tightest coupling first): {' → '.join(x[0] for x in proof_order)}")
    print(f"{'='*66}")

    return {
        "schema": "coverage_density_probe_v1",
        "computed_at": time.strftime("%Y-%m-%dT%H:%M:%SZ"),
        "params": {"W": W, "k": k, "T_max": T_max},
        "density_stats": {
            "goormaghtigh": {"total_hits": float(d_G.sum()), "nonzero": int((d_G > 0).sum()),
                             "max": float(d_G.max()), "first_20": d_G[:20].tolist()},
            "spherion": {"total_hits": float(d_S.sum()), "nonzero": int((d_S > 0).sum()),
                         "max": float(d_S.max()), "first_20": d_S[:20].tolist()},
            "lonely_runner": {"mean": float(d_L.mean()), "std": float(d_L.std()),
                              "min": float(d_L.min()), "max": float(d_L.max()),
                              "first_20": d_L[:20].tolist()},
        },
        "coupling_matrix": C.tolist(),
        "eigenvalues": eig.tolist(),
        "effective_rank": float(pr),
        "qr_oracle": {
            "qr_error": qr_error,
            "taus": taus,
            "lossless": qr_error < 1e-10,
        },
        "crossings": {
            "C03_Goor_LR":    {"corr": c03, "sidon_delta": delta_03},
            "C36_LR_Spher":   {"corr": c36, "sidon_delta": delta_36},
            "C06_Goor_Spher": {"corr": c06, "sidon_delta": delta_06},
        },
        "verdict": verdict,
        "proof_order": [x[0] for x in proof_order],
        "caveat": (
            "LonelyRunner density is time-averaged Φ(t,θ) — a continuous S¹ quantity "
            "sampled at w/W*T_max time steps; not directly comparable to discrete ℕ-valued "
            "densities. Rank collapse would indicate temporal oscillation structure matches "
            "the obstruction-count distribution, not type-theoretic identity."
        ),
    }


# ---------------------------------------------------------------------------
# CLI
# ---------------------------------------------------------------------------

def main() -> None:
    ap = argparse.ArgumentParser(description="CoverageSystem braid falsification gate")
    ap.add_argument("--W", type=int, default=200, help="Max target value (default 200)")
    ap.add_argument("--k", type=int, default=3, help="LonelyRunner runner count (default 3)")
    ap.add_argument("--T", type=float, default=5.0, help="LonelyRunner time window (default 5.0)")
    ap.add_argument("--json", action="store_true", help="Print full receipt JSON")
    ap.add_argument("--out", type=str, default=None, help="Write receipt to file")
    args = ap.parse_args()

    receipt = run_probe(W=args.W, k=args.k, T_max=args.T)

    out_path = args.out
    if out_path is None:
        out_path = str(ROOT.parent.parent / "shared-data/data/coverage_density_probe_receipt.json")

    Path(out_path).write_text(json.dumps(receipt, indent=2, ensure_ascii=False))
    print(f"\nReceipt: {out_path}")

    if args.json:
        print(json.dumps(receipt, indent=2))


if __name__ == "__main__":
    main()