"""
Lecture 09 - Advanced Optimizers: Momentum, AdaGrad, RMSProp, Adam

Reproduces the exact worked comparison from the lecture page: starting from
w=0 on L(w)=(w-3)^2 (same loss as Lecture 7), compute ONE update step under
plain SGD, Momentum, AdaGrad, and Adam, using the hypothetical prior-state
values given in the lecture (v_{t-1}=0.5 for Momentum, G_{t-1}=0.2 for
AdaGrad, m_{t-1}=0.3 & v_{t-1}=0.1 & t=2 for Adam). Also includes small,
runnable, general-purpose optimizer step functions (incl. RMSProp) usable
in a real training loop.

Run: python lecture-09-optimizers.py
"""
import numpy as np


def grad_L(w):
    return 2 * (w - 3)   # L(w) = (w-3)^2, same loss as Lecture 7


# ---- general-purpose one-step optimizer functions ----
def sgd_step(w, grad, lr):
    return w - lr * grad


def momentum_step(w, grad, lr, v_prev, beta=0.9):
    v = beta * v_prev - lr * grad
    return w + v, v


def adagrad_step(w, grad, lr, G_prev, eps=1e-8):
    G = G_prev + grad ** 2
    w_new = w - (lr / np.sqrt(G + eps)) * grad
    return w_new, G


def rmsprop_step(w, grad, lr, G_prev, beta=0.9, eps=1e-8):
    G = beta * G_prev + (1 - beta) * grad ** 2
    w_new = w - (lr / np.sqrt(G + eps)) * grad
    return w_new, G


def adam_step(w, grad, lr, m_prev, v_prev, t, beta1=0.9, beta2=0.999, eps=1e-8):
    m = beta1 * m_prev + (1 - beta1) * grad
    v = beta2 * v_prev + (1 - beta2) * grad ** 2
    m_hat = m / (1 - beta1 ** t)
    v_hat = v / (1 - beta2 ** t)
    w_new = w - lr * m_hat / (np.sqrt(v_hat) + eps)
    return w_new, m, v


if __name__ == "__main__":
    w0, lr = 0.0, 0.1
    g = grad_L(w0)
    print(f"Common starting point: w=0, eta=0.1, gradient = {g:.4f}\n")

    # ---- Plain SGD ----
    w_sgd = sgd_step(w0, g, lr)
    print(f"Plain SGD   -> w_new = {w_sgd:.4f}")

    # ---- Momentum (v_{t-1} = 0.5) ----
    w_mom, v = momentum_step(w0, g, lr, v_prev=0.5, beta=0.9)
    print(f"Momentum    -> w_new = {w_mom:.4f}   (v_t = {v:.4f})")

    # ---- AdaGrad (G_{t-1} = 0.2) ----
    w_ada, G = adagrad_step(w0, g, lr, G_prev=0.2)
    print(f"AdaGrad     -> w_new = {w_ada:.4f}   (G_t = {G:.4f})")

    # ---- Adam (m_{t-1}=0.3, v_{t-1}=0.1, t=2) ----
    w_adam, m, v2 = adam_step(w0, g, lr, m_prev=0.3, v_prev=0.1, t=2)
    print(f"Adam        -> w_new = {w_adam:.4f}   (m_t = {m:.4f}, v_t = {v2:.4f})")

    print("\n(These four numbers depend on the assumed prior-state values above,")
    print(" not on any inherent ranking of the optimizers -- see the lecture note.)")

    # ---- bonus: RMSProp step, for completeness (not in the headline comparison) ----
    w_rms, G_rms = rmsprop_step(w0, g, lr, G_prev=0.2, beta=0.9)
    print(f"\n[RMSProp with the same G_prev=0.2 for reference] -> w_new = {w_rms:.4f}")