"""Observe original simulator returns without changing its events or RNG calls."""
import ast
import json
from pathlib import Path

tree = ast.parse(Path('/bundle/code/run.py').read_text())
function = next(n for n in tree.body if isinstance(n, ast.FunctionDef) and n.name == 'simulate_stochastic')
for node in ast.walk(function):
    if isinstance(node, ast.Return):
        assert isinstance(node.value, ast.Dict)
        for label, expression in [('final_N', 'N_S + N_R'), ('final_R', 'N_R'), ('elapsed', 't')]:
            node.value.keys.append(ast.Constant(label))
            node.value.values.append(ast.parse(expression, mode='eval').body)
ast.fix_missing_locations(tree)
namespace = {'__name__': 'audited_simulator'}
exec(compile(tree, '/bundle/code/run.py', 'exec'), namespace)
sim = namespace['simulate_stochastic']
output = {}
for regimen, start in [('continuous', 10000), ('adaptive', 20000)]:
    rows = [sim(f0=0.0, mu=8e-5, regimen=regimen, seed=start+i) for i in range(500)]
    progressed = [r for r in rows if r['final_N'] >= 1200 and r['ttp'] < 300]
    extinct = [r for r in rows if r['cure']]
    censored = [r for r in rows if not r['cure'] and r['ttp'] == 300]
    output[regimen] = {
        'trials': len(rows), 'progression_observed': len(progressed),
        'extinction': len(extinct), 'horizon_without_observed_progression': len(censored),
        'censored_below_progression_threshold': sum(r['final_N'] < 1200 for r in censored),
        'censored_without_resistant_cells_at_end': sum(r['final_R'] == 0 for r in censored),
        'reported_mean_noncure_ttp': round(sum(r['ttp'] for r in rows if not r['cure']) / (len(rows)-len(extinct)), 1),
        'reported_mean_dose': round(sum(r['dose'] for r in rows)/len(rows), 1),
        'cure_pct': 100*len(extinct)/len(rows),
        'seed_first_censored': next((start+i for i,r in enumerate(rows) if r in censored), None),
    }
print(json.dumps(output, indent=2))