"""What a search budget cannot fix after early pruning
Deterministic, standard-library experiment. RESULTS contains the full data.
"""
import math

leaves = ((2, 3, 4), (5, 6, 7), (0, 8, 10))
heuristic_order = (0, 1, 2)

def search(mode, budget):
    allowed = heuristic_order[:{'greedy': 1, 'beam2': 2, 'uct': 3}[mode]]
    visits, totals = [0] * 3, [0.0] * 3
    best, used = None, 0
    for _ in range(budget):
        frontier = [a for a in allowed if visits[a] < len(leaves[a])]
        if not frontier:
            break
        if mode == 'uct':
            unseen = [a for a in frontier if visits[a] == 0]
            def score(a):
                bonus = math.sqrt(2 * math.log(max(1, used)) / visits[a])
                return totals[a] / visits[a] + bonus
            arm = unseen[0] if unseen else max(frontier, key=score)
        else:
            arm = min(frontier, key=lambda a: (visits[a], a))
        reward = leaves[arm][visits[arm]]
        visits[arm] += 1
        totals[arm] += reward
        used += 1
        best = reward if best is None else max(best, reward)
    return best, used

for budget in (2, 4, 6, 9):
    row = {mode: search(mode, budget) for mode in ('greedy', 'beam2', 'uct')}
    print('budget', budget, 'best_reward,used', row)
assert search('uct', 9) == (10, 9)
assert search('greedy', 9) == (4, 3)


sweep=[dict(budget=b,mode=m,best_reward=search(m,b)[0],used=search(m,b)[1]) for b in range(1,10) for m in ('greedy','beam2','uct')]
# Exhaustive three-branch allocation recovers all nine leaves regardless of ordering.
original_order=heuristic_order
permuted=[]
from itertools import permutations
for order in permutations(range(3)):
    heuristic_order=order
    permuted.append(dict(order=list(order),greedy=search('greedy',9)[0],beam2=search('beam2',9)[0],uct=search('uct',9)[0]))
heuristic_order=original_order
assert all(r['uct']==10 for r in permuted)
RESULTS=dict(leaves=leaves,heuristics=[9,6,2],sweep=sweep,order_sensitivity=permuted)
print('Heuristic order permutations:',len(permuted),'; full-budget UCT maximum always 10')