"""Exact enumeration of a one-step constrained decision and fixed-penalty alternatives; not an implementation of CPO or a real-world safety controller.
Run with Python 3.10+; standard library only. Writes results.json in the current directory.
"""
budget = 0.6
probabilities = [i / 100 for i in range(101)]

def outcome(p):
    reward = 4.0 + 6.0 * p
    cost = 3.0 * p
    return reward, cost

feasible = [p for p in probabilities if outcome(p)[1] <= budget + 1e-12]
best_p = max(feasible, key=lambda p: outcome(p)[0])
reward, cost = outcome(best_p)
assert abs(best_p - 0.2) < 1e-12
assert abs(reward - 5.2) < 1e-12
print(f"Constrained: p={best_p:.2f}, reward={reward:.2f}, cost={cost:.2f}")

for penalty in [0.5, 1.5, 2.0, 3.0]:
    scores = [outcome(p)[0] - penalty * outcome(p)[1] for p in probabilities]
    optimum = max(scores)
    maximizers = [p for p, score in zip(probabilities, scores)
                  if abs(score - optimum) < 1e-12]
    print(f"lambda={penalty:.1f}: optimal p range "
          f"[{min(maximizers):.2f}, {max(maximizers):.2f}]")

assert outcome(1.0)[1] > budget
print(f"At constrained optimum, cost-3 episodes occur with probability {best_p:.2f}")

# Extended experiment: deterministic parameter sweep and invariants.
frontier=[]
for i in range(31):
    limit=i/10;p=min(1.,limit/3);reward,cost=outcome(p)
    assert cost<=limit+1e-12
    frontier.append(dict(budget=limit,optimal_p=p,reward=reward,cost=cost,probability_cost_exceeds_budget=p if limit<3 else 0.))
penalties=[]
for lam in [0.,.5,1.,1.5,1.9,2.,2.1,2.5,3.,4.]:
    best_min,best_max=(1.,1.) if lam<2 else ((0.,0.) if lam>2 else (0.,1.))
    penalties.append(dict(penalty=lam,optimal_p_min=best_min,optimal_p_max=best_max))
data=dict(budget_frontier=frontier,penalty_sweep=penalties,selected_budget=.6,
          episode_cost_distribution=[dict(cost=0,probability=.8),dict(cost=3,probability=.2)])
assert abs(frontier[6]['optimal_p']-.2)<1e-12
print('Budget sweep: each extra cost unit buys 2 reward units until budget 3.')
print('At budget 0.6, the expected-cost-optimal policy still exceeds that budget in 20% of episodes.')

import json
from pathlib import Path
Path("results.json").write_text(json.dumps(data,indent=2,allow_nan=False)+"\n",encoding="utf-8")