module Bob.Learner

import Data.Bytes
import Data.Float32Bits
import Model.Architecture
import Representation.Schema
import Std.Float
import Std.Natural

family BobLearningSemantics : Type 0
constructor BobLearningSemanticsValue
field unrestricted bobLearningObjectiveIdentity : Bytes
field unrestricted bobLearningObjective : (family ModelObjective)
field unrestricted bobLearningLearnerIdentity : Bytes
field unrestricted bobLearningLearner : (family LearnerContract)

end-family

def bobObjectiveIdentity : Bytes =
  b"exact-full-256-cross-entropy"

def bobObjective : (family ModelObjective) =
  (constructor ModelObjective ModelExactVocabulary)

def bobLearnerIdentity : Bytes =
  b"adamw"

def bobLearner : (family LearnerContract) =
  (constructor LearnerContract LearnerAdamW)

def bobLearningSemantics : (family BobLearningSemantics) =
  (record
    BobLearningSemantics
    (bobLearningObjectiveIdentity = bobObjectiveIdentity)
    (bobLearningObjective = bobObjective)
    (bobLearningLearnerIdentity = bobLearnerIdentity)
    (bobLearningLearner = bobLearner))

-- AdamW hyperparameters.  Everything the optimizer kernel reads per step is
-- computed from these four values and the step number.
def bobLearningRate : F32 = 0.05
def bobBeta1 : F32 = 0.9
def bobBeta2 : F32 = 0.999
def bobEpsilon : F32 = 1e-8
def bobWeightDecay : F32 = 0.0

def bobF32One : F32 = 1.0

def bobF32Power =
  (lambda unrestricted base : F32 .
    (lambda unrestricted exponent : Nat .
      (nat-eliminate
        (lambda unrestricted current : Nat . F32)
        bobF32One
        (lambda unrestricted predecessor : Nat .
          (lambda unrestricted induction : F32 . (f32-multiply induction base)))
        exponent)))

-- The nine scalars the AdamW SM86 kernel reads, in its ABI order: gradient
-- scale, beta1, 1 - beta1, beta2, 1 - beta2, the decay multiplier
-- 1 - lr * weightDecay, the bias-corrected step size
-- lr * sqrt(1 - beta2^t) / (1 - beta1^t), the epsilon scaled by the same
-- sqrt(1 - beta2^t), and the constant one.  `step` counts from 1.
--
-- Everything is binary32, the kernel's own precision.  1 - beta2^t loses
-- about five significant digits to cancellation at small t (beta2^t is close
-- to one), so the step size carries a relative error near 1e-5; the kernel
-- applies the same tolerance to its running moments, and the value is what the
-- learner declares rather than a number computed elsewhere.
def bobAdamStepScalars =
  (lambda unrestricted step : Nat .
    (let unrestricted rootSecond =
      (f32-sqrt (f32-subtract bobF32One (bobF32Power bobBeta2 step)))
    in
      (bytes-builder-build
        (bytes-builder-append (bytes-builder-chunk (dataFloat32EncodeLE bobF32One))
        (bytes-builder-append (bytes-builder-chunk (dataFloat32EncodeLE bobBeta1))
        (bytes-builder-append (bytes-builder-chunk (dataFloat32EncodeLE (f32-subtract bobF32One bobBeta1)))
        (bytes-builder-append (bytes-builder-chunk (dataFloat32EncodeLE bobBeta2))
        (bytes-builder-append (bytes-builder-chunk (dataFloat32EncodeLE (f32-subtract bobF32One bobBeta2)))
        (bytes-builder-append (bytes-builder-chunk (dataFloat32EncodeLE (f32-subtract bobF32One (f32-multiply bobLearningRate bobWeightDecay))))
        (bytes-builder-append (bytes-builder-chunk (dataFloat32EncodeLE
          (f32-divide (f32-multiply bobLearningRate rootSecond)
                      (f32-subtract bobF32One (bobF32Power bobBeta1 step)))))
        (bytes-builder-append (bytes-builder-chunk (dataFloat32EncodeLE (f32-multiply bobEpsilon rootSecond)))
                              (bytes-builder-chunk (dataFloat32EncodeLE bobF32One)))))))))))))