"""Upper bounds on the most lines (sets) among n points of AG(d,3), by counting
over the parallel classes of hyperplanes. Exact integer arithmetic throughout.

For a set S of n points with L lines, and a parallel class W of hyperplanes
(3 parallel copies of AG(d-1,3)) cutting S into parts of sizes a, b, c:

  (1) summed over all classes, sum C(k,2) = c1 * C(n,2)
  (2) summed over all classes, sum C(k,3) = c2 * C(n,3) + (c1 - c2) * L
  (3) in each class, L <= gsub(a) + gsub(b) + gsub(c) + min(ab, bc, ca)

where c1 = (3^(d-1) - 1)/2 classes put two given points (or a line) in one
hyperplane, c2 = (3^(d-2) - 1)/2 classes do so for three points not on a
line, and gsub is the exact maximum one dimension down. To rule out L >= T,
keep only the splits allowed by (3) at L = T, and check that no choice of
splits for the classes meets (1) while reaching (2) at L = T. Larger L only
removes splits and raises the target, so it's ruled out too.

  python3 bound.py            # AG(4,3), n = 3..81, from the AG(3,3) maxima (third.c)
  python3 bound.py selftest   # AG(3,3) from the AG(2,3) maxima, to compare with third.c
  python3 bound.py dims       # AG(5,3) and up, n <= 27, by induction on d
"""
import sys
from math import comb
import numpy as np

# exact maxima in AG(3,3) for k = 0..27 points, from third.c (out/third.txt)
G3 = [0, 0, 0, 1, 1, 2, 3, 5, 8, 12, 12, 13, 14, 16, 19, 23, 26, 30, 36, 41,
      47, 54, 62, 71, 81, 92, 104, 117]
# exact maxima in AG(2,3), the 3x3 grid (9 points, 12 lines)
G2 = [0, 0, 0, 1, 1, 2, 3, 5, 8, 12]


def splits(n, cap):
    """Ways to cut n points into three parallel hyperplanes of `cap` points."""
    return [(a, b, n - a - b) for a in range(n + 1) for b in range(n - a + 1)
            if a >= b >= n - a - b and a <= cap]


def rules_out(d, gsub, n, target):
    """Can no n-set of AG(d,3) have `target` or more lines? (verdict, how, details)"""
    classes = (3 ** d - 1) // 2
    c1 = (3 ** (d - 1) - 1) // 2
    c2 = (3 ** (d - 2) - 1) // 2
    cap = len(gsub) - 1  # parts must fit where gsub is known (a hyperplane, or n <= 27)
    h = lambda p: gsub[p[0]] + gsub[p[1]] + gsub[p[2]] + min(p[0] * p[1], p[1] * p[2], p[0] * p[2])
    if n > 3 * cap:
        raise ValueError('gsub too short')
    ok = [p for p in splits(n, cap) if h(p) >= target]
    pair_need = c1 * comb(n, 2)
    trip_need = c2 * comb(n, 3) + (c1 - c2) * target
    if not ok:
        return True, 'no split allowed', {'splits': []}
    P = {p: sum(comb(k, 2) for k in p) for p in ok}
    T = {p: sum(comb(k, 3) for k in p) for p in ok}
    det = {'splits': ok}
    if pair_need > classes * max(P.values()):
        return True, 'pairs alone (too many)', det
    if pair_need < classes * min(P.values()):
        return True, 'pairs alone (too few)', det
    if classes * max(T.values()) < trip_need:
        return True, 'triples alone', det
    # exact: row[s] = most triples from k classes whose pairs add up to s (-1: impossible)
    row = np.full(pair_need + 1, -1, dtype=np.int64)
    row[0] = 0
    for _ in range(classes):
        nxt = np.full(pair_need + 1, -1, dtype=np.int64)
        for p in ok:
            if P[p] > pair_need:
                continue
            src = row[:pair_need + 1 - P[p]]
            cand = np.where(src >= 0, src + T[p], -1)
            np.maximum(nxt[P[p]:], cand, out=nxt[P[p]:])
        row = nxt
    reach = int(row[pair_need])
    det['trip_max_exact'] = reach
    return reach < trip_need, 'both together', det


def proved_max(d, gsub, n, start):
    """The largest line count the argument can't rule out, from `start` up."""
    t = start + 1
    while not rules_out(d, gsub, n, t)[0]:
        t += 1
    return t - 1, rules_out(d, gsub, n, t)


def built(n):
    """Lines in the best layout of this simple kind: a full 27-card hyperplane
    (117 lines) plus a and b points in the two parallel ones. Every pair across
    those two completes a line through the full hyperplane."""
    if n <= 27:
        return G3[n]
    return max(117 + G3[a] + G3[n - 27 - a] + a * (n - 27 - a)
               for a in range(max(0, n - 54), min(27, n - 27) + 1))


def main():
    arg = sys.argv[1] if len(sys.argv) > 1 else ''
    if arg == 'selftest':
        print('AG(3,3) from the 3x3 grid: proved maximum against the exhaustive count (third.c)')
        for n in range(3, 28):
            m, (_, how, det) = proved_max(3, G2, n, G3[n])
            print(f'n={n:2d}  exhaustive {G3[n]:3d}  proved at most {m:3d}  '
                  f'{"TIGHT" if m == G3[n] else "gap"}  ({how})')
    elif arg == 'dims':
        # Induction on d: if every set of k <= 27 points in AG(d-1,3) has at most
        # G3[k] lines, the argument in AG(d,3) gives the same for n <= 27 (parts
        # have at most 27 points, so G3 bounds them). The splits it allows don't
        # depend on d. 'no split allowed' then holds for every d; 'pairs alone
        # (too many)' needs c1/classes * C(n,2) > max pairs over the allowed
        # splits, and c1/classes = 1/3 - (2/3)/(3^d - 1) only grows with d. So
        # once a dimension D settles every n in one of those two ways, so does
        # every dimension after it.
        top = int(sys.argv[2]) if len(sys.argv) > 2 else 6
        for d in range(5, top + 1):
            hows = []
            for n in range(3, 28):
                v, how, _ = rules_out(d, G3, n, G3[n] + 1)
                if not v:
                    print(f'AG({d},3): n={n} NOT ruled out'); return
                hows.append(how)
            print(f'AG({d},3): every n <= 27 holds at most G3[n] lines '
                  f'({", ".join(f"{h}: {hows.count(h)}" for h in sorted(set(hows)))})', flush=True)
            if all(h in ('no split allowed', 'pairs alone (too many)') for h in hows):
                print(f'  so, by the monotone pair count, every AG(d,3) with d >= {d} does too')
                return
    else:
        print('AG(4,3): n points, the best layout found (a hyperplane, plus points in')
        print('the two parallel ones) against the most the counting argument allows')
        for n in range(3, 61):
            b = built(n)
            m, (_, how, det) = proved_max(4, G3, n, b)
            print(f'n={n:2d}  built {b:3d}  proved at most {m:3d}  '
                  f'{"TIGHT" if m == b else f"gap {m - b}":7s}  ({how}; splits allowed: {det["splits"]})',
                  flush=True)


if __name__ == '__main__':
    main()