#!/usr/bin/env python3
# SPDX-License-Identifier: CC-BY-4.0
# Creator: reinforcement-learning-editor; source: AI Encyclopedia CN, Chapter 13.
"""Deterministically generate/check Chapter 13 episode and update teaching data."""
import argparse, json, math
from pathlib import Path

ROUND = 9

def rnd(value):
    if not math.isfinite(value):
        raise ValueError('all numeric inputs and results must be finite')
    return round(value, ROUND)

def discounted_returns(rewards, gamma):
    """Return G_t; retain the unrounded binary64 recurrence and round only outputs."""
    running = 0.0
    returns = [0.0] * len(rewards)
    for index in range(len(rewards) - 1, -1, -1):
        running = rewards[index] + gamma * running
        returns[index] = rnd(running)
    return returns

def q_update(initial_q, alpha, reward, gamma, next_value, terminal=False):
    bootstrap = 0.0 if terminal else next_value
    target = reward + gamma * bootstrap
    td_error = target - initial_q
    updated_q = initial_q + alpha * td_error
    return {'initial_q': initial_q, 'alpha': alpha, 'reward': reward, 'gamma': gamma,
            'terminal': terminal, 'bootstrap_value': bootstrap, 'target': rnd(target),
            'td_error': rnd(td_error), 'updated_q': rnd(updated_q)}

def softmax_two(logits):
    high = max(logits)
    exps = [math.exp(x-high) for x in logits]
    total = sum(exps)
    return [rnd(x/total) for x in exps]

def payload():
    gamma, rewards = 0.9, [-1, 2, 5]
    q_learning = q_update(1.0, 0.5, 2, gamma, 4.0)
    terminal = q_update(1.0, 0.5, 5, gamma, 999.0, terminal=True)
    sarsa = q_update(1.0, 0.5, 2, gamma, 2.0)
    action_values = {'advance': 3.0, 'safe': 2.0, 'greedy': 4.0}
    action_order, logits_before = ['left', 'right'], [0.0, 0.0]
    probabilities_before = softmax_two(logits_before)
    sampled_index, sampled_return, policy_alpha = action_order.index('right'), 2.0, 0.1
    score_gradient = [rnd((1.0 if i == sampled_index else 0.0) - p) for i,p in enumerate(probabilities_before)]
    gradient_estimate = [rnd(sampled_return*x) for x in score_gradient]
    logits_after = [rnd(x + policy_alpha*g) for x,g in zip(logits_before, gradient_estimate)]
    return {
      'schema': 2, 'creator': 'reinforcement-learning-editor', 'license': 'CC BY 4.0',
      'license_url': 'https://creativecommons.org/licenses/by/4.0/',
      'source': 'AI Encyclopedia CN, Chapter 13',
      'provenance': 'Original fictional episode and TD/policy-update fixture; no people or deployment data.',
      'conventions': {
        'reward_indexing': 'trajectory rewards are [R1, R2, R3]; returns are [G0, G1, G2], so returns[t] begins with rewards[t] = R_(t+1)',
        'return': 'G_t = R_(t+1) + gamma R_(t+2) + ...; no reward follows the terminal transition',
        'q_learning': 'target = r + gamma max over every action in available_actions[s_next] of Q(s_next,a) for nonterminal transitions',
        'terminal_bootstrap': 'zero; the terminal available-action set is empty and no maximum is taken',
        'rounding': 'IEEE-754 binary64 calculations remain unrounded; only emitted derived fields are rounded to 9 decimal places with Python round (ties-to-even)',
        'trajectory_semantics': 'trajectory lists actions actually sampled in this episode; available_actions separately and exhaustively defines the fixture action set at every state'
      },
      'episode_fixture': {
        'states': ['start','bridge','goal','terminal'],
        'available_actions': {'start':['advance'], 'bridge':['advance'], 'goal':['advance','safe','greedy'], 'terminal':[]},
        'gamma': gamma,
        'trajectory': [
          {'state':'start','action':'advance','reward':-1,'next_state':'bridge','terminal':False},
          {'state':'bridge','action':'advance','reward':2,'next_state':'goal','terminal':False},
          {'state':'goal','action':'advance','reward':5,'next_state':'terminal','terminal':True}],
        'rewards': rewards, 'returns': discounted_returns(rewards,gamma)},
      'q_value_snapshot': {'goal': action_values},
      'q_learning_update': {**q_learning, 'state_action':['bridge','advance'],
        'next_state':'goal', 'next_state_action_values':action_values,
        'selected_bootstrap_action':'greedy'},
      'terminal_update': {**terminal, 'state_action':['goal','advance'],
        'next_state':'terminal', 'stored_next_value_ignored':999.0,
        'note':'Bootstrap is zero because this transition is terminal; no max over an empty action set is evaluated.'},
      'policy_contrast': {'description':'Counterfactual branch for the bridge transition, separate from the recorded advance trajectory: if SARSA next samples safe it uses value 2, while Q-learning uses the maximum 4.',
        'relation_to_episode_trajectory':'separate counterfactual update; does not assert that the recorded episode sampled safe',
        'sarsa_behavior_action':'safe','sarsa_update':sarsa,'q_learning_updated_q':q_learning['updated_q']},
      'policy_gradient_update': {
        'environment':'one nonterminal state with actions left/right; the sampled right action ends the episode with return 2',
        'parameterization':'two-logit softmax', 'action_order':action_order, 'alpha':policy_alpha,
        'sampled_action':action_order[sampled_index],'return':sampled_return,'gamma_power':1.0,
        'logits_before':logits_before,'probabilities_before':probabilities_before,
        'score_gradient_log_pi_sampled_action':score_gradient,
        'gradient_estimate':gradient_estimate,'logits_after':logits_after,
        'probabilities_after':softmax_two(logits_after)}
    }

def encode(): return (json.dumps(payload(),ensure_ascii=False,indent=2,sort_keys=True)+'\n').encode()
def default_output():
    here=Path(__file__).resolve()
    return here.parent if here.parent.name=='data' and here.parent.parent.name=='public' else here.parents[1]/'public/data'
def main():
    parser=argparse.ArgumentParser(); mode=parser.add_mutually_exclusive_group(required=True)
    mode.add_argument('--check',action='store_true'); mode.add_argument('--write',action='store_true')
    parser.add_argument('--output-dir',type=Path); args=parser.parse_args(); output=(args.output_dir or default_output())/'ch13-tiny-mdp.json'; expected=encode()
    if args.check:
        if not output.is_file() or output.read_bytes()!=expected: raise SystemExit(f'FAIL: {output} missing or differs')
        print(f'PASS: {output} deterministic; returns, 3 TD updates, complete action sets/Q snapshot, and numerical policy update verified')
    else:
        output.parent.mkdir(parents=True,exist_ok=True); output.write_bytes(expected); print(f'WROTE: {output}')
if __name__=='__main__': main()