module Learning.Learner.MomentNormalized

import Std.Float
import Data.Float32Bits

-- A CPU (Std.Float F32) learner that moves each parameter by its bias-corrected
-- first moment divided by the root of its bias-corrected second moment:
--
--   m' = beta1 m + (1 - beta1) g                       firstMomentValue
--   v' = beta2 v + (1 - beta2) (g g)                   secondMomentValue
--   p' = p beta                                        momentDecayPowerNext
--   m^ = m' / (1 - p1'),  v^ = v' / (1 - p2')          correctMomentBias
--   w' = w - (rate m^) / (sqrt v^ + epsilon)            normalizedStepValue
--
-- This is the update known as Adam (Kingma & Ba 2015, Algorithm 1) without
-- weight decay. It differs from Learning.Checked.AdamW, the SM86 owner, on
-- purpose: there the moments use fused multiply-adds, the reciprocal square
-- root is the MUFU approximation and the bias corrections are folded into a
-- precomputed step size; here every operation is a correctly rounded binary32
-- add, subtract, multiply, divide or sqrt in the order written above, so a CPU
-- run is reproducible bit for bit by any IEEE-754 implementation that keeps
-- that order (systems/manfred checks it against a Nim reference).
--
-- The complements 1 - beta are derived here in F32 rather than passed in:
-- 1 - 0.9 rounds to 0.100000024, not to the literal 0.1, and a learner that
-- took the complement as a separate hyperparameter could silently disagree
-- with its own beta. The bias powers beta^t are carried by the caller as
-- running products (momentDecayPowerNext), which needs no pow.

family MomentNormalizedHyper : Type 0
constructor MomentNormalizedHyperValue
field unrestricted momentNormalizedRate : F32
field unrestricted momentNormalizedBeta1 : F32
field unrestricted momentNormalizedBeta2 : F32
field unrestricted momentNormalizedEpsilon : F32

end-family

-- One parameter with its two moments, before or after an update.
family MomentNormalizedElement : Type 0
constructor MomentNormalizedElementValue
field unrestricted momentNormalizedWeight : F32
field unrestricted momentNormalizedFirst : F32
field unrestricted momentNormalizedSecond : F32

end-family

-- m' = beta m + (1 - beta) g
def firstMomentValue =
  (lambda unrestricted beta : F32 .
    (lambda unrestricted moment : F32 .
      (lambda unrestricted gradient : F32 .
        (f32-add
          (f32-multiply beta moment)
          (f32-multiply (f32-subtract stdF32One beta) gradient)))))

-- v' = beta v + (1 - beta) (g g)
def secondMomentValue =
  (lambda unrestricted beta : F32 .
    (lambda unrestricted moment : F32 .
      (lambda unrestricted gradient : F32 .
        (f32-add
          (f32-multiply beta moment)
          (f32-multiply (f32-subtract stdF32One beta) (f32-multiply gradient gradient))))))

-- beta^(t+1) from beta^t; start the product at one.
def momentDecayPowerNext =
  (lambda unrestricted beta : F32 .
    (lambda unrestricted power : F32 .
      (f32-multiply power beta)))

-- m / (1 - beta^t), with beta^t already advanced to this update.
def correctMomentBias =
  (lambda unrestricted power : F32 .
    (lambda unrestricted moment : F32 .
      (f32-divide moment (f32-subtract stdF32One power))))

-- w - (rate m^) / (sqrt v^ + epsilon)
def normalizedStepValue =
  (lambda unrestricted rate : F32 .
    (lambda unrestricted epsilon : F32 .
      (lambda unrestricted weight : F32 .
        (lambda unrestricted first : F32 .
          (lambda unrestricted second : F32 .
            (f32-subtract
              weight
              (f32-divide
                (f32-multiply rate first)
                (f32-add (f32-sqrt second) epsilon))))))))

-- One parameter's update, given the bias powers already advanced to this step
-- (power1 = beta1^t, power2 = beta2^t for the t-th update, t >= 1).
def momentNormalizedUpdate =
  (lambda unrestricted hyper : (family MomentNormalizedHyper) .
    (lambda unrestricted power1 : F32 .
      (lambda unrestricted power2 : F32 .
        (lambda unrestricted element : (family MomentNormalizedElement) .
          (lambda unrestricted gradient : F32 .
            (eliminate MomentNormalizedHyper
              (lambda unrestricted current : (family MomentNormalizedHyper) . (family MomentNormalizedElement))
              hyper
              (branch MomentNormalizedHyperValue rate beta1 beta2 epsilon .
                (eliminate MomentNormalizedElement
                  (lambda unrestricted current : (family MomentNormalizedElement) . (family MomentNormalizedElement))
                  element
                  (branch MomentNormalizedElementValue weight first second .
                    (let unrestricted first1 = (firstMomentValue beta1 first gradient) in
                      (let unrestricted second1 = (secondMomentValue beta2 second gradient) in
                        (constructor MomentNormalizedElement MomentNormalizedElementValue
                          (normalizedStepValue rate epsilon weight
                            (correctMomentBias power1 first1)
                            (correctMomentBias power2 second1))
                          first1
                          second1))))))))))))

def momentNormalizedWeightOf =
  (lambda unrestricted element : (family MomentNormalizedElement) .
    (eliminate MomentNormalizedElement (lambda unrestricted current : (family MomentNormalizedElement) . F32) element
      (branch MomentNormalizedElementValue weight first second . weight)))

def momentNormalizedFirstOf =
  (lambda unrestricted element : (family MomentNormalizedElement) .
    (eliminate MomentNormalizedElement (lambda unrestricted current : (family MomentNormalizedElement) . F32) element
      (branch MomentNormalizedElementValue weight first second . first)))

def momentNormalizedSecondOf =
  (lambda unrestricted element : (family MomentNormalizedElement) .
    (eliminate MomentNormalizedElement (lambda unrestricted current : (family MomentNormalizedElement) . F32) element
      (branch MomentNormalizedElementValue weight first second . second)))

def momentNormalizedBeta1Of =
  (lambda unrestricted hyper : (family MomentNormalizedHyper) .
    (eliminate MomentNormalizedHyper (lambda unrestricted current : (family MomentNormalizedHyper) . F32) hyper
      (branch MomentNormalizedHyperValue rate beta1 beta2 epsilon . beta1)))

def momentNormalizedBeta2Of =
  (lambda unrestricted hyper : (family MomentNormalizedHyper) .
    (eliminate MomentNormalizedHyper (lambda unrestricted current : (family MomentNormalizedHyper) . F32) hyper
      (branch MomentNormalizedHyperValue rate beta1 beta2 epsilon . beta2)))