#!/usr/bin/env python3
# =============================================================================
# worker_c4.py — Cycle 3: one distributed chunk of the C4' scan
# Author: Simon-Pierre Boucher — contact@spboucher.ai
# =============================================================================
# Scans gaps whose END prime lies in [lo, hi) and emits MERGEABLE partial sums:
#   S1/S2 (lag-1/lag-2 Pearson streaming sums), gap histogram, D = N2-N4,
#   chunk-local record gaps. A 5000-wide sieve buffer below lo recovers the
#   two gaps preceding the first counted one (max gap < 1e12 is ~540, and any
#   540-window below 1e12 contains a prime, so the buffer always suffices).
# Deterministic. Usage: python3 worker_c4.py <lo> <hi> <out.json>
# =============================================================================

import json
import platform
import sys
import time

import numpy as np

from core import primes_upto, sieve_segment

import os
BUFFER = 5000
SEG = int(os.environ.get("C4_SEG", 50_000_000))


def run(lo, hi, out_path, sieve_fn=None, tag=""):
    sieve = sieve_fn or sieve_segment
    t0 = time.time()
    base = primes_upto(int(hi ** 0.5) + 1)
    start = 2 if lo <= 2 else lo - BUFFER

    hist = np.zeros(6000, dtype=np.int64)
    S1 = dict(n=0, x=0.0, y=0.0, xx=0.0, yy=0.0, xy=0.0)
    S2 = dict(n=0, x=0.0, y=0.0, xx=0.0, yy=0.0, xy=0.0)
    D = 0
    records = []
    best = 0
    n_gaps = 0
    prev = None
    tail = []  # up to 2 gaps immediately preceding current segment's first gap
    s = start
    while s < hi:
        e = min(s + SEG, hi)
        primes = sieve(s, e, base)
        s = e
        if len(primes) == 0:
            continue
        if prev is not None:
            primes = np.concatenate(([prev], primes))
        if len(primes) >= 2:
            gaps = np.diff(primes)
            ends = primes[1:]
            starts = primes[:-1]
            inr = ends >= lo  # gaps counted by this worker (end prime >= lo)
            gin = gaps[inr]
            n_gaps += len(gin)
            hist += np.bincount(gin, minlength=len(hist))[: len(hist)]
            D += int((gin == 2).sum() - (gin == 4).sum())
            # pair sums: pair (g_{n-1}, g_n) / (g_{n-2}, g_n) counted at g_n
            gl = np.concatenate((np.array(tail, dtype=gaps.dtype), gaps))
            k = len(tail)
            idx = np.nonzero(inr)[0] + k          # positions of counted gaps in gl
            i1 = idx[idx >= 1]
            a = gl[i1 - 1].astype(np.float64); b = gl[i1].astype(np.float64)
            S1["n"] += len(a)
            S1["x"] += float(a.sum()); S1["y"] += float(b.sum())
            S1["xx"] += float((a * a).sum()); S1["yy"] += float((b * b).sum())
            S1["xy"] += float((a * b).sum())
            i2 = idx[idx >= 2]
            a2 = gl[i2 - 2].astype(np.float64); b2 = gl[i2].astype(np.float64)
            S2["n"] += len(a2)
            S2["x"] += float(a2.sum()); S2["y"] += float(b2.sum())
            S2["xx"] += float((a2 * a2).sum()); S2["yy"] += float((b2 * b2).sum())
            S2["xy"] += float((a2 * b2).sum())
            # chunk-local records among counted gaps
            if len(gin) and int(gin.max()) > best:
                sin = starts[inr]
                for i in np.nonzero(gin > best)[0]:
                    g, p = int(gin[i]), int(sin[i])
                    if g > best:
                        best = g
                        records.append([g, p])
            tail = [int(v) for v in gaps[-2:]]
        prev = primes[-1]

    nz = np.nonzero(hist)[0]
    out = {
        "lo": lo, "hi": hi, "n_gaps": int(n_gaps), "D": D,
        "S1": S1, "S2": S2,
        "hist": {int(g): int(hist[g]) for g in nz},
        "records": records,
        "host": platform.node() + tag,
        "seconds": round(time.time() - t0, 1),
    }
    with open(out_path, "w") as f:
        json.dump(out, f)
    print(f"done [{lo},{hi}) on {out['host']} in {out['seconds']}s", flush=True)


def main():
    lo, hi = int(float(sys.argv[1])), int(float(sys.argv[2]))
    out_path = sys.argv[3]
    from pathlib import Path as _P
    if _P(out_path).exists():
        print(f"skip existing {out_path}", flush=True)
        return
    if len(sys.argv) > 4 and sys.argv[4] == "gpu":
        from gpu_sieve import sieve_segment_gpu
        run(lo, hi, out_path, sieve_fn=sieve_segment_gpu, tag="/gpu")
    else:
        run(lo, hi, out_path)


if __name__ == "__main__":
    main()