"""Why the policy-gradient loss works, on the smallest possible example.

Run: python3 site/ml/foundations/examples/policy_gradient.py

A policy picks a number x from a bell curve centred at mu (spread 1). The reward is +1 when x lands
in the window [1, 2.5] and -1 everywhere else. Training should move mu into the window, using only
samples and their rewards: no formula for the reward, no derivative of it.
"""

import math
import random

LOW, HIGH = 1.0, 2.5        # the rewarded window
SIGMA = 1.0                 # the bell curve's spread (fixed)


def reward(x, offset=0.0):
    """+1 inside the window, -1 outside. `offset` adds the same amount to every reward."""
    return (1.0 if LOW <= x <= HIGH else -1.0) + offset


def push(x, mu):
    """Which way to move mu to make x more likely: the slope of log p(x) with respect to mu."""
    return (x - mu) / SIGMA ** 2


def estimate(mu, samples, baseline=False, offset=0.0):
    """The score-function estimate of how expected reward changes with mu."""
    rewards = [reward(x, offset) for x in samples]
    b = sum(rewards) / len(rewards) if baseline else 0.0     # subtract the batch's average reward
    return sum((r - b) * push(x, mu) for x, r in zip(samples, rewards)) / len(samples)


def exact_slope(mu):
    """The true slope of expected reward, from the bell curve's formula (only for checking)."""
    pdf = lambda z: math.exp(-z * z / 2) / math.sqrt(2 * math.pi)
    return 2 / SIGMA * (pdf((LOW - mu) / SIGMA) - pdf((HIGH - mu) / SIGMA))


def train(steps=25, n=20, lr=0.5, baseline=False, offset=0.0, seed=0):
    rng = random.Random(seed)
    mu, history = 0.0, []
    for _ in range(steps):
        samples = [rng.gauss(mu, SIGMA) for _ in range(n)]
        g = estimate(mu, samples, baseline, offset)
        history.append({"mu": mu, "samples": samples, "rewards": [reward(x, offset) for x in samples], "estimate": g})
        mu += lr * g                                           # step uphill on expected reward
    return mu, history


def spread(mu=0.0, n=20, trials=2000, baseline=False, offset=0.0, seed=1):
    """How much single-batch estimates scatter around the true slope."""
    rng = random.Random(seed)
    values = [estimate(mu, [rng.gauss(mu, SIGMA) for _ in range(n)], baseline, offset) for _ in range(trials)]
    mean = sum(values) / trials
    return mean, math.sqrt(sum((v - mean) ** 2 for v in values) / trials)


if __name__ == "__main__":
    print(f"true slope at mu = 0: {exact_slope(0.0):.3f}")
    for offset in (0, 20):
        for b in (False, True):
            mean, sd = spread(baseline=b, offset=offset)
            mu, _ = train(baseline=b, offset=offset)
            print(f"offset {offset:2d}, {'baseline' if b else 'raw     '}: estimates average {mean:.3f}, "
                  f"spread {sd:.3f}; mu after 25 steps {mu:.2f}")
    gamma = 0.9
    print("discounted credit for the 7 moves of one drop, reward at the end:",
          ", ".join(f"{gamma ** (6 - t):.2f}" for t in range(7)))