#!/usr/bin/env python3
"""
A401136 -- self-contained check program, Python 3.8 or later.

    Square array read by antidiagonals downwards:
    Y(n,k) = (product of the first n-1 primes)*prime(n)^k - 1 for n, k >= 1.

Run:    python A401136.py                 (computes antidiagonals 1..140)
        python A401136.py b401136.txt     (also compares with that b-file)

Standard library only.

The program computes the first 140 antidiagonals, 9870 terms, from the
definition, with its own prime sieve, and checks them:

 * the 36 terms of the entry's DATA field;
 * row 1 against 2^k - 1, and column 1 against the product of the first n
   primes minus 1;
 * the row recurrence Y(n,k) = prime(n)*Y(n,k-1) + prime(n) - 1 at every cell
   with k >= 2;
 * the SHA-256 fingerprint of the 9870 lines "m a(m)" against the fingerprint
   of the b-file uploaded to the OEIS (its comment lines excluded), so that the
   whole b-file is checked without downloading it.

If a b-file is given, or b401136.txt or b401136_1.txt is in the current
directory, every term is also compared with it, line by line.

Reading order: antidiagonal d holds Y(1,d), Y(2,d-1), ..., Y(d,1), so the row
index n increases along it ("downwards"), and Y(n,k) is term number
a((n+k-1)*(n+k-2)/2 + n).
"""

import hashlib
import os
import sys
import time

ANTIDIAGONALS = 140

DATA = [1, 3, 5, 7, 17, 29, 15, 53, 149, 209, 31, 161, 749, 1469, 2309, 63, 485,
        3749, 10289, 25409, 30029, 127, 1457, 18749, 72029, 279509, 390389,
        510509, 255, 4373, 93749, 504209, 3074609, 5075069, 8678669, 9699689]

# SHA-256 of the lines "m a(m)\n", m = 1..9870, of the b-file b401136.txt.
BFILE_SHA256 = "217e484a4c3efd7c15b5674981f97e1e8d278b47ec1de4185c6da858c3cc429c"

BFILE_NAMES = ["b401136.txt", "b401136_1.txt"]


def first_primes(count):
    """The first `count` primes, by the sieve of Eratosthenes."""
    limit = 16
    while True:
        flags = bytearray([1]) * limit
        flags[0:2] = b"\x00\x00"
        for p in range(2, int(limit ** 0.5) + 1):
            if flags[p]:
                flags[p * p::p] = bytearray(len(range(p * p, limit, p)))
        found = [p for p in range(limit) if flags[p]]
        if len(found) >= count:
            return found[:count]
        limit *= 2


def position(n, k):
    """The index m with a(m) = Y(n,k)."""
    return (n + k - 1) * (n + k - 2) // 2 + n


def compute(D, primes):
    """a(1), ..., a(D*(D+1)/2): Y(n,k) for n + k <= D + 1, by antidiagonals."""
    h = [1]                                      # h[n-1] = product of the first n-1 primes
    for p in primes[:D - 1]:
        h.append(h[-1] * p)
    terms = []
    for d in range(1, D + 1):                    # antidiagonal d: n + k = d + 1
        for n in range(1, d + 1):
            terms.append(h[n - 1] * primes[n - 1] ** (d + 1 - n) - 1)
    return terms


def read_bfile(path):
    """[(m, a(m)), ...] from a b-file, skipping comment and blank lines."""
    pairs = []
    with open(path, encoding="ascii") as f:
        for line in f:
            line = line.strip()
            if line and not line.startswith("#"):
                m, value = line.split()
                pairs.append((int(m), int(value)))
    return pairs


def main(argv):
    t0 = time.time()
    D = ANTIDIAGONALS
    primes = first_primes(D)
    terms = compute(D, primes)
    Y = lambda n, k: terms[position(n, k) - 1]
    failures = []

    def check(label, ok):
        print("%s: %s" % (label, "ok" if ok else "FAILED"))
        if not ok:
            failures.append(label)

    print("A401136: %d terms = antidiagonals 1..%d, largest %d digits"
          % (len(terms), D, max(len(str(v)) for v in terms)))
    check("DATA (%d terms)" % len(DATA), terms[:len(DATA)] == DATA)
    check("row 1 = 2^k - 1 for k = 1..%d" % D,
          all(Y(1, k) == 2 ** k - 1 for k in range(1, D + 1)))
    primorial, column_ok = 1, True
    for n in range(1, D + 1):
        primorial *= primes[n - 1]
        column_ok &= Y(n, 1) == primorial - 1
    check("column 1 = A002110(n) - 1 for n = 1..%d" % D, column_ok)
    cells = [(n, k) for n in range(1, D + 1) for k in range(2, D + 2 - n)]
    check("Y(n,k) = prime(n)*Y(n,k-1) + prime(n) - 1 at all %d cells with k >= 2" % len(cells),
          all(Y(n, k) == primes[n - 1] * Y(n, k - 1) + primes[n - 1] - 1 for n, k in cells))
    lines = "".join("%d %d\n" % (m, v) for m, v in enumerate(terms, 1))
    check("SHA-256 of the %d lines = the b-file uploaded to the OEIS" % len(terms),
          hashlib.sha256(lines.encode("ascii")).hexdigest() == BFILE_SHA256)

    path = argv[1] if len(argv) > 1 else next((f for f in BFILE_NAMES if os.path.exists(f)), None)
    if path is None:
        print("b-file: none in this directory (optional; the fingerprint above covers it)")
    elif not os.path.exists(path):
        check("b-file %s: not found" % path, False)
    else:
        pairs = read_bfile(path)
        bad = [(m, v) for m, v in pairs if not (1 <= m <= len(terms)) or terms[m - 1] != v]
        for m, v in bad[:5]:
            print("MISMATCH at a(%d): computed %s, b-file has %d"
                  % (m, terms[m - 1] if 1 <= m <= len(terms) else "nothing", v))
        complete = [m for m, _ in pairs] == list(range(1, len(terms) + 1))
        check("b-file %s: %d terms, all equal" % (path, len(pairs)), complete and not bad)

    if failures:
        print("FAILED: %d check(s) did not pass." % len(failures))
        return 1
    print("OK: all checks passed. Total %d s." % int(time.time() - t0))
    return 0


if __name__ == "__main__":
    sys.exit(main(sys.argv))