"""Exact finite-action entropy optimization for hand-specified values; not SAC training or a continuous-control benchmark.
Run with Python 3.10+; standard library only. Writes results.json in the current directory.
"""
import math

values = [1.0, 0.8, 0.0]

def distribution(q, alpha):
    shifted = [(value - max(q)) / alpha for value in q]
    weights = [math.exp(value) for value in shifted]
    return [value / sum(weights) for value in weights]

def score(probabilities, alpha):
    reward = sum(p * q for p, q in zip(probabilities, values))
    entropy = -sum(p * math.log(p) for p in probabilities if p)
    return reward, entropy, reward + alpha * entropy

for alpha in [0.05, 0.2, 1.0]:
    probabilities = distribution(values, alpha)
    reward, entropy, objective = score(probabilities, alpha)
    assert abs(sum(probabilities) - 1) < 1e-12
    assert objective >= score([1.0, 0.0, 0.0], alpha)[2]
    assert objective >= score([1 / 3] * 3, alpha)[2]
    print(f"alpha={alpha:.2f} p=" + str([round(p, 4) for p in probabilities]))
    print(f"  reward={reward:.4f} entropy={entropy:.4f} objective={objective:.4f}")

original = distribution(values, 0.2)
rescaled = distribution([10 * q for q in values], 2.0)
assert all(abs(a - b) < 1e-12 for a, b in zip(original, rescaled))
print("Reward/temperature rescaling check: passed")

# Extended experiment: deterministic parameter sweep and invariants.
sweep=[]
for i in range(41):
    alpha=0.025*(80**(i/40))
    p=distribution(values,alpha);reward,entropy,objective=score(p,alpha)
    sweep.append(dict(alpha=alpha,probabilities=p,reward=reward,entropy=entropy,objective=objective))
checks=[]
for alpha in [.05,.2,1.]:
    optimal=distribution(values,alpha); optimum=score(optimal,alpha)[2];max_error=0.;grid_best=-1e9
    for i in range(101):
        for j in range(101-i):
            p=[i/100,j/100,(100-i-j)/100]
            objective=score(p,alpha)[2]
            kl=sum(x*math.log(x/y) for x,y in zip(p,optimal) if x)
            max_error=max(max_error,abs((optimum-objective)-alpha*kl))
            assert objective<=optimum+1e-12
            grid_best=max(grid_best,objective)
    assert max_error<1e-12
    checks.append(dict(alpha=alpha,grid_points=5151,identity_max_error=max_error,grid_objective_gap=optimum-grid_best))
data=dict(values=values,temperature_sweep=sweep,simplex_checks=checks)
print('All 15,453 simplex candidates satisfy J(p*)−J(p) = alpha × KL(p || p*) within 1e−12.')

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