#!/usr/bin/env python3
"""
scan_quartet.py  -  Repair / discover the (K, delta) that completes a
valid Martini Quartet for a TRUSTED twin midpoint MT.

The twin midpoint MT is trusted (MT-1, MT+1 already prime). This tool
finds which K (for a fixed delta) or which even delta (for a fixed K)
makes the Martini pair MP +/- delta/2 prime, where MP = MT * K.

Useful when a published K or delta is suspected to be a typo: pin the
field you trust, scan the other.

Screening uses gmpy2 BPSW (fast, near-certain). Confirm any hit with
PFGW / verify_quartet.py before publishing.

Modes:
  --scan-k        fix --delta, search K in a window around --k
  --scan-delta    fix --k,     search even delta in a window around --delta
  --typo-only     restrict search to single-edit typos of the pinned-as-
                  suspect field (transpose / delete / substitute / insert)

Usage:
  py scan_quartet.py --mt-file mt1051.txt --k 3545 --delta 297914 --scan-k --k-window 12000
  py scan_quartet.py --mt-file mt1051.txt --k 3545 --delta 297914 --scan-delta --delta-window 200000
  py scan_quartet.py --mt-file mt1051.txt --k 3545 --delta 297914 --typo-only
"""

import argparse
import sys
import time

try:
    import gmpy2
    from gmpy2 import mpz, is_prime
except ImportError:
    sys.exit("gmpy2 not found. Install with:  pip install gmpy2 --break-system-packages")


def load_mt(args):
    if args.mt_file:
        raw = open(args.mt_file).read()
    elif args.mt:
        raw = args.mt
    else:
        sys.exit("Provide MT via --mt or --mt-file.")
    cleaned = "".join(c for c in raw if c.isdigit())
    if not cleaned:
        sys.exit("No digits found in MT input.")
    return mpz(cleaned)


def martini_ok(MT, K, D):
    """True if MP +/- D/2 are both prime, where MP = MT*K."""
    if D % 2:
        return False
    MP = MT * mpz(K)
    h = mpz(D) // 2
    return is_prime(MP - h) and is_prime(MP + h)


def typo_variants(value):
    """All single-edit variants of an integer's decimal string:
    adjacent transpositions, deletions, substitutions, insertions."""
    s = str(value)
    out = set()
    for i in range(len(s) - 1):                       # transpose adjacent
        l = list(s); l[i], l[i + 1] = l[i + 1], l[i]; out.add(int("".join(l)))
    for i in range(len(s)):                           # delete one
        if len(s) > 1:
            out.add(int(s[:i] + s[i + 1:]))
    for i in range(len(s)):                           # substitute one
        for d in "0123456789":
            if d != s[i]:
                out.add(int(s[:i] + d + s[i + 1:]))
    for i in range(len(s) + 1):                       # insert one
        for d in "0123456789":
            out.add(int(s[:i] + d + s[i:]))
    out.discard(int(s))
    return sorted(v for v in out if v > 0)


def report_hit(K, D, K0, D0):
    print("\n" + "=" * 56)
    print(" VALID QUARTET FOUND")
    print("=" * 56)
    print(f"   K     = {K}    (published: {K0})")
    print(f"   delta = {D}    (published: {D0})")
    if K != K0:
        print(f"   -> correction is in K: {K0} -> {K}")
    if D != D0:
        print(f"   -> correction is in delta: {D0} -> {D}")
    print("\n Confirm before publishing:")
    print(f"   py verify_quartet.py --mt-file <MT file> --k {K} --delta {D}")


def main():
    ap = argparse.ArgumentParser(description="Find the (K, delta) that completes a Martini quartet for a trusted MT.")
    g = ap.add_mutually_exclusive_group()
    g.add_argument("--mt", help="MT as decimal string")
    g.add_argument("--mt-file", help="file containing MT")
    ap.add_argument("--k", required=True, type=int, help="published / starting K")
    ap.add_argument("--delta", required=True, type=int, help="published / starting delta")
    ap.add_argument("--scan-k", action="store_true", help="fix delta, scan K")
    ap.add_argument("--scan-delta", action="store_true", help="fix K, scan delta")
    ap.add_argument("--typo-only", action="store_true", help="only test single-edit typos of the suspect field(s)")
    ap.add_argument("--k-window", type=int, default=10000, help="+/- range for K scan (default 10000)")
    ap.add_argument("--delta-window", type=int, default=200000, help="+/- range for delta scan (default 200000)")
    ap.add_argument("--all", action="store_true", help="list all hits instead of stopping at first")
    args = ap.parse_args()

    MT = load_mt(args)
    K0, D0 = args.k, args.delta

    print(f"MT: {MT.num_digits()} digits")
    if not (is_prime(MT - 1) and is_prime(MT + 1)):
        print("WARNING: MT-1 / MT+1 are NOT both prime -- MT itself may be wrong.")
    else:
        print("MT twin pair (MT-1, MT+1): both PRP  [trusted]")

    hits = []

    # ---- typo-only mode: cheap, try single edits of each field first ----
    if args.typo_only or (not args.scan_k and not args.scan_delta):
        print("\nSingle-edit scan: delta variants (K fixed) ...")
        t = time.time()
        for D in typo_variants(D0):
            if martini_ok(MT, K0, D):
                hits.append((K0, D))
                if not args.all:
                    break
        print(f"  done [{time.time()-t:.1f}s]")
        if not (hits and not args.all):
            print("Single-edit scan: K variants (delta fixed) ...")
            t = time.time()
            for K in typo_variants(K0):
                if martini_ok(MT, K, D0):
                    hits.append((K, D0))
                    if not args.all:
                        break
            print(f"  done [{time.time()-t:.1f}s]")

    # ---- windowed K scan ----
    if args.scan_k and not (hits and not args.all):
        lo, hi = max(1, K0 - args.k_window), K0 + args.k_window
        print(f"\nScanning K in [{lo}, {hi}] at delta={D0} ...")
        t = time.time()
        # search outward from K0 so the closest hit comes first
        order = sorted(range(lo, hi + 1), key=lambda k: abs(k - K0))
        for K in order:
            if martini_ok(MT, K, D0):
                hits.append((K, D0))
                if not args.all:
                    break
        print(f"  done [{time.time()-t:.1f}s]")

    # ---- windowed delta scan ----
    if args.scan_delta and not (hits and not args.all):
        D0e = D0 if D0 % 2 == 0 else D0 + 1
        lo, hi = max(2, D0e - args.delta_window), D0e + args.delta_window
        print(f"\nScanning even delta in [{lo}, {hi}] at K={K0} ...")
        t = time.time()
        MP = MT * mpz(K0)
        order = sorted(range(lo, hi + 1, 2), key=lambda d: abs(d - D0e))
        for D in order:
            h = mpz(D) // 2
            if is_prime(MP - h) and is_prime(MP + h):
                hits.append((K0, D))
                if not args.all:
                    break
        print(f"  done [{time.time()-t:.1f}s]")

    if not hits:
        print("\nNo valid quartet found in the searched space.")
        print("Try a wider --k-window / --delta-window, or check MT itself.")
        sys.exit(1)

    if args.all:
        print(f"\nAll hits ({len(hits)}):")
        for K, D in hits:
            print(f"   K={K}, delta={D}")
    else:
        K, D = hits[0]
        report_hit(K, D, K0, D0)


if __name__ == "__main__":
    main()
