"""Original exact finite-path calculation; no learned policy, gridworld benchmark, or empirical convergence claim.
Run with Python 3.10+; standard library only. Writes results.json in the current directory.
"""
gamma = 0.9
potential = {"S": -2.0, "M": -1.0, "G": 0.0}
paths = {
    "direct": [("S", "G", 1.0)],
    "detour": [("S", "M", 0.0), ("M", "G", 1.0)],
}

def path_return(path, mode):
    total = 0.0
    for step, (state, next_state, reward) in enumerate(path):
        if mode == "potential":
            reward += gamma * potential[next_state] - potential[state]
        elif mode == "bonus" and next_state == "M":
            reward += 2.0
        total += gamma ** step * reward
    return total

returns = {}
for name, path in paths.items():
    values = [path_return(path, mode) for mode in ("base", "potential", "bonus")]
    returns[name] = values
    assert abs((values[1] - values[0]) - 2.0) < 1e-12
    print(f"{name}: base={values[0]:.1f}, potential={values[1]:.1f}, bonus={values[2]:.1f}")
assert returns["direct"][0] > returns["detour"][0]
assert returns["direct"][1] > returns["detour"][1]
assert returns["direct"][2] < returns["detour"][2]

# Extended experiment: deterministic parameter sweep and invariants.
loop_sweep=[]
for loops in range(13):
    path=[('S','M',0.)]+[('M','M',0.)]*loops+[('M','G',1.)]
    base=path_return(path,'base'); shaped=path_return(path,'potential');bonus=path_return(path,'bonus')
    naive=0.
    for t,(s,n,r) in enumerate(path): naive+=gamma**t*(r+potential[n]-potential[s])
    offset=shaped-base
    assert abs(offset-2)<1e-12
    loop_sweep.append(dict(extra_loops=loops,transitions=len(path),base=base,potential=shaped,
                           arbitrary_bonus=bonus,correct_offset=offset,undiscounted_difference_offset=naive-base))
# Complete-path versus truncated-path boundary term.
truncated=[('S','M',0.)]
truncated_offset=path_return(truncated,'potential')-path_return(truncated,'base')
assert abs(truncated_offset-(2+gamma*(-1)))<1e-12
assert loop_sweep[-1]['arbitrary_bonus']>loop_sweep[0]['arbitrary_bonus']
data=dict(loop_sweep=loop_sweep,direct={mode:path_return(paths['direct'],mode) for mode in ['base','potential','bonus']},
          truncated_offset=truncated_offset,gamma=gamma,potential=potential)
print(f"Potential offset for every complete looping path: {loop_sweep[0]['correct_offset']:.1f}")
print(f"Twelve self-loops: base={loop_sweep[-1]['base']:.6f}, arbitrary bonus={loop_sweep[-1]['arbitrary_bonus']:.6f}")
print(f"Stopping after S→M leaves boundary term: shaped-minus-base={truncated_offset:.1f}")

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