module Learning.ScaledGradientUpdate

import Component.Pointwise
import Std.Fin
import Std.List
import Std.Natural

-- These two operations compose independently: callers choose the arithmetic
-- and the source of the gradient. The accumulation state is persistent, and
-- applying the parameter update consumes the newly accumulated value.
def accumulateScaledGradient =
  (lambda erased value : Type 0 .
    (lambda unrestricted add : (pi unrestricted left : value . (pi unrestricted right : value . value)) .
      (lambda unrestricted multiply : (pi unrestricted left : value . (pi unrestricted right : value . value)) .
        (lambda unrestricted decay : value .
          (lambda unrestricted previous : value .
            (lambda unrestricted gradient : value . (add (multiply decay previous) gradient)))))))

def subtractScaledUpdate =
  (lambda erased value : Type 0 .
    (lambda unrestricted subtract : (pi unrestricted left : value . (pi unrestricted right : value . value)) .
      (lambda unrestricted multiply : (pi unrestricted left : value . (pi unrestricted right : value . value)) .
        (lambda unrestricted rate : value .
          (lambda unrestricted parameter : value .
            (lambda unrestricted update : value . (subtract parameter (multiply rate update))))))))

def scaledGradientReadOnly =
  (lambda erased slots : Nat .
    (lambda unrestricted writable : (family StdFin slots) .
      (lambda unrestricted gradient : (family StdFin slots) .
        (lambda unrestricted decay : (family StdFin slots) .
          (lambda unrestricted rate : (family StdFin slots) .
            (naturalAnd
              (pointwiseSlotsSeparate slots writable gradient)
              (naturalAnd
                (pointwiseSlotsSeparate slots writable decay)
                (pointwiseSlotsSeparate slots writable rate))))))))

def scaledGradientSlotsIndependent =
  (lambda erased slots : Nat .
    (lambda unrestricted parameter : (family StdFin slots) .
      (lambda unrestricted accumulator : (family StdFin slots) .
        (lambda unrestricted gradient : (family StdFin slots) .
          (lambda unrestricted decay : (family StdFin slots) .
            (lambda unrestricted rate : (family StdFin slots) .
              (naturalAnd
                (pointwiseSlotsSeparate slots parameter accumulator)
                (naturalAnd
                  (scaledGradientReadOnly slots parameter gradient decay rate)
                  (scaledGradientReadOnly slots accumulator gradient decay rate)))))))))

-- The same equations can produce a typed state program instead of evaluating
-- numbers immediately. Positions are supplied by the state's owner. Assigning
-- the accumulator first makes its new value available to the parameter update;
-- the compiler must preserve that dependency and each separate rounding.
-- Writable slots are independent of each other and of the read-only inputs;
-- otherwise the first assignment could overwrite a later equation's input.
def scaledGradientAssignments =
  (lambda erased word : Type 0 .
    (lambda erased slots : Nat .
      (lambda unrestricted parameter : (family StdFin slots) .
        (lambda unrestricted accumulator : (family StdFin slots) .
          (lambda unrestricted gradient : (family StdFin slots) .
            (lambda unrestricted decay : (family StdFin slots) .
              (lambda unrestricted rate : (family StdFin slots) .
                (lambda erased independent : (equal Nat (scaledGradientSlotsIndependent slots parameter accumulator gradient decay rate) 1) .
                  (let unrestricted input =
                    (lambda unrestricted slot : (family StdFin slots) .
                      (constructor PointwiseExpression PointwiseInput word slots slot))
                    in
                    (let unrestricted multiply =
                      (pointwiseBinary word slots (constructor PointwiseBinary PointwiseMultiply))
                      in
                      (constructor
                        StdList
                        StdListCons
                        (family PointwiseAssignment word slots)
                        (constructor
                          PointwiseAssignment
                          PointwiseAssign
                          word
                          slots
                          accumulator
                          (accumulateScaledGradient
                            (family PointwiseExpression word slots)
                            (pointwiseBinary word slots (constructor PointwiseBinary PointwiseAdd))
                            multiply
                            (input decay)
                            (input accumulator)
                            (input gradient)))
                        (constructor
                          StdList
                          StdListCons
                          (family PointwiseAssignment word slots)
                          (constructor
                            PointwiseAssignment
                            PointwiseAssign
                            word
                            slots
                            parameter
                            (subtractScaledUpdate
                              (family PointwiseExpression word slots)
                              (pointwiseBinary
                                word
                                slots
                                (constructor PointwiseBinary PointwiseSubtract))
                              multiply
                              (input rate)
                              (input parameter)
                              (input accumulator)))
                          (constructor StdList StdListEmpty (family PointwiseAssignment word slots))))))))))))))