#!/usr/bin/env python3
"""
pfxor2 - windowed-palindrome keystream, bidirectional and constant-memory.

Key is (M, x, y): the window is defined by its RADIUS BOUNDS, not by a gap
count. Two scans run inward from x and from y; bit i pairs the i-th gap from
each end. Nothing is accumulated - each bit is emitted and the gaps discarded.

Consequences of keying on (x, y) rather than W:
  - memory is O(1) in the payload size
  - the two halves are independent processes; each segments internally
  - overshoot is free (truncate); undershoot is detected before writing
  - the pairing is identical to the accumulate-everything method, because a
    window's pairs are determined by its endpoints alone
"""
import argparse, hashlib, json, os, sys, time
import numpy as np
from multiprocessing import Process, Queue
from gmpy2 import mpz, is_prime
from sympy import primerange

SIEVE_LIMIT = 300_000          # legacy default

def sieve_limit_for(span):
    """Presieve cost is O(pi(limit)) per segment; benefit scales with span.
    Measured optimum tracks span/250 (see bench)."""
    return max(1000, min(2_000_000, int(span // 250)))
C2 = 0.6601618158468696


def _sieve_segment(M, base, span, sp, descending=False):
    """Radii in [base, base+span) surviving small-prime elimination, then BPSW."""
    parity = 1 if M % 2 == 0 else 0
    start = base + ((parity - base) % 2)
    n = (base + span - start + 1) // 2
    if n <= 0:
        return []
    alive = np.ones(n, dtype=bool)
    for p in sp:
        r = int(M % p)
        for t in (r, p - r):
            s = (t - start) % (2 * p)
            if s % 2:
                s = (s + p) % (2 * p)
            i0 = s // 2
            if i0 < n:
                alive[i0::p] = False
    offs = np.nonzero(alive)[0]
    if descending:
        offs = offs[::-1]
    out = []
    for o in offs:
        d = start + 2 * int(o)
        if is_prime(M - d) and is_prime(M + d):
            out.append(d)
    return out


def _scan(M, origin, direction, limit, seg, q, sp_limit):
    """Emit hits outward-to-inward from `origin`. direction=+1 from x, -1 from y."""
    sp = list(primerange(3, sieve_limit_for(seg)))
    M = mpz(M)
    pos = origin
    while True:
        if direction > 0:
            base, span = pos, seg
            if base >= limit:
                break
            span = min(span, limit - base)
        else:
            base = max(pos - seg, limit)
            span = pos - base
            if span <= 0:
                break
        hits = _sieve_segment(M, base, span, sp, descending=(direction < 0))
        if hits:
            q.put(hits)                      # one transfer per segment, not per hit
        pos = base + span if direction > 0 else base
    q.put(None)


def keystream_bidirectional(M, x, y, nbits, seg=20_000_000, verbose=False):
    qL, qR = Queue(maxsize=10000), Queue(maxsize=10000)
    pL = Process(target=_scan, args=(M, x, +1, y, seg, qL, SIEVE_LIMIT))
    pR = Process(target=_scan, args=(M, y, -1, x, seg, qR, SIEVE_LIMIT))
    pL.start(); pR.start()

    from collections import deque
    bufL, bufR = deque(), deque()
    doneL = doneR = False
    def pull(q, buf, done):
        if buf: return buf.popleft(), done
        if done: return None, done
        chunk = q.get()
        if chunk is None: return None, True
        buf.extend(chunk)
        return buf.popleft(), done

    bits = np.empty(nbits, dtype=np.uint8)
    got = ties = 0
    prevL, doneL = pull(qL, bufL, doneL); prevR, doneR = pull(qR, bufR, doneR)
    first_left, first_right = prevL, prevR
    lastL, lastR = prevL, prevR
    t0 = time.time()
    try:
        while got < nbits:
            a, doneL = pull(qL, bufL, doneL); b, doneR = pull(qR, bufR, doneR)
            if a is None or b is None:
                raise RuntimeError(
                    f"window exhausted: only {got:,} of {nbits:,} bits. "
                    f"Increase y.")
            if a >= b:
                raise RuntimeError(
                    f"scans crossed at d={a:,} with {got:,}/{nbits:,} bits. "
                    f"Increase y.")
            gL = a - prevL
            gR = prevR - b
            prevL, prevR = a, b
            lastL, lastR = a, b
            if gL == gR:
                ties += 1
                continue
            bits[got] = 1 if gL > gR else 0
            got += 1
            if verbose and got % 200000 == 0:
                print(f"  {got:,}/{nbits:,} bits, "
                      f"left at {a:,}, right at {b:,}, "
                      f"{time.time()-t0:.0f}s", file=sys.stderr)
    finally:
        pL.terminate(); pR.terminate(); pL.join(); pR.join()

    meta = {"x": int(x), "y": int(y), "ties": ties,
            "tie_rate": ties / max(got + ties, 1),
            "left_reached": int(lastL), "right_reached": int(lastR),
            "unused_middle": int(lastR - lastL),
            "first_hits": [int(first_left), int(first_right)],
            "ones": float(bits.mean())}
    return bits, meta


def suggest_y(M, x, nbits, margin=1.15):
    """Radius span needed for ~2*nbits gaps, with headroom."""
    lnM = float(np.log(float(M)))
    mean_gap = lnM ** 2 / (2 * C2)
    return int(x + 2 * nbits * mean_gap * margin)


def main():
    ap = argparse.ArgumentParser()
    ap.add_argument("infile"); ap.add_argument("outfile")
    ap.add_argument("-m", "--midpoint", required=True)
    ap.add_argument("-x", "--inner", required=True, help="inner radius bound")
    ap.add_argument("-y", "--outer", help="outer radius bound (omit to auto-size)")
    ap.add_argument("--seg", type=int, default=0, help="0 = auto from span")
    ap.add_argument("--meta"); ap.add_argument("-v", "--verbose", action="store_true")
    a = ap.parse_args()
    M = int(eval(str(a.midpoint))); x = int(eval(str(a.inner)))
    data = open(a.infile, "rb").read(); nbits = len(data) * 8
    y = int(eval(str(a.outer))) if a.outer else suggest_y(M, x, nbits)

    # Key is (M, x). Radius contributes log2(0.9*M / mean_gap) bits after
    # discounting hit-spacing quantisation; the search over the two is joint.
    lnM = float(np.log(float(M)))
    mean_gap = lnM**2 / (2*C2)
    keybits = np.log2(float(M)) + np.log2(0.9*float(M)/mean_gap)
    if keybits < 100:
        print(f"WARNING: combined keyspace only {keybits:.0f} bits.", file=sys.stderr)
    if x > 0.9*M:
        print("WARNING: x is in the skin - local density becomes measurable "
              "there and depth stops being hidden.", file=sys.stderr)
    print(f"payload {len(data):,} B ({nbits:,} bits)\n"
          f"M       {len(str(M))} digits\n"
          f"keyspace {keybits:.0f} bits (M {np.log2(float(M)):.0f} + radius "
          f"{keybits-np.log2(float(M)):.0f})\n"
          f"window  x = {x:,}\n        y = {y:,}\n"
          f"span    {y-x:,}", file=sys.stderr)

    seg = a.seg or max(2_000_000, min(80_000_000, (y - x) // 12))
    print(f"segment {seg:,}  presieve to {sieve_limit_for(seg):,}", file=sys.stderr)
    t0 = time.time()
    bits, meta = keystream_bidirectional(M, x, y, nbits, seg=seg, verbose=a.verbose)
    el = time.time() - t0
    out = bytes(np.frombuffer(data, dtype=np.uint8) ^ np.packbits(bits)[:len(data)])
    open(a.outfile, "wb").write(out)

    meta.update({"bytes": len(data), "seconds": round(el, 1),
                 "sha256_in": hashlib.sha256(data).hexdigest()[:16],
                 "sha256_out": hashlib.sha256(out).hexdigest()[:16]})
    print(f"\n{el:.1f}s ({len(data)/max(el,1e-9):.0f} B/s)\n"
          f"ties {meta['ties']:,} ({100*meta['tie_rate']:.3f}%)\n"
          f"ones {meta['ones']:.5f} (z {(meta['ones']-0.5)*2*np.sqrt(nbits):+.2f})\n"
          f"scans stopped at {meta['left_reached']:,} and {meta['right_reached']:,}, "
          f"{meta['unused_middle']:,} of span unused", file=sys.stderr)
    print(f"\nRADIUS RANGE CONSUMED: [{x:,} .. {y:,}] - never reuse or overlap.",
          file=sys.stderr)
    if a.meta:
        json.dump(meta, open(a.meta, "w"), indent=2)


if __name__ == "__main__":
    main()