#!/usr/bin/env python3
"""Bergwandeling: vier manieren om een smal dal af te lopen.

Het rekenvoorbeeld uit het artikel "Niet-lineair optimaliseren" op
https://fermat360.nl/artikelen/niet-lineair-optimaliseren. De figuren en de
getallen over de paden in dat artikel komen uit dit script.

Alleen de standaardbibliotheek van Python 3 is nodig:

    python3 bergwandeling.py          # tabel met de uitkomsten
    python3 bergwandeling.py --json   # alle paden, als JSON
"""

import json
import math
import random
import sys

# Het landschap is een smal dal: h(u, w) = 1/2 * (u^2 + STEILHEID * w^2), in
# de coordinaten van het dal zelf. Langs u loopt de bodem flauw af, langs w
# zijn de wanden steil. Het laagste punt ligt in (0, 0), op hoogte 0.
STEILHEID = 50
STAPPEN = 60
RUIS = 1.0  # standaardafwijking van de ruis op de helling: de "minibatch"
ZAAD = 7  # vaste startwaarde, zodat elke keer dezelfde ruis volgt
START = (-9.0, 1.6)  # ver weg op de flauwe bodem, een eind tegen de wand op

# Elke methode kiest haar staplengte uit dezelfde reeks: 10^(k/8), k = -24 ... 4.
STAPLENGTES = [10 ** (k / 8) for k in range(-24, 5)]


def draai(punt, hoek):
    """Draait een punt, of een helling, over `hoek` graden."""
    c, s = math.cos(math.radians(hoek)), math.sin(math.radians(hoek))
    x, y = punt
    return (c * x - s * y, s * x + c * y)


def hoogte(punt, hoek):
    """Het verlies op deze plek, in een dal dat `hoek` graden gedraaid ligt."""
    u, w = draai(punt, -hoek)
    return 0.5 * (u * u + STEILHEID * w * w)


def helling(punt, hoek):
    """De gradient: de richting waarin het hier het steilst omhoog gaat."""
    u, w = draai(punt, -hoek)
    return draai((u, STEILHEID * w), hoek)


def ruisreeks():
    """Normaal verdeelde ruis, met Box-Muller op random().

    random() geeft bij hetzelfde zaad in elke Python-versie dezelfde reeks;
    random.gauss() belooft dat niet.
    """
    toeval = random.Random(ZAAD)
    reeks = []
    for _ in range(STAPPEN):
        straal = math.sqrt(-2 * math.log(1 - toeval.random()))
        hoek = 2 * math.pi * toeval.random()
        reeks.append((RUIS * straal * math.cos(hoek), RUIS * straal * math.sin(hoek)))
    return reeks


# De vier methoden. Elke methode krijgt een staplengte eta en geeft een
# functie terug die uit de helling g en het stapnummer t de stap berekent.
# Wat een methode tussen twee stappen onthoudt, bewaart die functie zelf.


def sgd(eta):
    def stap(g, t):
        return tuple(eta * gi for gi in g)

    return stap


def momentum(eta, beta=0.9):
    v = (0.0, 0.0)

    def stap(g, t):
        nonlocal v
        v = tuple(beta * vi + gi for vi, gi in zip(v, g))  # vaart houden
        return tuple(eta * vi for vi in v)

    return stap


def adagrad(eta, eps=1e-8):
    som = (0.0, 0.0)

    def stap(g, t):
        nonlocal som
        som = tuple(si + gi * gi for si, gi in zip(som, g))  # groeit alleen
        return tuple(eta * gi / (math.sqrt(si) + eps) for gi, si in zip(g, som))

    return stap


def adam(eta, beta1=0.9, beta2=0.999, eps=1e-8):
    m = (0.0, 0.0)
    v = (0.0, 0.0)

    def stap(g, t):
        nonlocal m, v
        m = tuple(beta1 * mi + (1 - beta1) * gi for mi, gi in zip(m, g))
        v = tuple(beta2 * vi + (1 - beta2) * gi * gi for vi, gi in zip(v, g))
        # Biascorrectie: m en v beginnen op nul en zijn eerst te klein.
        m_hat = tuple(mi / (1 - beta1**t) for mi in m)
        v_hat = tuple(vi / (1 - beta2**t) for vi in v)
        return tuple(eta * mh / (math.sqrt(vh) + eps) for mh, vh in zip(m_hat, v_hat))

    return stap


METHODEN = {"SGD": sgd, "Momentum": momentum, "AdaGrad": adagrad, "Adam": adam}


def wandel(methode, eta, hoek):
    """Zet STAPPEN stappen bergaf; geeft het hele pad terug, met het startpunt."""
    stap = methode(eta)
    ruis = [draai(r, hoek) for r in ruisreeks()]  # de ruis draait mee met het dal
    p = draai(START, hoek)
    pad = [p]
    for t in range(1, STAPPEN + 1):
        g = helling(p, hoek)
        g = (g[0] + ruis[t - 1][0], g[1] + ruis[t - 1][1])  # de helling, geschat
        dx, dy = stap(g, t)
        p = (p[0] - dx, p[1] - dy)  # nieuwe plek = oude plek - stap
        pad.append(p)
    return pad


def beste_wandeling(naam, hoek):
    """Probeert elke staplengte en houdt de wandeling die het laagst eindigt.

    Laagst eindigen = de gemiddelde hoogte over de laatste tien stappen.
    """
    beste = None
    for eta in STAPLENGTES:
        pad = wandel(METHODEN[naam], eta, hoek)
        hoogtes = [hoogte(p, hoek) for p in pad]
        eind = sum(hoogtes[-10:]) / 10
        if math.isfinite(eind) and (beste is None or eind < beste["eindhoogte"]):
            onder_een = next((i for i, h in enumerate(hoogtes) if h < 1), None)
            beste = {
                "methode": naam,
                "staplengte": eta,
                "eindhoogte": eind,
                "onder_een": onder_een,
                "pad": pad,
            }
    return beste


def main():
    dalen = {"langs de as": 0, "scheef": 45}
    uitkomst = {dal: [beste_wandeling(m, hoek) for m in METHODEN] for dal, hoek in dalen.items()}
    if "--json" in sys.argv:
        json.dump(
            {"steilheid": STEILHEID, "stappen": STAPPEN, "ruis": RUIS, "dalen": dalen, "wandelingen": uitkomst},
            sys.stdout,
            indent=1,
        )
        print()
        return
    for dal, wandelingen in uitkomst.items():
        print(f"Dal {dal} ({dalen[dal]} graden), {STAPPEN} stappen")
        print(f"  {'methode':<10}{'staplengte':>11}{'hoogte':>9}{'onder 1':>9}")
        for w in wandelingen:
            onder = "niet" if w["onder_een"] is None else w["onder_een"]
            print(f"  {w['methode']:<10}{w['staplengte']:>11.4f}{w['eindhoogte']:>9.2f}{onder:>9}")
        print()


if __name__ == "__main__":
    main()