"""WebArena’s evaluation lesson: check the final state
Deterministic, standard-library experiment. RESULTS contains the full data.
"""
TEMPLATES = [
    ('valid', {'blue-mug': 2}, 'Added successfully', '/cart', True),
    ('quiet', {'blue-mug': 2}, '', '/catalog', True),
    ('color', {'red-mug': 2}, 'Added successfully', '/cart', False),
    ('count', {'blue-mug': 1}, 'Added successfully', '/cart', False),
    ('extra', {'blue-mug': 2, 'red-mug': 1}, 'Added successfully', '/cart', False),
]
cases = [dict(name=f'{name}-{i}', cart=cart.copy(), message=msg,
              url=url, label=label)
         for i in range(4) for name, cart, msg, url, label in TEMPLATES]

def confusion(check):
    result = dict(TP=0, FP=0, FN=0, TN=0)
    for case in cases:
        predicted, actual = check(case), case['label']
        key = ('TP' if actual else 'FP') if predicted else ('FN' if actual else 'TN')
        result[key] += 1
    return result

checks = {
    'message': lambda c: 'successfully' in c['message'],
    'url': lambda c: c['url'] == '/cart',
    'state': lambda c: c['cart'] == {'blue-mug': 2},
}
for name, check in checks.items():
    print(name, confusion(check))
assert len(cases) == 20
assert confusion(checks['state']) == dict(TP=8, FP=0, FN=0, TN=12)


rows = []
for case in cases:
    rows.append(dict(case, predictions={name:check(case) for name,check in checks.items()}))
weak = lambda c: c['cart'].get('blue-mug') == 2
assert confusion(weak) == dict(TP=8,FP=4,FN=0,TN=8)
RESULTS=dict(cases=rows, confusion={name:confusion(check) for name,check in checks.items()}, positive_only=confusion(weak))
print('Positive-only predicate',confusion(weak))
