#!/usr/bin/env python3
"""Muon in het klein: de Newton-Schulz-stap op een 2x2-matrix.

Het rekenvoorbeeld uit het artikel "Niet-lineair optimaliseren" op
https://fermat360.nl/artikelen/niet-lineair-optimaliseren. Er gaat een stap
in met een sterke richting (10) en een zwakke (0,1); na vijf rondes zijn beide
richtingen ongeveer even sterk.

Alleen de standaardbibliotheek van Python 3 is nodig:

    python3 newton_schulz.py          # de sterktes na elke ronde
    python3 newton_schulz.py --json   # hetzelfde, als JSON
"""

import json
import math
import sys

# De coefficienten uit de Muon-implementatie van Keller Jordan e.a. (2024).
# Ze zijn gekozen op snelheid, niet op exactheid.
A, B, C = 3.4445, -4.7750, 2.0315
RONDES = 5


def maal(p, q):
    """Het product van twee 2x2-matrices."""
    return [[sum(p[i][k] * q[k][j] for k in range(2)) for j in range(2)] for i in range(2)]


def getransponeerd(p):
    return [[p[0][0], p[1][0]], [p[0][1], p[1][1]]]


def optellen(*termen):
    """Een som van matrices met een getal ervoor: optellen((2, P), (3, Q)) = 2P + 3Q."""
    return [[sum(c * m[i][j] for c, m in termen) for j in range(2)] for i in range(2)]


def draaiing(hoek):
    """Een draaiing over `hoek` radialen."""
    return [[math.cos(hoek), -math.sin(hoek)], [math.sin(hoek), math.cos(hoek)]]


def sterktes(p):
    """De singuliere waarden van een 2x2-matrix, groot naar klein."""
    q = maal(getransponeerd(p), p)
    spoor = q[0][0] + q[1][1]
    det = q[0][0] * q[1][1] - q[0][1] * q[1][0]
    d = math.sqrt(max(spoor * spoor / 4 - det, 0))
    return [math.sqrt(spoor / 2 + d), math.sqrt(max(spoor / 2 - d, 0))]


def newton_schulz(g, rondes=RONDES):
    """X <- A*X + (B*XX' + C*(XX')^2)*X, na delen door de Frobeniusnorm.

    Geeft alle tussenstanden terug, te beginnen bij de genormaliseerde stap.
    """
    norm = math.sqrt(sum(x * x for rij in g for x in rij))
    x = [[e / norm for e in rij] for rij in g]
    standen = [x]
    for _ in range(rondes):
        a = maal(x, getransponeerd(x))
        b = optellen((B, a), (C, maal(a, a)))
        x = optellen((A, x), (1, maal(b, x)))
        standen.append(x)
    return standen


# De stap voor een matrix: sterktes 10 en 0,1, in willekeurige richtingen gedraaid.
STAP = maal(maal(draaiing(0.6), [[10, 0], [0, 0.1]]), draaiing(-1.1))


def main():
    voor = sterktes(STAP)
    rondes = [sterktes(x) for x in newton_schulz(STAP)]
    if "--json" in sys.argv:
        json.dump({"coefficienten": [A, B, C], "voor": voor, "rondes": rondes}, sys.stdout, indent=1)
        print()
        return
    print(f"stap zelf         {voor[0]:8.4f}{voor[1]:8.4f}")
    for i, (sterk, zwak) in enumerate(rondes):
        label = "genormaliseerd" if i == 0 else f"na ronde {i}"
        print(f"{label:<18}{sterk:8.4f}{zwak:8.4f}")


if __name__ == "__main__":
    main()